"""Main weighted-RDPG sweep (paper Sec. 4.1 design, Theorems 3.1 / 3.2).

One full eigendecomposition per replicate yields, simultaneously:
  * (2,inf) estimation error for every embedding dimension d = r + k,
  * the Lemma 2.1 two-term decomposition (base term + trailing term),
  * the Theorem 3.1 delocalization statistic max_{alpha>r} |u_hat_{j alpha}|.
"""
import json
import sys
import time
import numpy as np
import rdpg

R = 5
DIMS = [1, 2, 3, 4, 5, 6, 7, 10, 20, 40]
NGRID = [250, 500, 1000, 2000, 4000, 8000]
REPS = {250: 24, 500: 24, 1000: 16, 2000: 10, 4000: 5, 8000: 3}
REPS_ALT = {250: 16, 500: 16, 1000: 10, 2000: 6, 4000: 3, 8000: 2}


def one_rep(n, rng, kind, scale, rho, binary=False, sbm=False):
    if sbm:
        A, Xt, _ = rdpg.sbm(n, R, rng)
    elif binary:
        A, Xt = rdpg.binary_rdpg(n, R, rng, rho=rho)
    else:
        A, Xt = rdpg.weighted_rdpg(n, R, rng, rho=rho, kind=kind, scale=scale)
    s, U = rdpg.full_spectrum(A)
    rec = {"n": n}
    # --- Theorem 3.1: delocalization of every eigenvector with alpha > r
    trail = np.abs(U[:, R:])
    rec["deloc_max_all"] = float(trail.max())
    rec["deloc_max_rp1"] = float(np.abs(U[:, R]).max())
    rec["deloc_max_top40"] = float(np.abs(U[:, R:40]).max())
    rec["deloc_ipr_rp1"] = float((U[:, R] ** 4).sum())        # inverse participation ratio
    rec["signal_2inf"] = float(rdpg.two_inf(U[:, :R]))
    rec["s_hat"] = [float(x) for x in s[:41]]
    # --- Lemma 2.1 / Theorem 3.2: estimation error at every embedding dimension
    Xh_r = rdpg.ase_from_spectrum(s, U, R)
    base, _ = rdpg.err_2inf(Xh_r, Xt)
    rec["base_2inf"] = base
    rec["err"], rec["trail_2inf"], rec["decomp_slack"] = {}, {}, {}
    for d in DIMS:
        Xh = rdpg.ase_from_spectrum(s, U, d)
        e, _ = rdpg.err_2inf(Xh, Xt)
        rec["err"][str(d)] = e
        if d > R:
            t = rdpg.two_inf(Xh[:, R:d])
            rec["trail_2inf"][str(d)] = t
            rec["decomp_slack"][str(d)] = base + t - e      # must be >= 0 (Lemma 2.1)
    return rec


def sweep(tag, kind="normal", scale=1.0, rho=1.0, binary=False, sbm=False,
          ngrid=None, reps=None, seed=1):
    ngrid = ngrid or NGRID
    reps = reps or REPS
    rng = np.random.default_rng(seed)
    out = []
    for n in ngrid:
        t0 = time.time()
        for _ in range(reps[n]):
            out.append(one_rep(n, rng, kind, scale, rho, binary, sbm))
        print(f"  {tag} n={n} reps={reps[n]}  {time.time()-t0:.1f}s", flush=True)
    return out


if __name__ == "__main__":
    which = sys.argv[1]
    jobs = {
        # PRIMARY grid: sigma = 0.1 (the noise level of the paper's own Eq. 16
        # experiment, N(0, 0.1^2)).  This keeps every signal eigenvalue
        # s_j ~ n/30 far above the BBP/Wigner detection threshold sigma sqrt(n)
        # over the whole n grid (ratio 5.3 at n=250 up to 29.8 at n=8000).
        "normal":      dict(tag="normal", kind="normal", scale=0.1, seed=11, reps=REPS),
        "laplace":     dict(tag="laplace", kind="laplace", scale=0.1, seed=12, reps=REPS_ALT),
        "exponential": dict(tag="exponential", kind="exponential", scale=0.1, seed=13, reps=REPS_ALT),
        "poisson":     dict(tag="poisson", kind="poisson", scale=0.1, seed=14, reps=REPS_ALT),
        # FIG-1 replication at the paper's unit noise scale
        "normal1":      dict(tag="normal1", kind="normal", scale=1.0, seed=21, reps=REPS_ALT),
        "laplace1":     dict(tag="laplace1", kind="laplace", scale=1.0, seed=22, reps=REPS_ALT),
        "exponential1": dict(tag="exponential1", kind="exponential", scale=1.0, seed=23, reps=REPS_ALT),
        "poisson1":     dict(tag="poisson1", kind="poisson", scale=1.0, seed=24, reps=REPS_ALT),
        # binary networks (Conjecture 1) and the heavy-tailed negative control
        "binary":      dict(tag="binary", binary=True, seed=15, reps=REPS_ALT),
        "sbm":         dict(tag="sbm", sbm=True, seed=16, reps=REPS_ALT),
        "cauchy":      dict(tag="cauchy", kind="cauchy", scale=0.1, seed=17, reps=REPS_ALT),
    }
    j = jobs[which]
    tag = j.pop("tag")
    res = sweep(tag, **j)
    with open(f"outputs/main_{tag}.json", "w") as f:
        json.dump({"tag": tag, "r": R, "dims": DIMS, "runs": res}, f)
    print("wrote", tag, len(res))