Advent of Code - 2021 day 17 - visualization

viz.py

#!/usr/bin/env python3

import sys
import re
import math
import curses
import time

import aocutil

def load_data():
    with aocutil.read_data() as df:
        data = df.read().strip()

    x1, x2, y1, y2 = list(map(int, re.findall(r'-?\d+', data)))

    return min(x1, x2), max(x1, x2), min(y1, y2), max(y1, y2)

def step(dx0, dy0):
    x, y = 0, 0
    dx, dy = dx0, dy0
    while True:
        x, y = x + dx, y + dy
        if dx > 0:
            dx -= 1
        dy -= 1
        yield (x, y)

def within_target_area(x, y, xmin, xmax, ymin, ymax):
    return (xmin <= x <= xmax) and (ymin <= y <= ymax)

def crossed_target_area(x, y, xmin, xmax, ymin, ymax):
    return x > xmax or y < ymin

def compute_velocities():
    xmin, xmax, ymin, ymax = load_data()

    # Compute range of dx, dy
    dxmin = math.ceil((-1 + math.sqrt(1 + 8*xmin)) / 2)
    dxmax = xmax
    dymin = ymin
    dymax = -ymin - 1

    velocities = []

    for dx in range(dxmin, dxmax + 1):
        for dy in range(dymin, dymax + 1):
            for x, y in step(dx, dy):
                if crossed_target_area(x, y, xmin, xmax, ymin, ymax):
                    break

                if within_target_area(x, y, xmin, xmax, ymin, ymax):
                    velocities.append((dx, dy))
                    break

    return velocities

def _main(stdscr):
    xmin, xmax, ymin, ymax = load_data()
    velocities = [(dx, dy) for (dx, dy) in compute_velocities()
                  if not within_target_area(dx, dy, xmin, xmax, ymin, ymax)]
    curses.curs_set(0)

    y_start = max(v[1] for v in velocities) + 1

    def fire(dx, dy):
        stdscr.clear()

        # Draw submarine at y_start, 0
        stdscr.addstr(y_start, 0, 'S')

        # Draw target area
        for y in range(ymin, ymax+1):
            stdscr.addstr(y_start - y, xmin, '.' * (xmax - xmin + 1))

        # Draw velocity
        stdscr.addstr(0, 0, f'({dx:+d}, {dy:+d})')

        stdscr.refresh()
        time.sleep(0.5)

        for x, y in step(dx, dy):
            ys = y_start - y
            char = 'o'
            if ys < 0:
                ys = 0
                char = '^'

            if within_target_area(x, y, xmin, xmax, ymin, ymax):
                stdscr.addstr(ys, x, '*')
                break
            else:
                stdscr.addstr(ys, x, char)

            stdscr.refresh()
            time.sleep(0.1)

    for dx, dy in velocities:
        fire(dx, dy)
        stdscr.refresh()
        time.sleep(0.5)


if __name__ == '__main__':
    curses.wrapper(_main)