from tools.aoc import AOCDay
from typing import Any


def HASH(s: str) -> int:
    this_hash = 0
    for c in s:
        this_hash = (this_hash + ord(c)) * 17 % 256

    return this_hash


class Box:
    def __init__(self, nr: int):
        self.nr = nr
        self.labels = []
        self.focal_lengths = []

    def add(self, label: str, focal_length: int) -> None:
        if label in self.labels:
            self.focal_lengths[self.labels.index(label)] = focal_length
        else:
            self.labels.append(label)
            self.focal_lengths.append(focal_length)

    def remove(self, label: str) -> None:
        if label not in self.labels:
            return

        p = self.labels.index(label)
        self.labels.pop(p)
        self.focal_lengths.pop(p)

    def get_focus_power(self) -> int:
        return sum((self.nr + 1) * (x + 1) * l for x, l in enumerate(self.focal_lengths))


class Day(AOCDay):
    inputs = [
        [
            (1320, "input15_test"),
            (513214, "input15_dennis"),
            (513158, "input15"),
        ],
        [
            (145, "input15_test"),
            (258826, "input15_dennis"),
            (200277, "input15"),
        ]
    ]

    def part1(self) -> Any:
        return sum(HASH(c) for c in self.getInputAsArraySplit(","))

    def part2(self) -> Any:
        boxes = {}
        for s in self.getInputAsArraySplit(","):
            if '-' in s:
                label = s[:-1]
                label_hash = HASH(label)
                boxes.setdefault(label_hash, Box(label_hash)).remove(label)
            else:
                label, f_str = s.split("=")
                focal_length = int(f_str)
                label_hash = HASH(label)
                boxes.setdefault(label_hash, Box(label_hash)).add(label, focal_length)

        return sum(x.get_focus_power() for x in boxes.values())


if __name__ == '__main__':
    day = Day(2023, 15)
    day.run(verbose=True)