from __future__ import annotations
import itertools
from collections import deque
from enum import Enum
from tools.aoc import AOCDay
from typing import Any


class ItemType(int, Enum):
    GENERATOR = 0
    MICROCHIP = 1


class Item:
    def __init__(self, element: str, typ: ItemType):
        self.element = element
        self.type = typ

    def compatible(self, other: Item | Generator | Microchip) -> bool:
        if self.type == other.type or self.element == other.element:
            return True
        else:
            return False

    def __str__(self):
        return "%s(%s)" % (self.__class__.__name__, self.element)

    def __repr__(self):
        return str(self)


class Generator(Item):
    def __init__(self, element: str):
        super().__init__(element, ItemType.GENERATOR)


class Microchip(Item):
    def __init__(self, element: str):
        super().__init__(element, ItemType.MICROCHIP)


def is_valid(group: list[Generator | Microchip]) -> bool:
    generators = set()
    microchips = set()
    for item in group:
        if item.type == ItemType.GENERATOR:
            generators.add(item.element)
        else:
            microchips.add(item.element)

    if not generators or not microchips:
        return True

    return all(m in generators for m in microchips)


def get_legal_moves(
    from_floor: int, floors: list[list[Generator | Microchip]], debug: bool = False
) -> set[tuple[int, Microchip | Generator, ...]]:
    targets = []
    if 0 < from_floor:
        targets.append(from_floor - 1)
    if from_floor < len(floors) - 1:
        targets.append(from_floor + 1)

    moves = set()
    if len(floors[from_floor]) == 1:
        for target in targets:
            if is_valid(floors[target] + [floors[from_floor][0]]):
                moves.add((target, floors[from_floor][0]))
    else:
        for a, b in itertools.combinations(floors[from_floor], 2):
            if not a.compatible(b):
                continue

            for target in targets:
                if is_valid(floors[target] + [a]):
                    moves.add((target, a))
                if is_valid(floors[target] + [b]):
                    moves.add((target, b))
                if is_valid(floors[target] + [a, b]):
                    moves.add((target, a, b))

    return moves


def move_items(
    floors: list[list[Generator | Microchip]], from_floor: int, move: tuple[int, Microchip | Generator, ...]
) -> None:
    to_floor = move[0]
    for item in move[1:]:
        floors[from_floor].remove(item)
        floors[to_floor].append(item)


def floor_hash(floor: int, floors: list[list[Generator | Microchip]]) -> str:
    return (
        str(floor)
        + "@"
        + ";".join(
            "%d:%s"
            % (
                i,
                ",".join(
                    x.__class__.__name__[0] + x.element[0]
                    for x in sorted(floors[i], key=lambda f: (f.__class__.__name__, f.element))
                ),
            )
            for i in range(len(floors))
        )
    )


def copy_floors(floors: list[list[Generator | Microchip]]) -> list[list[Generator | Microchip]]:
    return [list(floors[x]) for x in range(len(floors))]


def all_on_floor(floors: list[list[Generator | Microchip]], target_floor: int) -> bool:
    all_floors_empty = True
    for i in range(len(floors)):
        if i != target_floor and len(floors[i]) > 0:
            all_floors_empty = False

    return all_floors_empty


def get_min_steps(floors: list[list[Generator | Microchip]]) -> int:
    q = deque([(0, 3, floors)])
    seen = set()
    while q:
        dist, cur_floor_id, cur_floors = q.popleft()
        if cur_floor_id == 0 and all_on_floor(cur_floors, cur_floor_id):
            return dist

        cur_floor_hash = floor_hash(cur_floor_id, cur_floors)
        if cur_floor_hash in seen:
            continue
        seen.add(cur_floor_hash)

        for move in get_legal_moves(cur_floor_id, cur_floors):
            nxt_floors = copy_floors(cur_floors)
            move_items(nxt_floors, cur_floor_id, move)
            q.append((dist + 1, move[0], nxt_floors))

    return 0


class Day(AOCDay):
    inputs = [
        [
            (11, "input11_test"),
            (37, "input11"),
        ],
        [
            (61, "input11"),
        ],
    ]

    def parse_input(self, p2: bool = False) -> list[list[Generator | Microchip]]:
        floors = []
        for line in reversed(self.getInput()):
            _, contents = line[:-1].split(" contains ")
            if contents == "nothing relevant":
                floors.append([])
                continue

            floor_contents = []
            contents = contents.replace(", and ", ", ")
            contents = contents.replace(" and ", ", ")
            contents = contents.split(", ")
            for content in contents:
                if not content:
                    continue
                _, element, typ = content.split()
                if typ == "generator":
                    floor_contents.append(Generator(element))
                else:
                    element = element.split("-")[0]
                    floor_contents.append(Microchip(element))

            floors.append(floor_contents)

        if p2:
            floors[3].append(Generator("elerium"))
            floors[3].append(Microchip("elerium"))
            floors[3].append(Generator("dilithium"))
            floors[3].append(Microchip("dilithium"))

        return floors

    def part1(self) -> Any:
        return get_min_steps(self.parse_input())

    def part2(self) -> Any:
        return get_min_steps(self.parse_input(p2=True))


if __name__ == "__main__":
    day = Day(2016, 11)
    day.run(verbose=True)