# This does not work for the sample :D

def solve(data: str):
    # split input
    blocks = data.split("\n\n")
    shape_blocks = blocks[:-1]
    region_block = blocks[-1]

    # for every shape, get the number of occupied cells
    shape_area = []
    for shape_block in shape_blocks:
        shape_area.append(shape_block.count('#'))

    fit_regions = 0

    # for every region, check if the area is sufficient to fit all shapes
    for region_data in region_block.splitlines():
        size_data, shape_data = region_data.split(': ')

        # get region size
        m, n = [int(dim) for dim in size_data.split('x')]

        # get area needed to fit all shapes, without considering arrangement
        area_needed = 0
        for id, freq in enumerate(map(int, shape_data.split(' '))):
            area_needed += shape_area[id] * freq
        
        # if the region area is sufficient, count it as a fit (!!!)
        if m * n > area_needed:
            fit_regions += 1
    
    return fit_regions

from collections import defaultdict

# build the adjacency graph from the input data
def build_graph(data: str):
    rules_data = data.splitlines()
    graph = defaultdict(list)
    for rule_data in rules_data:
        from_node, to_nodes = rule_data.split(": ")
        graph[from_node].extend(to_nodes.split(' '))
    return graph

# simply dfs every path from 'you' until we reach 'out'
# this is what I initially used for part 1
def count_paths_dfs(data: str, start = "you", end = "out"):
    graph = build_graph(data)

    paths_to_end = 0
    stack = [start]     # currently active paths
    while stack:
        node = stack.pop()
        if node == end:
            # we reached end, no need to check further
            paths_to_end += 1
            continue
        # every node connected to current node is a potential path
        stack.extend(graph[node])
    
    return paths_to_end

# memoized version of counting paths from start to end
def count_paths_memo(graph: defaultdict[list], start = "you", end = "out"):
    # I used to have some code to check for cycles, but thankfully it wasn't needed

    # its pretty much same logic as count_paths_dfs, except recursive and with memoization
    memo = {}
    def backtrack(curr: str):
        if curr in memo:
            return memo[curr]        
        if curr == end:
            ans = 1
        else:
            ans = sum(backtrack(conn) for conn in graph[curr])
        memo[curr] = ans
        return ans
    
    return backtrack(start)

# Optimized part 1 using memoization (not necessary, but faster)
def part1_with_memoiz(data: str, start = "you", end = "out"):
    graph = build_graph(data)
    return count_paths_memo(graph, start, end)

# Part 2 requires counting paths from "svr" to "out", 
#   including only those paths that go through both "fft" and "dac"
# Instead of adding booleans / int to my memoization function, 
#   I chose to compute all path arrangements and add them up
def part2(data: str):
    graph = build_graph(data)

    # getting from start to one requirement
    svr_to_fft = count_paths_memo(graph, "svr", "fft")
    svr_to_dac = count_paths_memo(graph, "svr", "dac")
    # getting from one requirement to the other
    fft_to_dac = count_paths_memo(graph, "fft", "dac")
    dac_to_fft = count_paths_memo(graph, "dac", "fft")
    # the final push to out
    fft_to_out = count_paths_memo(graph, "fft", "out")
    dac_to_out = count_paths_memo(graph, "dac", "out")

    # total paths = (start -> fft -> dac -> out) + (start -> dac -> fft -> out)
    svr_to_out = (
        svr_to_dac * dac_to_fft * fft_to_out +
        svr_to_fft * fft_to_dac * dac_to_out
    )
    return svr_to_out

sample_p1 = """aaa: you hhh
you: bbb ccc
bbb: ddd eee
ccc: ddd eee fff
ddd: ggg
eee: out
fff: out
ggg: out
hhh: ccc fff iii
iii: out"""
assert (t := count_paths_dfs(sample_p1)) == 5, f"Expected 5, got {t}"
assert (t := part1_with_memoiz(sample_p1)) == 5, f"Expected 5, got {t}"

sample_p2 = """svr: aaa bbb
aaa: fft
fft: ccc
bbb: tty
tty: ccc
ccc: ddd eee
ddd: hub
hub: fff
eee: dac
dac: fff
fff: ggg hhh
ggg: out
hhh: out"""
assert (t := part2(sample_p2)) == 2, f"Expected 2, got {t}"

# combinations helps us get all pairs of red tiles
# pairwise helps us get all adjacent pairs of red tiles
from itertools import combinations, pairwise
from collections import deque

def get_area(p1, p2):
    x1, y1 = p1
    x2, y2 = p2
    return (abs(x2 - x1) + 1) * (abs(y2 - y1) + 1)

@timer()
def part1(data: str):
    red_tiles = [tuple(map(int, line.split(','))) for line in data.splitlines()]

    max_area = 0
    for r1, r2 in combinations(red_tiles, 2):
        max_area = max(max_area, get_area(r1, r2))
    return max_area

def compress_coordinates(coords: list[tuple[int, int]]):
    i_of_interest_set = set()
    j_of_interest_set = set()

    for i, j in coords:
        i_of_interest_set.update([i-1, i, i+1])
        j_of_interest_set.update([j-1, j, j+1])
    
    sorted_is = sorted(list(i_of_interest_set))
    sorted_js = sorted(list(j_of_interest_set))
    m = len(sorted_is)
    n = len(sorted_js)

    compress_i_map = {val: id for id, val in enumerate(sorted_is)}
    compress_j_map = {val: id for id, val in enumerate(sorted_js)}
    return m, n, compress_i_map, compress_j_map


@timer()
def part2(data: str):
    red_tiles = [tuple(map(int, line.split(','))) for line in data.splitlines()]
    m, n, compress_i_map, compress_j_map = compress_coordinates(red_tiles)

    def get_compressed_coords(og_coords: tuple[int, int]):
        return compress_i_map[og_coords[0]], compress_j_map[og_coords[1]]

    # make the grid
    grid = [['.']*n for _ in range(m)]
    for i, j in red_tiles:
        ci, cj = get_compressed_coords((i, j))
        grid[ci][cj] = '#'

    # make the walls of the polygon enclosed by the red tiles
    def make_line(r1, r2):
        r1_i, r1_j = get_compressed_coords(r1)
        r2_i, r2_j = get_compressed_coords(r2)

        wall_char = '#'
        if r1_i == r2_i:
            # same row
            row = r1_i
            for col in range(min(r1_j, r2_j) + 1, max(r1_j, r2_j)):
                grid[row][col] = wall_char
        elif r1_j == r2_j:
            # same col
            col = r1_j
            for row in range(min(r1_i, r2_i) + 1, max(r1_i, r2_i)):
                grid[row][col] = wall_char

    for r1, r2 in pairwise(red_tiles):
        make_line(r1, r2)
    make_line(red_tiles[-1], red_tiles[0])

    # whiteout the exterior
    DIRECTIONS = [(0, 1), (1, 0), (0, -1), (-1, 0)]
    queue = deque([(0, 0)])
    while queue:
        i, j = queue.popleft()
        if grid[i][j] == ' ':
            continue
        grid[i][j] = ' '

        for di, dj in DIRECTIONS:
            ni, nj = i+di, j+dj
            if not (0 <= ni < m) or not (0 <= nj < n):
                continue
            if grid[ni][nj] != '.':
                continue
            queue.append((ni, nj))
    
    # DEBUG: print the grid
    # print()
    # for row in grid:
    #     print(''.join(row))
    
    # check if the rectangle formed by the corners r1, r2 is contained within the red-green polygon
    def check_contained(r1, r2):
        r1_i, r1_j = get_compressed_coords(r1)
        r2_i, r2_j = get_compressed_coords(r2)
        
        min_i, max_i = min(r1_i, r2_i), max(r1_i, r2_i)
        min_j, max_j = min(r1_j, r2_j), max(r1_j, r2_j)

        for i in range(min_i, max_i+1):
            for j in range(min_j, max_j+1):
                if grid[i][j] == ' ':
                    return False
        return True

    # get the max area of a rectangle that is contained within the red-green polygon
    #   and has red tiles at two of its corners
    max_area = 0
    for r1, r2 in combinations(red_tiles, 2):
        area = get_area(r1, r2)
        if area <= max_area:
            # dont bother
            continue
        if not check_contained(r1, r2):
            continue
        max_area = area

    return max_area

sample = """7,1
11,1
11,7
9,7
9,5
2,5
2,3
7,3"""
assert part1(sample) == 50
assert part2(sample) == 24

import heapq as hq
from itertools import combinations
from collections import Counter

# Disjoint Set Union (or Union-Find) data structure
#   with path compression and union by rank
class DSU:
    def __init__(self, size: int) -> None:
        self.parent = list(range(size)) # ith value is the parent node to node i
        self.rank = [1] * size          # max depth of subtree rooted here (used for union by rank)
        
    def find(self, x: int) -> int:
        # if the node is not its own parent, we need to recurse on parent
        if x != self.parent[x]:
            # path compression
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    
    def isConnected(self, x: int, y: int) -> bool:
        return self.find(x) == self.find(y)
    
    # returns a boolean whether or not union was needed
    def union(self, x: int, y: int) -> bool:
        rootX = self.find(x)
        rootY = self.find(y)        
        
        if rootX == rootY:
            # no union needed
            return False
        
        if self.rank[rootX] > self.rank[rootY]:
            # rootX has deeper subtree, therefore set it as parent to rootY (and its subtree)
            self.parent[rootY] = rootX
        elif self.rank[rootX] < self.rank[rootY]:
            # rootY has deeper subtree, therefore set it as parent to rootX (and its subtree)
            self.parent[rootX] = rootY
        else:
            # both subtrees are of equal depth, therefore choose either one of them as the parent of the other
            # here we chose rootX as the parent of rootY, therefore rootX depth increases by 1
            self.parent[rootY] = rootX
            self.rank[rootX] += 1
        
        # union complete
        return True

# parse puzzle input into junction coordinates (x, y, z)
def parse_junctions(data: str):
    return [tuple(map(int, line.split(','))) for line in data.splitlines()]

# get distance between two points in 3D space
# skipping the sqrt because we only care about relative distances
def get_dist(p1, p2):
    x1, y1, z1 = p1
    x2, y2, z2 = p2
    return (x2 - x1) ** 2 + (y2 - y1) ** 2 + (z2 - z1) ** 2

def part1(data: str, max_connections: int = 1000):
    junctions = parse_junctions(data)
    n = len(junctions)

    # <max_connections> shortest connections
    closest_connections = []
    # iterate over all pairs of junctions
    # combinations is just some syntactic sugar to avoid a nested for loop (still O(n^2))
    for i, j in combinations(range(n), 2):
        dist = get_dist(junctions[i], junctions[j])
        # keep <max_connections> closest connections
        #   have to use -dist to simulate max heap
        if len(closest_connections) < max_connections:
            hq.heappush(closest_connections, (-dist, i, j))
        else:
            hq.heappushpop(closest_connections, (-dist, i, j))

    # union all the closest connections
    dsu = DSU(n)
    for _, i, j in closest_connections:
        dsu.union(i, j)

    # count all circuit sizes
    circuit_sizes = Counter()
    for i in range(n):
        circuit_sizes[dsu.find(i)] += 1

    # multiply the sizes of the 3 largest circuits
    ans = 1
    for _, f in circuit_sizes.most_common(3):
        ans *= f
    return ans

def part2(data: str):
    junctions = parse_junctions(data)
    n = len(junctions)

    # all n^2 junction connections
    all_connections = []
    # combinations is just some syntactic sugar to avoid a nested for loop (still O(n^2))
    for i, j in combinations(range(n), 2):
        dist = get_dist(junctions[i], junctions[j])
        all_connections.append((dist, i, j))
    
    # create min heap of all connections
    # Heapify later to benefit from the heapify algorithm (O(n) instead of O(n log n))
    hq.heapify(all_connections)
    
    # connect junctions until all are connected
    dsu = DSU(n)
    unions = 0
    while all_connections:
        _, i, j = hq.heappop(all_connections)
        unions += dsu.union(i, j)
        # if we have n-1 unions, that means all n junctions are connected
        if unions == n-1:
            return junctions[i][0] * junctions[j][0]
    
    assert False, "It's not possible that all junctions never connect"

sample = """162,817,812
57,618,57
906,360,560
592,479,940
352,342,300
466,668,158
542,29,236
431,825,988
739,650,466
52,470,668
216,146,977
819,987,18
117,168,530
805,96,715
346,949,466
970,615,88
941,993,340
862,61,35
984,92,344
425,690,689"""
assert part1(sample, max_connections=10) == 40
assert part2(sample) == 25272