from aocd.models import Puzzle
from aocd import submit
from functools import reduce
from operator import itemgetter

import re


def parse_almanac(data):
    NUMBER_RE = re.compile(r"\d+")
    stanzas = data.split("\n\n")
    return [int(m.group(0)) for m in NUMBER_RE.finditer(stanzas[0])], [
        sorted(
            (
                (nums[1], nums[1] + nums[2], nums[0] - nums[1])
                for line in stanza.splitlines()
                if len(nums := [int(m.group(0)) for m in NUMBER_RE.finditer(line)]) == 3
            ),
            key=itemgetter(0),
        )
        for stanza in stanzas[1:]
    ]


def remap_almanac(ranges, mappings):
    for start, end in ranges:
        for start2, end2, offset in mappings:
            if start2 >= end or start >= end2:
                continue
            if start < start2:
                yield start, start2
                start = start2
            end2 = min(end, end2)
            yield start + offset, end2 + offset
            start = end2
        if start < end:
            yield start, end


def part_a(data):
    seeds, mappings = parse_almanac(data)
    return min(
        map(itemgetter(0), reduce(remap_almanac, mappings, ((x, x + 1) for x in seeds)))
    )


def part_b(data):
    seeds, mappings = parse_almanac(data)
    return min(
        map(
            itemgetter(0),
            reduce(
                remap_almanac,
                mappings,
                ((x, x + y) for x, y in zip(*[iter(seeds)] * 2)),
            ),
        )
    )


if __name__ == "__main__":
    # get puzzle and parse data
    puzzle = Puzzle(year=2023, day=5)

    answer_a = part_a(puzzle.input_data)
    print(f"{answer_a}")
    submit(answer_a, part="a", day=5, year=2023)

    # part b:
    answer_b = part_b(puzzle.input_data)
    print(f"{answer_b}")
    submit(answer_b, part="b", day=5, year=2023)