import itertools
import math
import random
import time
import turtle
import typing
import collections.abc

Func: typing.TypeAlias = collections.abc.Callable
DistTable: typing.TypeAlias = list[list[float]]
IndSeq: typing.TypeAlias = list[int]

def tree_weight(x) -> int:
    if x is None:
        return 0
    elif isinstance(x, BinaryTree):
        return x.weight
    return 1

class BinaryTree:
    def __init__(self, a, b=None):
        self.x = self.y = self.parent = None
        self.a = a
        self.b = b
        self.weight = tree_weight(a) + tree_weight(b)
        self.usage = 0

    def __str__(self) -> str:
        return str(self.a) if self.b is None else f"({self.a} {self.b})"

    def merge(self, other):
        a, b = self, other
        while a.parent:
            a = a.parent
        while b.parent:
            b = b.parent
        if a is b:
            return a
        parent = BinaryTree(a, b)
        a.parent = b.parent = parent
        return parent

    def use(self):
        assert self.usage < self.weight
        self.usage += 1
        if self.parent:
            self.parent.use()

def distance_table[T](points: list[T], metric: Func[[T, T], float]) -> DistTable:
    return [[metric(x, y) for y in points] for x in points]

def show_mat(distances: DistTable) -> str:
    rows = []
    for i in range(len(distances)):
        first, last = i == 0, i == len(distances) - 1
        row = "/" if first else "\\" if last else "|"
        row += " ".join(f"{d:5.2f}" for d in distances[i])
        row += "\\" if first else "/" if last else "|"
        rows.append(row)
    return "\n".join(rows)

def score_nearest(distances: DistTable, sample: IndSeq) -> float:
    return sum(min(distances[i][j] for j in sample) for i in range(len(distances)))

def score_total(distances: DistTable, sample: IndSeq) -> float:
    return sum(sum(
        0 if i in sample else distances[i][j] for j in sample
    ) for i in range(len(distances)))

def score_spread(distances: DistTable, sample: IndSeq) -> float:
    return -sum(sum(distances[sample[i]][j] for j in sample[:i]) for i in range(len(sample)))

def greedy_min_seq(distances: DistTable, score_func: Func[[DistTable, IndSeq], float]) -> IndSeq:
    seq = []
    options = set(range(len(distances)))
    while options:
        best = min(options, key=lambda o: score_func(distances, seq + [o]))
        seq.append(best)
        options.remove(best)
    return seq

def furthest_nb(distances: DistTable) -> IndSeq:
    seq = []
    options = set(range(len(distances)))
    start = min(options, key=lambda o: score_nearest(distances, [o]))
    seq.append(start)
    options.remove(start)
    while options:
        dl = distances[seq[-1]]
        best = max(options, key=lambda o: dl[o])
        seq.append(best)
        options.remove(best)
    return seq

def distance_hierarchy(distances: DistTable) -> BinaryTree:
    n = len(distances)
    forest = [BinaryTree(i) for i in range(n)]
    p = None
    for (i, j) in sorted(
        ((i, j) for i in range(n) for j in range(n)),
        key=lambda ij: distances[ij[0]][ij[1]]
    ):
        p = forest[i].merge(forest[j])
    return p

def scattered_hierarchy(hierarchy: BinaryTree) -> IndSeq:
    seq = []
    while hierarchy.usage < hierarchy.weight:
        fb = hierarchy
        while fb.b is not None:
            ua = fb.a.usage / fb.a.weight
            ub = fb.b.usage / fb.b.weight
            if ua < ub:
                fb = fb.a
            elif ua == ub and fb.a.weight > fb.b.weight:
                fb = fb.a
            else:
                fb = fb.b
        fb.use()
        seq.append(fb.a)
    return seq

def total_len(distances: DistTable, seq: IndSeq) -> float:
    return sum(distances[i][j] for (i, j) in itertools.pairwise(seq))

N = 100
points: list[tuple[float, float]] = [(random.randint(0, N // 2), random.randint(0, N // 3)) for _ in range(N)]
print(points)
t: turtle.Turtle = turtle.Turtle()
t.hideturtle()
t.pen(speed=10)
dt: DistTable = distance_table(points, math.dist)
print(show_mat(dt))
h: BinaryTree = distance_hierarchy(dt)
print(h)
METHODS = [
    ("graph", lambda: scattered_hierarchy(h)),
    ("nearest-nb sum", lambda: greedy_min_seq(dt, score_nearest)),
    ("all dist sum", lambda: greedy_min_seq(dt, score_total)),
    ("sample spread", lambda: greedy_min_seq(dt, score_spread)),
    ("furthest-nb", lambda: furthest_nb(dt)),
]
input("ready?")
for (name, method) in METHODS:
    to = time.time()
    seq = method()
    tf = time.time()
    l = total_len(dt, seq)
    print(f"{name} method took {tf - to} s to find {seq}, length {l}")
    t.clear()
    for (i, (x, y)) in enumerate(points[k] for k in seq):
        t.teleport(2000 / N * (x - 0.25 * N), 2000 / N * (y - 0.17 * N))
        t.dot(20 / math.sqrt(i + 1))
        t.write(i)
    input("next?")