Advent of Code - 2021 day 16 - visualization

viz.py

#!/usr/bin/env -S python3 -u

import sys
import time

from functools import reduce
from io import StringIO

import aocutil

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

    return ''.join(format(int(x, 16), '04b') for x in data)

class Packet:
    def __init__(self, level=0):
        self.level = level
        self.version = None
        self.type = None
        self.value = None
        self.length = 0
        self.ltypeid = None
        self.subpackets = None
        self.packets = []

    def parse(self, sio):
        self.version = int(sio.read(3), 2)
        self.type = int(sio.read(3), 2)
        self.length = 6
        # print(f'Got packet v{self.version} type {self.type}:', end=' ')
        if self.type == 4:
            # Literal
            value = ''
            while True:
                packet = sio.read(5)
                self.packets.append(packet)

                notlast = int(packet[0])
                value += packet[1:]
                self.length += 5
                if not notlast:
                    break

            self.value = int(value, 2)
            # print(self.value)
        else:
            # Operator type
            self.ltypeid = int(sio.read(1))
            if self.ltypeid:
                # Number of subpackets
                self.subpackets = int(sio.read(11), 2)
                # print(f'{self.ltypeid}:{self.subpackets}')
                self.length += 12
                for _ in range(self.subpackets):
                    # print(f'Parsing subpacket #{_} @ L{self.level + 1}')
                    packet = Packet(self.level + 1)
                    self.length += packet.parse(sio)
                    self.packets.append(packet)
            else:
                # Number of bits
                self.subpackets = int(sio.read(15), 2)
                self.length += 16
                read = 0
                # print(f'{self.ltypeid}:{self.subpackets}')
                while read < self.subpackets:
                    # print(f'Parsing subpacket @ L{self.level + 1}')
                    packet = Packet(self.level + 1)
                    read += packet.parse(sio)
                    # print(f'Got {read} bits, need {self.subpackets}')
                    self.packets.append(packet)

                self.length += read

        # print(f'Packet @ L{self.level} has length {self.length}')
        return self.length

    def dump(self):
        margin = '    ' * self.level
        print(f'{margin}V{self.version} T{self.type}', end=' ')
        if self.value is not None:
            print(self.value)
        else:
            print(f'{self.ltypeid}:{self.subpackets}')
            for packet in self.packets:
                packet.dump()

    def version_sum(self):
        return self.version + sum(p.version_sum() for p in self.packets if isinstance(p, Packet))

    def evaluate(self):
        if self.type != 4:
            subvals = [p.evaluate() for p in self.packets]

        if self.type == 0:
            # Sum packet
            return sum(subvals)
        elif self.type == 1:
            # Product packet
            return reduce(lambda x, y: x*y, subvals)
        elif self.type == 2:
            return min(subvals)
        elif self.type == 3:
            return max(subvals)
        elif self.type == 4:
            return self.value
        elif self.type == 5:
            return int(subvals[0] > subvals[1])
        elif self.type == 6:
            return int(subvals[0] < subvals[1])
        elif self.type == 7:
            return int(subvals[0] == subvals[1])

    def viz(self):
        margin = '    ' * self.level
        print(margin, end='')

        def print_string(string, color):
            print(f'\033[3{color}m', end='')
            for char in string:
                print(char, end='')
                time.sleep(0.02)
            print('\033[m', end=' ')

        def print_num(packet):
            more = int(packet[0])
            color = [1, 2][more]
            print(f'\033[3{color}m{packet[0]}\033[m', end='')
            print_string(packet[1:], 4)

        print_string(format(self.version, '03b'), 3)
        print_string(format(self.type, '03b'), 6)

        if self.value is not None:
            for packet in self.packets:
                print_num(packet)
            print(f"\033[92m// {self.value}")
        else:
            ltypeid = str(self.ltypeid)
            color = [5, 6][self.ltypeid]
            packet = ['Sum', 'Product', 'Minimum', 'Maximum', 'Literal', 'Greater Than', 'Less Than', 'Equal To'][self.type]
            print(f'\033[3{color}m{ltypeid}\033[m', end='')
            if self.ltypeid:
                print_string(format(self.subpackets, '011b'), 4)
                print(f'\033[92m// {packet} packet ({self.subpackets} subpackets)')
            else:
                print_string(format(self.subpackets, '015b'), 4)
                print(f'\033[92m// {packet} packet ({self.subpackets} bits in subpackets)')

            for packet in self.packets:
                packet.viz()

def _main():
    data = load_data()

    sio = StringIO(data)
    packet = Packet()
    packet.parse(sio)

    packet.viz()

if __name__ == '__main__':
    _main()