import re
from abc import ABC
from copy import copy
from dataclasses import dataclass
from typing import List, Tuple, Union, TypedDict

from aoc import BaseAssignment

Coordinate = Tuple[int, int]

def mirrored_index(index: int, fold: int) -> int:
    return fold + (fold - index)

@dataclass
class Sheet:
    dots: List[Coordinate]
    width: int = 0
    height: int = 0

    def __post_init__(self):
        x_coordinates, y_coordinates = zip(*self.dots)

        self.width = max(x_coordinates) + 1
        self.height = max(y_coordinates) + 1


    def fold_x(self, position: int):
        for x, y in copy(self.dots):
            mirrored_x = mirrored_index(x, position)
            if mirrored_x < position:
                self.dots.append((mirrored_x, y))
                self.dots.remove((x, y))

        self.width = self.width // 2
        self.dots = list(set(self.dots))

    def fold_y(self, position: int):
        for x, y in copy(self.dots):
            mirrored_y = mirrored_index(y, position)
            if mirrored_y < position:
                self.dots.append((x, mirrored_y))
                self.dots.remove((x, y))

        self.height = self.height // 2
        self.dots = list(set(self.dots))

    def __repr__(self):
        dots = set(self.dots)
        return '\n'.join([
            '',
            *[
                ''.join([
                    '█' if (x, y) in dots else ' '
                    for x in range(self.width)
                ])
                for y in range(self.height)
            ],
            '',
        ])

class Instruction(TypedDict):
    direction: str
    position: int

coordinate_regex = re.compile('(\d+),(\d+)')
instruction_regex = re.compile('fold along (?P<direction>[xy])=(?P<position>\d+)')

class Assignment(BaseAssignment, ABC):
    def parse_item(self, item: str) -> Union[Coordinate,Instruction]:
        if coordinate := coordinate_regex.match(item):
            return tuple(int(i) for i in coordinate.groups())

        if instruction := instruction_regex.match(item):
            return {
                **instruction.groupdict(),
                'position': int(instruction.groupdict()['position'])
            }

    def read_input(self, example = False) -> Tuple[Sheet, List[Instruction]]:
        input = super().read_input(example)

        coordinates: List[Coordinate] = []

        while True:
            coordinate = next(input)
            if coordinate is None:
                break

            coordinates.append(coordinate)

        sheet = Sheet(dots=coordinates)

        return sheet, list(input)


class AssignmentOne(Assignment):
    example_result = 17

    def run(self, input: Tuple[Sheet, List[Instruction]]) -> int:
        sheet, instructions = input
        getattr(sheet, f'fold_{instructions[0]["direction"]}')(instructions[0]["position"])
        return len(sheet.dots)

class AssignmentTwo(Assignment):
    def run(self, input: Tuple[Sheet, List[Instruction]]) -> Sheet:
        sheet, instructions = input
        for instruction in instructions:
            getattr(sheet, f'fold_{instruction["direction"]}')(instruction["position"])
        return sheet