from dataclasses import dataclass
from functools import lru_cache, reduce
from typing import List, Generator, Dict, Optional

from aoc import BaseAssignment


class Assignment(BaseAssignment):
    def parse_item(self, item: str) -> int:
        return int(item)


class AssignmentOne(Assignment):
    def run(self, input: Generator) -> int:
        adapters = sorted(input)
        device_adapter = max(adapters) + 3

        adapter_list = [
            0, *adapters, device_adapter
        ]

        differences = [
            adapter_list[index + 1] - adapter
            for index, adapter in enumerate(adapter_list)
            if adapter != device_adapter
        ]

        return differences.count(1) * differences.count(3)


@dataclass
class Node:
    value: int
    next: List['Node']
    paths: int = 0

    def __repr__(self):
        return str(f'<Node: {self.value}, paths: {self.paths}>')

    def __post_init__(self):
        self.paths = reduce(lambda count, node: count + 1 if len(node.next) == 0 else count + node.paths, self.next, 0)


class AssignmentTwo(Assignment):
    def generate_graph(
            self, input: List, pointers: Dict[int, Node] = None) -> Optional[Node]:
        if pointers is None:
            pointers = {}

        if len(input) == 0:
            return

        value, next = input[0], input[1:]

        pointers[value] = Node(value=value, next=[
            (
                pointers[next_value]
                if next_value in pointers
                else self.generate_graph(
                    input=input[index + 1:],
                    pointers=pointers
                )
            )
            for index, next_value in enumerate(next[:3])
            if next_value - value <= 3
        ])

        return pointers[value]

    def run(self, input: Generator) -> int:
        adapters = sorted(input)
        device_adapter = max(adapters) + 3

        return self.generate_graph([0, *adapters, device_adapter]).paths