from collections import defaultdict

from tools.aoc import AOCDay
from typing import Any


def flatten(what: set) -> set:
    found_something = True
    while found_something:
        found_something = False
        liqlist = []
        for c in what:
            for c2 in liqlist:
                if c[0] <= c2[0] and c[1] >= c2[1]:
                    found_something = True
                    c2[0] = c[0]
                    c2[1] = c[1]
                elif c[0] >= c2[0] and c[1] <= c2[1]:
                    found_something = True
                    break
                elif c2[0] <= c[1] < c2[1]:
                    found_something = True
                    c2[0] = c[0]
                    break
                elif c2[0] <= c[0] < c2[1]:
                    found_something = True
                    c2[1] = c[1]
                    break
            else:
                liqlist.append([c[0], c[1]])

        what = set((x[0], x[1]) for x in liqlist)

    return what


class Day(AOCDay):
    inputs = [
        [
            (26, "input15_test"),
            (6275922, "input15"),
        ],
        [
            (56000011, "input15_test"),
            (11747175442119, "input15"),
        ]
    ]

    def get_sensor_beacon_dict(self):
        sensors = {}
        for line in self.getInput():
            sp = line.split(" ")
            sx, sy, bx, by = sp[2], sp[3], sp[8], sp[9]
            sx = int(sx[2:-1])
            sy = int(sy[2:-1])
            bx = int(bx[2:-1])
            by = int(by[2:])
            sensors[(sx, sy)] = (bx, by)

        return sensors

    def part1(self) -> Any:
        liq = 10 if self._current_test_file.endswith("_test") else 2000000
        liqset = set()
        for sensor, beacon in self.get_sensor_beacon_dict().items():
            d = abs(sensor[0] - beacon[0]) + abs(sensor[1] - beacon[1])
            if sensor[1] - d > liq > sensor[1] + d:
                continue

            r = 2 * (d - abs(sensor[1] - liq)) + 1
            if r < 1:
                continue
            x_min, x_max = sensor[0] - r // 2, sensor[0] + r // 2
            liqset.add((x_min, x_max))

        liqset = flatten(liqset)
        return sum(x[1] - x[0] for x in liqset)

    def part2(self) -> Any:
        max_c = 20 if self._current_test_file.endswith("_test") else 4000000
        occ = defaultdict(set)
        for sensor, beacon in self.get_sensor_beacon_dict().items():
            d = abs(sensor[0] - beacon[0]) + abs(sensor[1] - beacon[1])
            for y in range(-d, d + 1):
                r = 2 * (d - abs(y)) + 1
                if 0 <= sensor[1] + y <= max_c:
                    occ[sensor[1] + y].add((max(0, sensor[0] - r // 2), min(sensor[0] + r // 2, max_c)))

        for y, occs in occ.items():
            occs = flatten(occs)
            assert len(occs) in [1, 2]
            if len(occs) == 2:
                f = list(occs)
                min_y = min(x[1] for x in f)
                max_x = max(x[0] for x in f)
                if max_x - min_y == 2:
                    return (min_y + 1) * 4000000 + y

        return ""


if __name__ == '__main__':
    day = Day(2022, 15)
    day.run(verbose=True)