# -*- coding: utf-8 -*-
import re
from abc import ABC
from dataclasses import dataclass, field
from typing import Tuple, Iterator, Set, List
from typing import Union

from aoc import BaseAssignment
from aoc.datastructures import Coordinate


@dataclass
class Sensor:
    coordinate: Coordinate
    nearest: Coordinate
    radius: int = None

    def __hash__(self):
        return hash(self.coordinate)

    def __post_init__(self):
        self.radius = self.coordinate.manhattan_distance(self.nearest)

    def within_radius(self, coordinate: Coordinate) -> bool:
        distance = self.coordinate.manhattan_distance(coordinate)
        return distance <= self.radius

    def x_coordinates_within_radius_at(
        self, y: int, min_x: int = None, max_x: int = None
    ) -> Union[range, list]:

        coordinate = Coordinate(self.coordinate.x, y)

        if not self.within_radius(coordinate):
            return []

        dy = abs(self.coordinate.y - y)
        left_over = self.radius - dy

        left_x = self.coordinate.x - left_over
        right_x = self.coordinate.x + left_over

        return range(
            max(left_x, min_x if min_x is not None else left_x - 1),
            min(right_x, max_x if max_x is not None else right_x + 1) + 1,
        )

    def area_inside(self) -> Set[Coordinate]:
        coordinates = set()

        for y in range(
            self.coordinate.y - (self.radius - 1), self.coordinate.y + self.radius
        ):
            for x in self.x_coordinates_within_radius_at(y):
                coordinates.add(Coordinate(x, y))

        return coordinates


@dataclass
class Map:
    sensors: List[Sensor] = field(default_factory=list)
    beacons: Set[Coordinate] = field(default_factory=set)

    def __post_init__(self):
        all_coordinates = self.beacons | {sensor.coordinate for sensor in self.sensors}

        min_x = min(c.x for c in all_coordinates)
        max_x = max(c.x for c in all_coordinates)
        min_y = min(c.y for c in all_coordinates)
        max_y = max(c.y for c in all_coordinates)

        self.width = list(range(min_x, max_x + 1))
        self.height = list(range(min_y, max_y + 1))


input_pattern = re.compile("x=(-?[-0-9]+), y=(-?[0-9]+)")


class Assignment(BaseAssignment[int, Tuple[Sensor, Coordinate]], ABC):
    def parse_item(self, line: str) -> Tuple[Sensor, Coordinate]:
        match = input_pattern.findall(line)

        if len(match) != 2:
            raise RuntimeError()

        sensor_match, beacon_match = match

        beacon = Coordinate(int(beacon_match[0]), int(beacon_match[1]))
        sensor = Sensor(Coordinate(int(sensor_match[0]), int(sensor_match[1])), beacon)

        return sensor, beacon

    def create_map(self, input: Iterator[Tuple[Sensor, Coordinate]]) -> Map:
        sensors = list()
        beacons = set()

        for sensor, beacon in input:
            sensors.append(sensor)
            beacons.add(beacon)

        return Map(sensors, beacons)


class AssignmentOne(Assignment):
    example_result = 26
    example_kwargs = {"y": 10}

    # Wrong: 4062163

    def run(self, input: Iterator[Tuple[Sensor, Coordinate]], y=2000000):
        map = self.create_map(input)

        x_coordinates = {
            x
            for sensor in map.sensors
            for x in sensor.x_coordinates_within_radius_at(y)
        }

        return len(x_coordinates) - len(
            [beacon.x for beacon in map.beacons if beacon.y == y]
        )


class AssignmentTwo(Assignment):
    example_result = 56000011
    example_kwargs = {"max": 20}

    def run(self, input: Iterator[Tuple[Sensor, Coordinate]], max=4000000):
        map = self.create_map(input)
        coordinates = {Coordinate(x, y) for x in range(max) for y in range(max)}

        print("Coordinates generated")

        for sensor in map.sensors:
            coordinates -= sensor.area_inside()

        distress_beacon = (coordinates - map.beacons).pop()

        return (distress_beacon.x * 4000000) + distress_beacon.y