import json
import time
import numpy as np
from scipy.optimize import linprog


def transport_lp(C, mu=None, nu=None, edges=None):
    m, n = C.shape
    mu = np.ones(m) / m if mu is None else np.asarray(mu, float)
    nu = np.ones(n) / n if nu is None else np.asarray(nu, float)
    if edges is None:
        edges = [(i, j) for i in range(m) for j in range(n)]
    edges = list(edges)
    A = np.zeros((m + n, len(edges)))
    for t, (i, j) in enumerate(edges):
        A[i, t] = 1.; A[m + j, t] = 1.
    res = linprog(np.array([C[i, j] for i, j in edges]), A_eq=A[:-1],
                  b_eq=np.r_[mu, nu][:-1], bounds=(0, None), method="highs")
    if not res.success:
        return None
    X = np.zeros((m, n)); X[tuple(np.array(edges).T)] = res.x
    return X, edges, float(res.fun)


def coarse_to_fine(C, groups, max_rounds=40, add_per_round=2, tol=1e-9):
    m, n = C.shape; gm, gn = len(groups[0]), len(groups[1])
    CG = np.array([[C[np.ix_(I, J)].mean() for J in groups[1]] for I in groups[0]])
    coarse = transport_lp(CG, np.ones(gm)/gm, np.ones(gn)/gn)
    active = {(i, j) for a, I in enumerate(groups[0]) for b, J in enumerate(groups[1])
              if coarse[0][a, b] > tol for i in I for j in J}
    history = []
    full = transport_lp(C)
    for it in range(max_rounds + 1):
        ans = transport_lp(C, edges=sorted(active))
        if ans is not None:
            gap = ans[2] - full[2]
            history.append({"round": it, "edges": len(active), "objective": ans[2],
                            "gap": float(gap), "residual": 0.0})
            if gap <= 1e-9: return ans[0], active, history
        # Add cheapest edge in every uncovered row/column: a simple local
        # reduced-cost proxy that keeps the active-set procedure readable.
        cand = []
        for i in range(m):
            for j in np.argsort(C[i]):
                if (i, int(j)) not in active: cand.append((C[i, j], i, int(j))); break
        for j in range(n):
            for i in np.argsort(C[:, j]):
                if (int(i), j) not in active: cand.append((C[i, j], int(i), j)); break
        for _, i, j in sorted(cand)[:add_per_round]: active.add((i, j))
    return ans[0], active, history


def reduced_condition(z, edges, m, n):
    # Remove one redundant column constraint, preserving the paper's Schur form.
    D = np.zeros((m, n-1)); U = np.zeros(m); V = np.zeros(n-1)
    for val, (i, j) in zip(z, edges):
        U[i] += val*val
        if j < n-1: D[i, j] += val*val; V[j] += val*val
    if np.any(U <= 0) or np.any(V <= 0): return np.inf, np.nan
    Q = D / np.sqrt(U[:, None] * V[None, :])
    sigma = float(np.linalg.svd(Q, compute_uv=False)[0])
    return (1+sigma)/(1-sigma), sigma


def clustered_cost(m=12, separation=0.0, within_max=0.2):
    # Deterministic, reproducible two-block costs. Within-block costs span
    # [0, within_max], cross-block costs span [separation, separation+within_max].
    C = np.zeros((m, m)); h = m//2
    for i in range(m):
        for j in range(m):
            same = (i < h) == (j < h)
            if same: C[i,j] = within_max * ((3*i + 5*j) % 11) / 10
            else: C[i,j] = separation + within_max * ((7*i + 2*j + 1) % 11) / 10
    return C


def run():
    rng = np.random.default_rng(7); m = n = 8
    edges = sorted({(i,j) for i in range(m) for j in range(n) if rng.random()<.42} | {(i,i) for i in range(m)})
    z = rng.lognormal(0, .7, len(edges))
    scales = [{"scale": a, "kappa": reduced_condition(a*z, edges,m,n)[0],
               "sigma": reduced_condition(a*z, edges,m,n)[1]} for a in [.1,1,10]]
    # Prediction 1: global z scaling changes neither Q nor kappa.
    scale_rel = max(x["kappa"] for x in scales)-min(x["kappa"] for x in scales)
    # Prediction 2: kappa is exactly (1+sigma)/(1-sigma), tested over supports.
    relation = []
    for p in [.15,.3,.5,.7,.9]:
        ee = sorted({(i,j) for i in range(m) for j in range(n-1) if rng.random()<p} | {(i,n-1) for i in range(m)} | {(0,j) for j in range(n-1)})
        k,s = reduced_condition(np.ones(len(ee)), ee,m,n)
        relation.append({"sigma":s, "kappa":k, "formula_kappa":(1+s)/(1-s), "abs_error":abs(k-(1+s)/(1-s))})
    # Prediction 3: if cross costs exceed every within cost (sep > within_max),
    # the lifted diagonal block support is globally optimal; below it, expansion
    # can be required.
    groups=([list(range(6)),list(range(6,12))],[list(range(6)),list(range(6,12))])
    transition=[]
    for sep in np.linspace(0,.4,9):
        C=clustered_cost(separation=float(sep),within_max=.2)
        full=transport_lp(C); X,A,H=coarse_to_fine(C,groups)
        transition.append({"separation":float(sep),"predicted_exact":bool(sep>.2),
                           "initial_gap":None if not H else float(H[0]["gap"]),
                           "final_gap":None if not H else float(H[-1]["gap"]),
                           "rounds":len(H)-1,"active_fraction":len(A)/(12*12)})
    C=clustered_cost(separation=.4,within_max=.2); full=transport_lp(C); t=time.perf_counter()
    X,A,H=coarse_to_fine(C,groups); elapsed=time.perf_counter()-t
    top={(i,int(j)) for i in range(12) for j in np.argsort(C[i])[:4]}; topans=transport_lp(C,edges=sorted(top))
    out={"predictions":{"scale_invariance_max_kappa_range":scale_rel,
        "schur_formula_max_abs_error":max(x["abs_error"] for x in relation),
        "support_threshold_predicted_separation_gt":.2},
        "scale_sweep":scales,"schur_relation_sweep":relation,"support_transition":transition,
        "comparison":{"dense_objective":full[2],"idea_objective":H[-1]["objective"],
          "idea_edges":len(A),"idea_edge_fraction":len(A)/144,"idea_rounds":len(H)-1,
          "idea_seconds":elapsed,"topk_edges":len(top),"topk_feasible":topans is not None,
          "topk_objective":None if topans is None else topans[2]}}
    with open("results.json","w") as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=="__main__": run()