"""Claim 1 (Theorem 2): training on an admissible subset attains the minimax rate
O(sqrt(d/(N lambda_k))). Tested as an exponent fit on n and, separately, on d --
a rate claim has two exponents and fitting only one leaves the other unchecked.
"""
import json, numpy as np
from screen_exp import population, draw, estimate_subspace, subspace_error, screen

RES = json.load(open("screen_results.json"))


def rate_vs_n():
    d, k, N, m = 30, 6, 99, 33
    B, srcs = population(d, k, N)
    sel = screen(srcs, m)
    ns, errs = [], []
    for n in (20, 40, 80, 160, 320, 640):
        e = []
        for rep in range(8):
            rng = np.random.default_rng(700+rep)
            ds = [draw(srcs[i], n, rng) for i in sel]
            e.append(subspace_error(estimate_subspace(ds, k), B))
        ns.append(n); errs.append(float(np.mean(e)))
        print("   n/src=%-4d  subspace err=%.5f" % (n, errs[-1]), flush=True)
    sl, ic = np.polyfit(np.log(ns), np.log(errs), 1)
    r2 = 1-np.var(np.log(errs)-(sl*np.log(ns)+ic))/np.var(np.log(errs))
    print("  exponent on n: %.4f (predicted -0.5), R2 %.4f" % (sl, r2), flush=True)
    return {"n_per_source": ns, "errors": [round(x, 6) for x in errs],
            "fitted_exponent": round(float(sl), 4), "predicted": -0.5,
            "r2": round(float(r2), 4)}


def rate_vs_d():
    k, N, m, n = 6, 99, 33, 200
    ds_, errs = [], []
    for d in (12, 20, 30, 45, 60):
        B, srcs = population(d, k, N, seed=d)
        sel = screen(srcs, m)
        e = []
        for rep in range(6):
            rng = np.random.default_rng(800+rep)
            data = [draw(srcs[i], n, rng) for i in sel]
            e.append(subspace_error(estimate_subspace(data, k), B))
        ds_.append(d); errs.append(float(np.mean(e)))
        print("   d=%-4d  subspace err=%.5f" % (d, errs[-1]), flush=True)
    sl, ic = np.polyfit(np.log(ds_), np.log(errs), 1)
    r2 = 1-np.var(np.log(errs)-(sl*np.log(ds_)+ic))/np.var(np.log(errs))
    print("  exponent on d: %.4f (predicted +0.5), R2 %.4f" % (sl, r2), flush=True)
    return {"d": ds_, "errors": [round(x, 6) for x in errs],
            "fitted_exponent": round(float(sl), 4), "predicted": 0.5,
            "r2": round(float(r2), 4)}


if __name__ == "__main__":
    RES["claim1_minimax_rate"] = {"vs_n": rate_vs_n(), "vs_d": rate_vs_d(),
                                  "note": "sqrt(d/(N lambda_k)) predicts err ~ n^-1/2 and ~ d^+1/2"}
    json.dump(RES, open("screen_results.json", "w"), indent=1)
    print("DONE")