import itertools as it
from typing import Any, Iterator

import numpy as np
import parse  # type: ignore
from numpy.typing import NDArray

from ..base import BaseSolver


class Solver(BaseSolver):
    def part1(
        self, sensor_to_beacon: dict[tuple[int, int], tuple[int, int]], row: int
    ) -> int:
        no_beacons_row_l: list[NDArray[np.floating[Any]]] = []

        for (sx, sy), (bx, by) in sensor_to_beacon.items():
            d = abs(sx - bx) + abs(sy - by)  # closest

            no_beacons_row_l.append(sx - np.arange(0, d - abs(sy - row) + 1))  # type: ignore
            no_beacons_row_l.append(sx + np.arange(0, d - abs(sy - row) + 1))  # type: ignore

        beacons_at_row = set(bx for (bx, by) in sensor_to_beacon.values() if by == row)
        no_beacons_row = set(it.chain(*no_beacons_row_l)).difference(beacons_at_row)  # type: ignore

        return len(no_beacons_row)

    def part2_intervals(
        self, sensor_to_beacon: dict[tuple[int, int], tuple[int, int]], xy_max: int
    ) -> tuple[int, int, int]:
        for y in self.progress.wrap(range(xy_max + 1)):
            its: list[tuple[int, int]] = []
            for (sx, sy), (bx, by) in sensor_to_beacon.items():
                d = abs(sx - bx) + abs(sy - by)
                dx = d - abs(sy - y)

                if dx >= 0:
                    its.append((max(0, sx - dx), min(sx + dx, xy_max)))

            its = sorted(its)
            _, e = its[0]

            for si, ei in its[1:]:
                if si > e + 1:
                    return si - 1, y, 4_000_000 * (si - 1) + y
                if ei > e:
                    e = ei

        return (0, 0, 0)

    def part2_cplex(
        self, sensor_to_beacon: dict[tuple[int, int], tuple[int, int]], xy_max: int
    ) -> tuple[int, int, int]:
        from docplex.mp.model import Model

        m = Model()

        x, y = m.continuous_var_list(2, ub=xy_max, name=["x", "y"])

        for (sx, sy), (bx, by) in sensor_to_beacon.items():
            d = abs(sx - bx) + abs(sy - by)
            m.add_constraint(
                m.abs(x - sx) + m.abs(y - sy) >= d + 1,  # type: ignore
                ctname=f"ct_{sx}_{sy}",
            )

        m.set_objective("min", x + y)

        s = m.solve()
        assert s is not None

        vx = int(s.get_value(x))
        vy = int(s.get_value(y))
        return vx, vy, 4_000_000 * vx + vy

    def solve(self, input: str) -> Iterator[Any]:
        lines = input.splitlines()

        sensor_to_beacon: dict[tuple[int, int], tuple[int, int]] = {}

        for line in lines:
            r: dict[str, str] = parse.parse(  # type: ignore
                "Sensor at x={sx}, y={sy}: closest beacon is at x={bx}, y={by}", line
            )
            sensor_to_beacon[int(r["sx"]), int(r["sy"])] = (int(r["bx"]), int(r["by"]))

        xy_max = 4_000_000 if max(sensor_to_beacon) > (1_000, 0) else 20
        row = 2_000_000 if max(sensor_to_beacon) > (1_000, 0) else 10

        yield self.part1(sensor_to_beacon, row)

        # x, y, a2 = part2_cplex(sensor_to_beacon, xy_max)
        x, y, a2 = self.part2_intervals(sensor_to_beacon, xy_max)
        self.logger.info(f"answer 2 is {a2} (x={x}, y={y})")
        yield a2