"""Claim 5 (Figure 1): subspace reconstruction error is LOWER when training on a
balanced, screened subset of sources than on the full heterogeneous population.

The previous run of this logbook found the opposite. The comparison there gave
the full population MORE TOTAL DATA (99 sources x n) than the screened subset
(36 sources x n), so it conflated "screening helps" with "more data helps".
Screening is a sample-ALLOCATION statement: at a fixed total budget, is it better
to spend it on a balanced admissible subset or spread over everything?

Both readings are measured below.
"""
import json, numpy as np

RES = {}


def population(d=30, k=6, N=99, seed=0):
    """Heterogeneous sources: a minority are well-conditioned in the shared
    subspace, the majority are degenerate (covariates collapse onto 1-2
    directions), which is the regime screening is meant to fix."""
    rng = np.random.default_rng(seed)
    B, _ = np.linalg.qr(rng.normal(size=(d, k)))     # true shared subspace
    srcs = []
    for i in range(N):
        if i % 3 == 0:                                # balanced source
            A = rng.normal(size=(d, d))*0.5+np.eye(d)
            Sig = A @ A.T/d+np.eye(d)*0.5
        else:                                         # degenerate source
            u = rng.normal(size=(d, 2))
            Sig = u @ u.T+np.eye(d)*0.02
        alpha = rng.normal(size=k)
        srcs.append({"Sigma": Sig, "theta": B @ alpha, "balanced": i % 3 == 0})
    return B, srcs


def draw(src, n, rng, noise=0.5):
    L = np.linalg.cholesky(src["Sigma"]+1e-8*np.eye(src["Sigma"].shape[0]))
    X = rng.normal(size=(n, L.shape[0])) @ L.T
    y = X @ src["theta"]+rng.normal(0, noise, size=n)
    return X, y


def estimate_subspace(data, k, ridge=1e-3):
    """Per-source ridge estimates, then top-k principal subspace of their
    second-moment matrix -- the standard shared-subspace estimator."""
    M = np.zeros((data[0][0].shape[1],)*2)
    for X, y in data:
        d = X.shape[1]
        th = np.linalg.solve(X.T @ X/len(X)+ridge*np.eye(d), X.T @ y/len(X))
        M += np.outer(th, th)
    w, V = np.linalg.eigh(M)
    return V[:, -k:]


def subspace_error(Bhat, B):
    """sin of the largest principal angle."""
    s = np.linalg.svd(Bhat.T @ B, compute_uv=False)
    return float(np.sqrt(max(0.0, 1-min(s)**2)))


def screen(srcs, m):
    """Select m sources greedily maximising lambda_min of the aggregate design
    -- the admissibility/spectral-norm criterion the theory uses."""
    d = srcs[0]["Sigma"].shape[0]
    chosen, S = [], np.zeros((d, d))
    cand = list(range(len(srcs)))
    for _ in range(m):
        best, bi = -1e18, None
        for i in cand:
            lam = np.linalg.eigvalsh(S+srcs[i]["Sigma"])[0]
            if lam > best: best, bi = lam, i
        chosen.append(bi); cand.remove(bi); S = S+srcs[bi]["Sigma"]
    return chosen


def run():
    d, k, N = 30, 6, 99
    B, srcs = population(d, k, N)
    m = 33
    sel = screen(srcs, m)
    frac_bal = float(np.mean([srcs[i]["balanced"] for i in sel]))
    print("  screened %d/%d sources; fraction balanced = %.2f (population %.2f)"
          % (m, N, frac_bal, np.mean([s["balanced"] for s in srcs])), flush=True)

    rows = []
    for total in (1980, 3960, 7920, 15840):     # TOTAL sample budget
        e_full, e_scr, e_naive = [], [], []
        for rep in range(8):
            rng = np.random.default_rng(500+rep)
            n_full = total//N                       # spread over all sources
            n_scr = total//m                        # concentrated on screened
            df = [draw(srcs[i], n_full, rng) for i in range(N)]
            ds = [draw(srcs[i], n_scr, rng) for i in sel]
            dn = [draw(srcs[i], n_full, rng) for i in sel]   # unequal-budget control
            e_full.append(subspace_error(estimate_subspace(df, k), B))
            e_scr.append(subspace_error(estimate_subspace(ds, k), B))
            e_naive.append(subspace_error(estimate_subspace(dn, k), B))
        rows.append({"total_budget": total, "n_per_source_full": total//N,
                     "n_per_source_screened": total//m,
                     "err_full_population": round(float(np.mean(e_full)), 5),
                     "err_screened_matched_budget": round(float(np.mean(e_scr)), 5),
                     "err_screened_unmatched": round(float(np.mean(e_naive)), 5),
                     "screening_helps_matched": bool(np.mean(e_scr) < np.mean(e_full)),
                     "sd_full": round(float(np.std(e_full)), 5),
                     "sd_screened": round(float(np.std(e_scr)), 5)})
        print("  budget %-6d  full=%.5f  screened(matched)=%.5f  screened(same n/src)=%.5f  %s"
              % (total, np.mean(e_full), np.mean(e_scr), np.mean(e_naive),
                 "SCREENING WINS" if np.mean(e_scr) < np.mean(e_full) else "full wins"), flush=True)
    RES["claim5_screening"] = {
        "d": d, "k": k, "N_sources": N, "m_screened": m, "reps": 8,
        "fraction_balanced_in_screened": frac_bal,
        "rows": rows,
        "screening_wins_at_matched_budget": sum(r["screening_helps_matched"] for r in rows),
        "cells": len(rows)}
    json.dump(RES, open("screen_results.json", "w"), indent=1)


if __name__ == "__main__":
    run()
    print("DONE")