import numpy as np
from wedge_pool import all_pairs_dist, wedge_split, connected, sse
from collections import deque


def connected_graph(adj):
    seen={0}; q=deque([0])
    while q:
        v=q.popleft()
        for z in adj[v]:
            if z not in seen: seen.add(z); q.append(z)
    return len(seen)==len(adj)


def main():
    rng=np.random.default_rng(123)
    failures=[]; checks=0
    for trial in range(500):
        n=int(rng.integers(5,14))
        adj=[set() for _ in range(n)]
        # random connected backbone plus extra edges
        for i in range(1,n):
            j=int(rng.integers(i)); adj[i].add(j); adj[j].add(i)
        for i in range(n):
            for j in range(i):
                if rng.random()<.18: adj[i].add(j); adj[j].add(i)
        adj=[sorted(x) for x in adj]
        R=list(range(n)); D=all_pairs_dist(adj,R)
        for u in range(n):
            for w in range(u+1,n):
                A,B=wedge_split(adj,R,u,w,D)
                if A and B:
                    checks+=1
                    if not connected(adj,A) or not connected(adj,B):
                        failures.append((trial,n,u,w,A,B)); break
            if failures: break
        if failures: break
    # Independent mean identity: SSE around c = SSE around mean + |R| ||c-mu||^2.
    X=rng.normal(size=(17,3)); R=list(range(17)); c=rng.normal(size=3)
    cost,mu=sse(X,R); lhs=float(((X-c)**2).sum()); rhs=cost+len(R)*float(((c-mu)**2).sum())
    print({'random_graph_split_checks':checks,'connectivity_failures':len(failures),
           'first_failure':failures[:1], 'mean_identity_abs_error':abs(lhs-rhs),
           'mean_cost':cost,'arbitrary_cost':lhs})

if __name__=='__main__': main()