"""Targeted sweeps and negative controls.

  ksweep   - paper Fig. 2:  error vs k at fixed n (predicted sqrt(k) growth)
  sigma    - sigma-dependence of the over-specification term (Theorem 3.2)
  rho      - rho-dependence of the under-specification floor (Theorem 3.2, k<0)
  ktrue    - |k|-dependence of the floor vs the EXACT Lemma 2.1 bound
  rankctrl - negative control: same d, but the discarded direction carries no signal
  sparse   - sparse binary RDPG, rho_n = n^{-gamma} (Conjecture 1 / Sec. 4.2)
  hubctrl  - negative control for Conjecture 1: variance condition c <= E|E|^2 broken
"""
import json
import sys
import numpy as np
import rdpg

R = 5


def _ase_err(A, Xt, dims, r=R):
    s, U = rdpg.full_spectrum(A)
    out = {}
    for d in dims:
        Xh = rdpg.ase_from_spectrum(s, U, d)
        e, _ = rdpg.err_2inf(Xh, Xt)
        out[d] = {"err": e}
        if d > r:
            out[d]["trail"] = rdpg.two_inf(Xh[:, r:d])
            out[d]["U_trail_2inf"] = rdpg.two_inf(U[:, r:d])
            out[d]["s_hat_d"] = float(abs(s[d - 1]))
    return out, s, U


def ksweep(rng, reps=8):
    """Paper Fig. 2: exponential noise, error as a function of embedding dim."""
    dims = [3, 4, 5, 6, 7, 8, 10, 14, 20, 30, 40, 60, 80]
    res = {}
    for n in [1000, 2000, 4000]:
        acc = {d: [] for d in dims}
        tacc = {d: [] for d in dims if d > R}
        for _ in range(reps):
            A, Xt = rdpg.weighted_rdpg(n, R, rng, kind="exponential", scale=0.1)
            o, _, _ = _ase_err(A, Xt, dims)
            for d in dims:
                acc[d].append(o[d]["err"])
                if d > R:
                    tacc[d].append(o[d]["trail"])
        res[n] = {"dims": dims,
                  "err_mean": [float(np.mean(acc[d])) for d in dims],
                  "err_se": [float(np.std(acc[d], ddof=1) / np.sqrt(reps)) for d in dims],
                  "trail_mean": {str(d): float(np.mean(tacc[d])) for d in dims if d > R}}
        ks = [d - R for d in dims if d > R]
        res[n]["k_slope_err"] = rdpg.loglog_slope(
            ks, [np.mean(acc[d]) for d in dims if d > R])
        res[n]["k_slope_trail"] = rdpg.loglog_slope(
            ks, [np.mean(tacc[d]) for d in dims if d > R])
        res[n]["argmin_dim"] = int(dims[int(np.argmin([np.mean(acc[d]) for d in dims]))])
    return res


def sigma(rng, reps=8):
    """Theorem 3.2 writes the over-specification term as sqrt(sigma^2 k)(...)/n^{1/4},
    i.e. sigma^1.  Lemma 2.1 builds that term as ||Uhat_{r+1:r+k}||_{2,inf}*||E||^{1/2},
    and ||E|| ~ 2 sigma sqrt(n), which predicts sigma^{1/2}.  We measure the exponent."""
    n, d = 2000, 10
    sigmas = [0.025, 0.05, 0.1, 0.2, 0.4, 0.8]
    T, EV, SV, EX = [], [], [], []
    for sg in sigmas:
        t, ev, sv, ex = [], [], [], []
        for _ in range(reps):
            A, Xt = rdpg.weighted_rdpg(n, R, rng, kind="normal", scale=sg)
            o, s, U = _ase_err(A, Xt, [R, d])
            t.append(o[d]["trail"])
            ev.append(o[d]["U_trail_2inf"])
            sv.append(np.sqrt(o[d]["s_hat_d"]))
            ex.append(o[d]["err"] - o[R]["err"])
        T.append(float(np.mean(t))); EV.append(float(np.mean(ev)))
        SV.append(float(np.mean(sv))); EX.append(float(np.mean(ex)))
    return {"n": n, "d": d, "k": d - R, "sigmas": sigmas,
            "trail_mean": T, "U_trail_2inf_mean": EV, "sqrt_s_hat_mean": SV,
            "excess_err_mean": EX,
            "exp_trail": rdpg.loglog_slope(sigmas, T),
            "exp_U": rdpg.loglog_slope(sigmas, EV),
            "exp_sqrt_s": rdpg.loglog_slope(sigmas, SV),
            "exp_excess": rdpg.loglog_slope(sigmas, EX),
            "bbp_ratio": [float(n / 30 / (sg * np.sqrt(n))) for sg in sigmas]}


def rho(rng, reps=8):
    """Theorem 3.2 (k<0): floor >~ sqrt(|k| rho_n).  Sweep rho at fixed n, |k|=2."""
    n, d = 4000, 3
    rhos = [1.0, 0.5, 0.25, 0.125, 0.0625, 0.03125]
    E, B = [], []
    for rh in rhos:
        e, b = [], []
        for _ in range(reps):
            X = rdpg.dirichlet_latent(n, R, rng)
            A, Xt = rdpg.weighted_rdpg(n, R, rng, rho=rh, kind="normal",
                                       scale=0.02, X=X)
            o, _, _ = _ase_err(A, Xt, [d])
            s_nz = np.sort(np.linalg.eigvalsh(rh * (X.T @ X)))[::-1]
            e.append(o[d]["err"])
            b.append(float(np.sqrt(s_nz[d:].sum() / n)))
        E.append(float(np.mean(e))); B.append(float(np.mean(b)))
    return {"n": n, "d": d, "k": d - R, "rhos": rhos, "floor_mean": E,
            "lemma21_bound_mean": B,
            "ratio": [float(a / b) for a, b in zip(E, B)],
            "sqrt_k_rho": [float(np.sqrt(2 * rh)) for rh in rhos],
            "exp_floor": rdpg.loglog_slope(rhos, E),
            "exp_bound": rdpg.loglog_slope(rhos, B)}


def ktrue(rng, reps=8):
    """|k|-dependence of the under-specification floor vs the exact Lemma 2.1 bound."""
    n = 4000
    dims = [1, 2, 3, 4]
    E = {d: [] for d in dims}
    B = {d: [] for d in dims}
    for _ in range(reps):
        X = rdpg.dirichlet_latent(n, R, rng)
        A, Xt = rdpg.weighted_rdpg(n, R, rng, kind="normal", scale=0.1, X=X)
        o, _, _ = _ase_err(A, Xt, dims)
        s_nz = np.sort(np.linalg.eigvalsh(X.T @ X))[::-1]
        for d in dims:
            E[d].append(o[d]["err"])
            B[d].append(float(np.sqrt(s_nz[d:].sum() / n)))
    return {"n": n, "dims": dims,
            "floor_mean": [float(np.mean(E[d])) for d in dims],
            "lemma21_bound_mean": [float(np.mean(B[d])) for d in dims],
            "ratio": [float(np.mean(E[d]) / np.mean(B[d])) for d in dims],
            "sqrt_k_rho": [float(np.sqrt(R - d)) for d in dims]}


def rankctrl(rng, reps=8):
    """NEGATIVE CONTROL for the k<0 lower bound.  At the SAME embedding
    dimension d=4 we compare (i) a true rank-5 P (one signal direction is
    discarded -> floor) with (ii) a true rank-4 P (nothing is discarded ->
    no floor).  If a floor appeared in (ii) as well, the effect would be an
    artefact of the dimension count rather than of discarded signal."""
    dims = [4]
    out = {"n": [], "rank5_d4": [], "rank4_d4": [], "rank5_d5": []}
    for n in [500, 1000, 2000, 4000]:
        a, b, c = [], [], []
        for _ in range(reps):
            X5 = rdpg.dirichlet_latent(n, 5, rng)
            A5, Xt5 = rdpg.weighted_rdpg(n, 5, rng, kind="normal", scale=0.1, X=X5)
            o5, _, _ = _ase_err(A5, Xt5, [4, 5])
            X4 = rdpg.dirichlet_latent(n, 4, rng)
            A4, Xt4 = rdpg.weighted_rdpg(n, 4, rng, kind="normal", scale=0.1, X=X4)
            o4, _, _ = _ase_err(A4, Xt4, [4], r=4)
            a.append(o5[4]["err"]); b.append(o4[4]["err"]); c.append(o5[5]["err"])
        out["n"].append(n)
        out["rank5_d4"].append(float(np.mean(a)))
        out["rank4_d4"].append(float(np.mean(b)))
        out["rank5_d5"].append(float(np.mean(c)))
    for key in ["rank5_d4", "rank4_d4", "rank5_d5"]:
        out["slope_" + key] = rdpg.loglog_slope(out["n"], out[key])
    return out


def sparse(rng, reps=6):
    """Sparse binary RDPG, rho_n = n^{-gamma} (paper Eq. 18).  The k<0 floor is
    predicted to be sqrt(|k| rho_n) ~ n^{-gamma/2}: a *decaying* floor whose
    slope is set by gamma.  This is a sharp, falsifiable prediction."""
    dims = [3, 4, 5, 6, 7, 10, 20]
    ngrid = [500, 1000, 2000, 4000, 8000]
    res = {}
    for gam in [0.0, 0.25, 0.5]:
        per = {d: [] for d in dims}
        bnd, dl = [], []
        for n in ngrid:
            acc = {d: [] for d in dims}
            bb, dd = [], []
            for _ in range(reps):
                X = rdpg.dirichlet_latent(n, R, rng)
                rh = n ** (-gam)
                A, Xt = rdpg.binary_rdpg(n, R, rng, rho=rh, X=X)
                o, s, U = _ase_err(A, Xt, dims)
                for d in dims:
                    acc[d].append(o[d]["err"])
                s_nz = np.sort(np.linalg.eigvalsh(rh * (X.T @ X)))[::-1]
                bb.append(float(np.sqrt(s_nz[3:].sum() / n)))
                dd.append(float(np.abs(U[:, R]).max()))
            for d in dims:
                per[d].append(float(np.mean(acc[d])))
            bnd.append(float(np.mean(bb)))
            dl.append(float(np.mean(dd)))
        res[str(gam)] = {"ngrid": ngrid,
                         "err": {str(d): per[d] for d in dims},
                         "lemma21_bound_d3": bnd,
                         "deloc_rp1": dl,
                         "slopes": {str(d): rdpg.loglog_slope(ngrid, per[d]) for d in dims},
                         "slope_bound_d3": rdpg.loglog_slope(ngrid, bnd),
                         "slope_deloc": rdpg.loglog_slope(ngrid, dl),
                         "predicted_floor_slope": -gam / 2}
    return res


def hubctrl(rng, reps=6):
    """NEGATIVE CONTROL for Conjecture 1.  The conjecture relaxes Assumption A7
    to c <= E|E_ij|^2 <= C, i.e. entry variances bounded AWAY FROM ZERO.  A
    bounded-degree binary graph (rho_n = c/n) violates the lower bound c, and
    there eigenvector localisation is expected.  We compare the delocalisation
    statistic in the dense (conjecture-compliant) and bounded-degree regimes."""
    ngrid = [1000, 2000, 4000]
    out = {"ngrid": ngrid, "dense": [], "bounded_degree": [],
           "dense_ipr": [], "bounded_degree_ipr": []}
    for n in ngrid:
        a, b, ai, bi = [], [], [], []
        for _ in range(reps):
            A, _ = rdpg.binary_rdpg(n, R, rng, rho=1.0)
            s, U = rdpg.full_spectrum(A)
            a.append(float(np.abs(U[:, R]).max())); ai.append(float((U[:, R] ** 4).sum()))
            A, _ = rdpg.binary_rdpg(n, R, rng, rho=12.0 / n)
            s, U = rdpg.full_spectrum(A)
            b.append(float(np.abs(U[:, R]).max())); bi.append(float((U[:, R] ** 4).sum()))
        out["dense"].append(float(np.mean(a)))
        out["bounded_degree"].append(float(np.mean(b)))
        out["dense_ipr"].append(float(np.mean(ai)))
        out["bounded_degree_ipr"].append(float(np.mean(bi)))
    out["slope_dense"] = rdpg.loglog_slope(ngrid, out["dense"])
    out["slope_bounded_degree"] = rdpg.loglog_slope(ngrid, out["bounded_degree"])
    out["slope_dense_ipr"] = rdpg.loglog_slope(ngrid, out["dense_ipr"])
    out["slope_bd_ipr"] = rdpg.loglog_slope(ngrid, out["bounded_degree_ipr"])
    return out


def decompctrl(rng, reps=8):
    """NEGATIVE CONTROL for the two-term decomposition of Lemma 2.1.  We rebuild
    the d-dimensional embedding with its trailing columns SET TO ZERO,
    [Xhat_{1:r} | 0], and re-solve the same O_d Procrustes problem.  If the
    n^{-1/4} degradation really comes from the trailing block, this ablated
    embedding must fall back exactly onto the n^{-1/2} base curve."""
    dims = [10, 20]
    out = {"ngrid": [], "base": [], "full": {str(d): [] for d in dims},
           "ablated": {str(d): [] for d in dims}}
    for n in [500, 1000, 2000, 4000]:
        b, f, a = [], {d: [] for d in dims}, {d: [] for d in dims}
        for _ in range(reps):
            A, Xt = rdpg.weighted_rdpg(n, R, rng, kind="normal", scale=0.1)
            s, U = rdpg.full_spectrum(A)
            Xr = rdpg.ase_from_spectrum(s, U, R)
            b.append(rdpg.err_2inf(Xr, Xt)[0])
            for d in dims:
                Xh = rdpg.ase_from_spectrum(s, U, d)
                f[d].append(rdpg.err_2inf(Xh, Xt)[0])
                Xz = Xh.copy()
                Xz[:, R:] = 0.0
                a[d].append(rdpg.err_2inf(Xz, Xt)[0])
        out["ngrid"].append(n)
        out["base"].append(float(np.mean(b)))
        for d in dims:
            out["full"][str(d)].append(float(np.mean(f[d])))
            out["ablated"][str(d)].append(float(np.mean(a[d])))
    out["slope_base"] = rdpg.loglog_slope(out["ngrid"], out["base"])
    out["slope_full"] = {k: rdpg.loglog_slope(out["ngrid"], v) for k, v in out["full"].items()}
    out["slope_ablated"] = {k: rdpg.loglog_slope(out["ngrid"], v) for k, v in out["ablated"].items()}
    return out


if __name__ == "__main__":
    which = sys.argv[1]
    rng = np.random.default_rng(hash(which) % (2 ** 31))
    fn = {"ksweep": ksweep, "sigma": sigma, "rho": rho, "ktrue": ktrue,
          "rankctrl": rankctrl, "sparse": sparse, "hubctrl": hubctrl,
          "decompctrl": decompctrl}[which]
    res = fn(rng)
    with open(f"outputs/extra_{which}.json", "w") as f:
        json.dump(res, f, indent=1)
    print(which, "done")