import math
import sys
from itertools import combinations

from vorlesung.L08_graphen.graph import AdjacencyListGraph


# ---------------------------------------------------------------------------
# b) Union-Find
# ---------------------------------------------------------------------------

class DisjointSet:
    """Union-Find mit Pfadkompression und Union-by-Size."""

    def __init__(self):
        self._parent = {}
        self._size = {}

    def make_set(self, v):
        self._parent[v] = v
        self._size[v] = 1

    def find_set(self, v):
        if self._parent[v] != v:
            self._parent[v] = self.find_set(self._parent[v])  # Pfadkompression
        return self._parent[v]

    def union(self, u, v):
        """Vereinigt die Schaltkreise von u und v.
        Gibt True zurück, wenn sie vorher getrennt waren."""
        ru, rv = self.find_set(u), self.find_set(v)
        if ru == rv:
            return False
        if self._size[ru] < self._size[rv]:
            ru, rv = rv, ru
        self._parent[rv] = ru
        self._size[ru] += self._size[rv]
        return True

    def size(self, v):
        return self._size[self.find_set(v)]

    def circuit_sizes(self):
        """Alle Schaltkreisgrößen, absteigend sortiert."""
        return sorted(
            [self._size[v] for v in self._parent if self._parent[v] == v],
            reverse=True
        )


# ---------------------------------------------------------------------------
# Hilfsfunktionen
# ---------------------------------------------------------------------------

def euclidean(a, b):
    return math.sqrt(sum((ai - bi) ** 2 for ai, bi in zip(a, b)))


def parse_boxes(text):
    boxes = []
    for line in text.strip().splitlines():
        line = line.strip()
        if line:
            boxes.append(tuple(int(x) for x in line.split(',')))
    return boxes


# ---------------------------------------------------------------------------
# c–e) Kruskal
# ---------------------------------------------------------------------------

def solve(text, k=1000):
    boxes = parse_boxes(text)
    n = len(boxes)
    print(f"Anzahl Boxen: {n}")
    print(f"Potenzielle Verbindungen: {n * (n - 1) // 2}")

    # a) Alle Paare nach Distanz sortiert
    edges = sorted(
        ((euclidean(a, b), a, b) for a, b in combinations(boxes, 2)),
        key=lambda e: e[0]
    )

    ds = DisjointSet()
    for b in boxes:
        ds.make_set(b)

    # --- Teil 1: die k kürzesten Paare verbinden (inkl. No-Ops) ---
    for _, u, v in edges[:k]:
        ds.union(u, v)
    idx = k

    sizes = ds.circuit_sizes()
    top3 = sizes[:3]
    result1 = top3[0] * top3[1] * top3[2]
    print(f"\n--- Teil 1 (k={k}) ---")
    print(f"Schaltkreisgrößen: {sizes}")
    print(f"Produkt der 3 größten: {top3[0]} × {top3[1]} × {top3[2]} = {result1}")

    # --- Teil 2: weiter bis ein einziger Schaltkreis ---
    last_u = last_v = None
    while idx < len(edges):
        _, u, v = edges[idx]
        idx += 1
        if ds.union(u, v):
            last_u, last_v = u, v
            if ds.size(last_u) == n:
                break

    result2 = last_u[0] * last_v[0]
    print(f"\n--- Teil 2 ---")
    print(f"Letzte Verbindung: {last_u} -- {last_v}")
    print(f"Produkt der X-Koordinaten: {last_u[0]} × {last_v[0]} = {result2}")

    # --- Verifikation mit graph.mst_kruskal() ---
    # mst_kruskal() liefert alle MST-Kanten aufsteigend nach Gewicht sortiert.
    # Die letzte Kante ist die schwerste im MST = die zuletzt hinzugefügte
    # Kante bei Kruskal = identisch mit dem Ergebnis aus Teil 2.
    # connect() ist gerichtet; beide Richtungen werden eingetragen, damit
    # all_edges() alle Paare sieht. Der Same-Set-Check in mst_kruskal()
    # filtert die Rückrichtung als No-Op heraus.
    g = AdjacencyListGraph()
    for b in boxes:
        g.insert_vertex(f"{b[0]},{b[1]},{b[2]}")
    for a, b in combinations(boxes, 2):
        name_a = f"{a[0]},{a[1]},{a[2]}"
        name_b = f"{b[0]},{b[1]},{b[2]}"
        d = euclidean(a, b)
        g.connect(name_a, name_b, d)
        g.connect(name_b, name_a, d)
    mst_edges, _ = g.mst_kruskal()
    last = mst_edges[-1]
    gu = tuple(int(x) for x in last[0].split(','))
    gv = tuple(int(x) for x in last[1].split(','))
    result2_graph = gu[0] * gv[0]
    assert result2_graph == result2, (
        f"Verifikation fehlgeschlagen: graph.mst_kruskal() → {result2_graph}, "
        f"eigene Lösung → {result2}"
    )
    print(f"Verifikation mit graph.mst_kruskal() bestanden.")

    return result1, result2


# ---------------------------------------------------------------------------
# Beispieldaten (AoC 2025, Tag 8 – 20 Junction Boxes)
# Erwartetes Ergebnis: Teil 1 (k=10) → 40, Teil 2 → 25272
# ---------------------------------------------------------------------------

BEISPIEL = """
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
"""


if __name__ == '__main__':
    if len(sys.argv) == 1:
        print("=== Selbsttest mit Beispieldaten ===")
        r1, r2 = solve(BEISPIEL, k=10)
        assert r1 == 40,    f"Teil 1 erwartet 40, erhalten {r1}"
        assert r2 == 25272, f"Teil 2 erwartet 25272, erhalten {r2}"
        print("Selbsttest bestanden.")
    else:
        filename = sys.argv[1]
        k = int(sys.argv[2]) if len(sys.argv) > 2 else 1000
        with open(filename) as f:
            text = f.read()
        solve(text, k)