"""Claim 6: the alpha-dependent tradeoff between drift and stability, on the
10-D Gaussian mixture. Low alpha (little real data re-injected each generation)
should drift/collapse; higher alpha should stabilise.

Previously the logbook only hash-pinned the paper's Figure 1 panels and said so.
The 10-D Gaussian-mixture half needs no image data, so it is run here.

Self-consuming loop: generation 0 fits a GMM to real samples; each later
generation trains on a mixture of `alpha` real data and `1-alpha` of its own
previous output, which is the setting model collapse is defined in.
"""
import json, numpy as np
from sklearn.mixture import GaussianMixture

RES = {}
D, K = 10, 4


def true_gmm(seed=0):
    rng = np.random.default_rng(seed)
    mu = rng.normal(size=(K, D))*2.5
    A = rng.normal(size=(K, D, D))*0.35
    cov = np.array([a @ a.T+0.35*np.eye(D) for a in A])
    w = rng.dirichlet(np.ones(K)*4)
    return mu, cov, w


def sample_true(mu, cov, w, n, rng):
    c = rng.choice(K, size=n, p=w)
    X = np.zeros((n, D))
    for k in range(K):
        m = c == k
        if m.sum():
            X[m] = rng.multivariate_normal(mu[k], cov[k], size=int(m.sum()))
    return X


def moment_error(X, mu, cov, w):
    """Distance to the TRUE distribution via first two moments (closed form)."""
    tm = (w[:, None]*mu).sum(0)
    tc = sum(w[k]*(cov[k]+np.outer(mu[k]-tm, mu[k]-tm)) for k in range(K))
    em = X.mean(0); ec = np.cov(X.T)
    return float(np.linalg.norm(em-tm)+np.linalg.norm(ec-tc, "fro"))


def run(alpha, gens=80, n=150, seed=0):
    """n is deliberately SMALL: with n=2000 the per-generation variance loss is
    ~1/n and collapse never manifests (variance stayed at 116% of true over 25
    generations). Collapse is driven by sampling noise compounding, so it needs
    small n and many generations."""
    rng = np.random.default_rng(seed)
    mu, cov, w = true_gmm(seed)
    real = sample_true(mu, cov, w, n, rng)
    model = GaussianMixture(K, covariance_type="full", reg_covar=1e-6,
                            random_state=0, max_iter=200).fit(real)
    errs = [moment_error(model.sample(n)[0], mu, cov, w)]
    traces = [float(np.trace(np.cov(model.sample(n)[0].T)))]
    for g in range(gens):
        gen = model.sample(n)[0]
        nreal = int(round(alpha*n))
        train = np.vstack([sample_true(mu, cov, w, nreal, rng), gen[:n-nreal]]) if nreal else gen
        model = GaussianMixture(K, covariance_type="full", reg_covar=1e-6,
                                random_state=0, max_iter=200).fit(train)
        S = model.sample(n)[0]
        errs.append(moment_error(S, mu, cov, w))
        traces.append(float(np.trace(np.cov(S.T))))
    return np.array(errs), np.array(traces)


def main():
    mu, cov, w = true_gmm(0)
    tm = (w[:, None]*mu).sum(0)
    tc = sum(w[k]*(cov[k]+np.outer(mu[k]-tm, mu[k]-tm)) for k in range(K))
    true_tr = float(np.trace(tc))
    rows = []
    for alpha in (0.0, 0.1, 0.25, 0.5, 1.0):
        E, T = [], []
        for s in range(5):
            e, tr = run(alpha, seed=s)
            E.append(e); T.append(tr)
        E = np.array(E); T = np.array(T)
        rows.append({"alpha": alpha, "seeds": 5, "generations": E.shape[1]-1,
                     "err_gen0": round(float(E[:, 0].mean()), 4),
                     "err_final": round(float(E[:, -1].mean()), 4),
                     "err_growth": round(float(E[:, -1].mean()/max(E[:, 0].mean(), 1e-9)), 3),
                     "variance_trace_true": round(true_tr, 3),
                     "variance_trace_final": round(float(T[:, -1].mean()), 3),
                     "variance_retained": round(float(T[:, -1].mean()/true_tr), 4)})
        print("  alpha=%.2f  moment error %.4f -> %.4f (%.2fx)   variance trace %.2f -> %.2f (%.1f%% of true)"
              % (alpha, E[:, 0].mean(), E[:, -1].mean(), rows[-1]["err_growth"],
                 true_tr, T[:, -1].mean(), 100*rows[-1]["variance_retained"]), flush=True)
    RES["claim6_alpha_tradeoff"] = {"D": D, "K": K, "rows": rows,
        "monotone_error_in_alpha": all(rows[i+1]["err_final"] <= rows[i]["err_final"]+1e-9
                                       for i in range(len(rows)-1)),
        "monotone_variance_in_alpha": all(rows[i+1]["variance_retained"] >= rows[i]["variance_retained"]-1e-9
                                          for i in range(len(rows)-1)),
        "collapse_at_alpha0": rows[0]["variance_retained"],
        "stable_at_alpha1": rows[-1]["variance_retained"]}
    R = RES["claim6_alpha_tradeoff"]
    print("  error monotone decreasing in alpha: %s | variance retention monotone increasing: %s"
          % (R["monotone_error_in_alpha"], R["monotone_variance_in_alpha"]), flush=True)
    json.dump(RES, open("collapse_results.json", "w"), indent=1)


if __name__ == "__main__":
    main(); print("DONE")