from tools.aoc import AOCDay
from tools.coordinate import Coordinate
from tools.grid import Grid
from typing import Any

SIDE_LOOKUP = {
    (0, -1): {
        "next": (1, 0),
        "turn_look": (-1, 0),
    },
    (0, 1): {
        "next": (-1, 0),
        "turn_look": (1, 0),
    },
    (1, 0): {
        "next": (0, 1),
        "turn_look": (0, -1),
    },
    (-1, 0): {
        "next": (0, -1),
        "turn_look": (0, 1),
    },
}


def get_sides(region: set[Coordinate]) -> int:
    sides = 0
    top_left = cur_pos = min(region)
    start_look = next_look = (0, -1)
    while cur_pos != top_left or start_look != next_look or sides < 4:
        if cur_pos + next_look not in region:
            if cur_pos + SIDE_LOOKUP[next_look]["next"] in region:
                cur_pos += SIDE_LOOKUP[next_look]["next"]
                continue
            else:
                sides += 1
                next_look = SIDE_LOOKUP[next_look]["next"]
        else:
            sides += 1
            cur_pos += next_look
            next_look = SIDE_LOOKUP[next_look]["turn_look"]

    return sides


class Day(AOCDay):
    inputs = [
        [
            (1930, "input12_test"),
            (1424006, "input12"),
        ],
        [
            (80, "input12_test2"),
            (236, "input12_test3"),
            (368, "input12_test4"),
            (1206, "input12_test"),
            (858684, "input12"),
        ],
    ]

    def get_regions(self) -> list[set[Coordinate]]:
        regions = []
        garden = Grid.from_data(self.getInput())
        seen = set()
        for x in garden.rangeX():
            for y in garden.rangeY():
                pos = Coordinate(x, y)
                if pos in seen:
                    continue

                region = set(garden.getRegion(pos, includeDiagonal=False))
                seen |= region
                regions.append(region)

        return regions

    def part1(self) -> Any:
        price = 0
        for region in self.get_regions():
            perimeter = sum(1 for cell in region for n in cell.getNeighbours(includeDiagonal=False) if n not in region)
            price += perimeter * len(region)

        return price

    def part2(self) -> Any:
        price = 0
        for region in self.get_regions():
            region_sides = get_sides(region)
            region_plots = len(region)

            g = Grid()
            for cell in region:
                g.set(cell)
            g.minX -= 1
            g.minY -= 1
            g.maxX += 1
            g.maxY += 1
            region |= set(g.getRegion(Coordinate(g.minX, g.minY), includeDiagonal=True))
            for rx in g.rangeX():
                for ry in g.rangeY():
                    r_pos = Coordinate(rx, ry)
                    if r_pos in region:
                        continue
                    inner_region = set(g.getRegion(r_pos, includeDiagonal=False))
                    region |= inner_region
                    region_sides += get_sides(inner_region)

            price += region_plots * region_sides

        return price


if __name__ == "__main__":
    day = Day(2024, 12)
    day.run(verbose=True)