import json, math, random, time
from collections import defaultdict, deque
import numpy as np

SEED = 2954
rng = np.random.default_rng(SEED)


def build_graph(n, kind, density=0.35):
    # Bipartite vertices are x=0..n-1 and y=n..2n-1.
    edges = []
    if kind == "chain":
        # A connected sparse graph with many local cycles.
        for x in range(n):
            edges.append((x, x))
            if x + 1 < n: edges.append((x, x + 1))
            if x + 2 < n: edges.append((x, x + 2))
    elif kind == "grid":
        # Dense enough connected random support, with reproducible seed.
        for x in range(n):
            for y in range(n):
                if (x == y or (x + 1) % n == y or (x + 2) % n == y or rng.random() < density):
                    edges.append((x, y))
    else:
        for x in range(n):
            for y in range(n):
                if rng.random() < density or (x == y): edges.append((x, y))
    # remove duplicates and isolate-free components are handled by forest
    return list(dict.fromkeys(edges))


def fundamental_basis(n, edges):
    """Return cycle coefficient vectors and cycle rank for a bipartite graph.
    Each vector gives signed edge coefficients for the oriented edge x -> y.
    """
    m = len(edges)
    adj = defaultdict(list)
    for i, (x, y) in enumerate(edges):
        u, v = x, n + y
        adj[u].append((v, i, 1))       # traversing x -> y agrees with edge orientation
        adj[v].append((u, i, -1))      # traversing y -> x reverses it
    parent = {}
    parent_edge = {}
    parent_sign = {}
    component = {}
    tree = set()
    for root in list(range(2*n)):
        if root in component or not adj[root]: continue
        component[root] = root
        stack = [root]
        while stack:
            u = stack.pop()
            for v, ei, sign in adj[u]:
                if v not in component:
                    component[v] = root
                    parent[v] = u
                    parent_edge[v] = ei
                    parent_sign[v] = sign
                    tree.add(ei)
                    stack.append(v)
    non_tree = [i for i in range(m) if i not in tree]
    cycles = []
    # Return the tree path start -> end with signs relative to x -> y.
    def tree_path(start, end):
        anc = {}
        u = start
        while True:
            anc[u] = True
            if u not in parent: break
            u = parent[u]
        start_up = []
        u = start
        end_up = []
        u = end
        while u not in anc:
            end_up.append((parent_edge[u], parent_sign[u]))
            u = parent[u]
        lca = u
        start_up = []
        u = start
        while u != lca:
            start_up.append((parent_edge[u], -parent_sign[u]))
            u = parent[u]
        end_down = [(ei, sign) for ei, sign in reversed(end_up)]
        return start_up + end_down

    cycles = []
    for ei in non_tree:
        x, y = edges[ei]
        coeff = np.zeros(m, dtype=float)
        coeff[ei] = 1.0  # x -> y
        for ej, sign in tree_path(n + y, x):
            coeff[ej] += sign
        cycles.append(coeff)
    active_vertices = len(component)
    comps = len(set(component.values()))
    rank = m - active_vertices + comps
    return np.asarray(cycles), rank


def exhaustive_simple_cycles(n, edges, cap=5000):
    """Enumerate oriented simple-cycle coefficient vectors, deduplicated."""
    adj = defaultdict(list)
    for i, (x, y) in enumerate(edges):
        u, v = x, n + y
        adj[u].append((v, i, 1))
        adj[v].append((u, i, -1))
    found = {}
    def dfs(start, u, visited, coeff):
        if len(found) >= cap:
            return
        for v, ei, sign in adj[u]:
            if v == start and len(visited) >= 4:
                key = tuple(sorted(np.flatnonzero(coeff).tolist()))
                if key not in found:
                    c = coeff.copy()
                    c[ei] += sign
                    found[key] = c
            elif v not in visited and v >= start:
                coeff[ei] += sign
                dfs(start, v, visited | {v}, coeff)
                coeff[ei] -= sign
    for start in range(2*n):
        dfs(start, start, {start}, np.zeros(len(edges)))
    return list(found.values())


def main():
    np.set_printoptions(precision=6, suppress=True)
    report = {"seed": SEED, "graphs": []}
    for kind,n in [("chain",12),("grid",9),("random",12)]:
        edges=build_graph(n,kind)
        basis,rank=fundamental_basis(n,edges)
        # Compatible field a = node potential on x minus node potential on y.
        ux=rng.normal(size=n); vy=rng.normal(size=n)
        a=np.array([ux[x]-vy[y] for x,y in edges])
        assert np.max(np.abs(basis @ a)) < 1e-10, "invalid fundamental-cycle orientation"
        compatible_max=float(np.max(np.abs(basis @ a))) if len(basis) else 0.0
        # Perturbed field: basis residual is the exact independent constraint vector.
        noise=rng.normal(scale=.1,size=len(edges)); ap=a+noise
        bres=basis@ap
        # Rank and cycle-space reconstruction: random cycle vectors are row-space combinations.
        if len(basis):
            coeff=rng.normal(size=(min(20,len(basis)*2),len(basis)))
            held=coeff@bres
            held95=float(np.percentile(np.abs(held),95))
            basis95=float(np.percentile(np.abs(bres),95))
        else: held95=basis95=0.0
        exhaustive=exhaustive_simple_cycles(n,edges)
        exhaustive_compat_max = (max((abs(float(c @ a)) for c in exhaustive), default=0.0))
        exhaustive_perturbed_95 = (float(np.percentile([abs(float(c @ ap)) for c in exhaustive], 95))
                                   if exhaustive else 0.0)
        # Cost proxy is number of residual scalar sums, plus actual timing.
        reps=10
        t0=time.perf_counter()
        for _ in range(reps): _=basis@ap
        tb=time.perf_counter()-t0
        t0=time.perf_counter()
        # Exhaustive cycle residuals represented by summing edge values; count only.
        for _ in range(reps):
            _=[float(c @ ap) for c in exhaustive]
        te=time.perf_counter()-t0
        report["graphs"].append({"kind":kind,"n":n,"edges":len(edges),"cycle_rank":rank,
          "basis_constraints":len(basis),"simple_cycles":len(exhaustive),
          "compatible_max_basis_residual":compatible_max,"exhaustive_compatible_max":exhaustive_compat_max,
          "perturbed_basis_95":basis95,"exhaustive_perturbed_95":exhaustive_perturbed_95,
          "heldout_linear_combination_95":held95,"basis_time_sec":tb,"exhaustive_time_sec":te,
          "speedup":(te/tb if tb else None)})
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()