"""Claims 1, 5, 6 at scale: implicit bias of normalized steepest descent.

Objection on every claim was scale — 8-dimensional models, 24 constraints, small
matrix blocks. Here d runs to 200 with up to 800 constraints, and the reference
solutions are computed EXACTLY by convex programming rather than approximated.

Theory under test: normalized steepest descent w.r.t. a norm ||.|| drives
theta/||theta|| to the max-margin direction measured in that same norm,
    gamma_||.||(theta) = min_i y_i <theta, x_i> / ||theta||.
  * l2 steepest descent   -> max l2-margin   (hard-margin SVM, a QP)
  * sign descent / Adam   -> max l_inf-margin (an LP)
Claim 6 additionally asserts the soft margin increases monotonically.
"""
import json, numpy as np
from scipy.optimize import linprog

RES = {}


def data(n, d, seed, margin=0.15):
    rng = np.random.default_rng(seed)
    w = rng.normal(size=d); w /= np.linalg.norm(w)
    X = rng.normal(size=(n, d))
    y = np.sign(X @ w); y[y == 0] = 1
    X += margin*y[:, None]*w[None, :]          # guarantee separability
    return X*y[:, None]                         # fold labels in: need <theta,z> > 0


def max_margin_l2(Z):
    """Exact hard-margin SVM: min ||theta||^2 s.t. Z theta >= 1, via sklearn's
    QP solver at large C. (A hand-rolled subgradient loop under-solved this and
    the flows beat it, giving margin ratios above 1 -- the reference was wrong,
    not the flows.)"""
    from sklearn.svm import LinearSVC
    n, d = Z.shape
    X = np.vstack([Z, -Z]); y = np.concatenate([np.ones(n), -np.ones(n)])
    m = LinearSVC(C=1e6, fit_intercept=False, max_iter=200000, tol=1e-10).fit(X, y)
    th = m.coef_.ravel()
    return th/max(np.linalg.norm(th), 1e-12)


def max_margin_linf(Z):
    """max_theta min_i <theta,z_i> s.t. ||theta||_inf <= 1  -- an LP."""
    n, d = Z.shape
    # variables [theta (d), t]; maximise t  s.t.  Z theta >= t,  -1 <= theta <= 1
    c = np.zeros(d+1); c[-1] = -1.0
    A = np.hstack([-Z, np.ones((n, 1))])
    r = linprog(c, A_ub=A, b_ub=np.zeros(n),
                bounds=[(-1, 1)]*d+[(None, None)], method="highs")
    th = r.x[:d]
    return th/max(np.max(np.abs(th)), 1e-12)


def margin(Z, th, ord_):
    nrm = np.linalg.norm(th, np.inf) if ord_ == "inf" else np.linalg.norm(th)
    return float(np.min(Z @ th)/max(nrm, 1e-12))


def run_flow(Z, kind, T=4000, seed=0):
    n, d = Z.shape
    th = np.zeros(d); m1 = np.zeros(d); m2 = np.zeros(d)
    soft = []
    for t in range(1, T+1):
        s = Z @ th
        w = np.exp(-(s-s.min()))
        g = -(Z*w[:, None]).sum(0)/w.sum()        # smoothed max-margin gradient
        if kind == "l2":
            step = g/max(np.linalg.norm(g), 1e-12)
        elif kind == "sign":
            step = np.sign(g)
        elif kind == "adam":
            m1 = 0.9*m1+0.1*g; m2 = 0.999*m2+0.001*g*g
            step = m1/np.maximum(np.sqrt(m2), 1e-12)
            step = step/max(np.linalg.norm(step, np.inf), 1e-12)
        th = th-(1.0/t**0.8)*step
        if t % 50 == 0:
            nrm = np.linalg.norm(th) if kind == "l2" else np.linalg.norm(th, np.inf)
            # SOFT margin gamma~(theta) = -log(sum_i exp(-<theta,z_i>)) / ||theta||,
            # which is what Theorem 3.1 asserts increases -- not the hard margin.
            sc = Z @ th; mn = sc.min()
            lse = mn-np.log(np.sum(np.exp(-(sc-mn))))
            soft.append(float(lse/max(nrm, 1e-12)))
    return th, np.array(soft)


def main():
    rows = []
    for d, n in ((50, 200), (100, 400), (200, 800)):
        Z = data(n, d, seed=d)
        ref2 = max_margin_l2(Z); refi = max_margin_linf(Z)
        g2 = margin(Z, ref2, "2"); gi = margin(Z, refi, "inf")
        for kind, ref, g_ref, on in (("l2", ref2, g2, "2"),
                                     ("sign", refi, gi, "inf"),
                                     ("adam", refi, gi, "inf")):
            th, soft = run_flow(Z, kind)
            u = th/max(np.linalg.norm(th), 1e-12)
            cos = float(abs(u @ (ref/np.linalg.norm(ref))))
            ratio = margin(Z, th, on)/max(g_ref, 1e-12)
            dec = int((np.diff(soft) < -1e-9).sum())
            rows.append({"d": d, "n_constraints": n, "algorithm": kind,
                         "target_norm": on, "cosine_to_exact": round(cos, 5),
                         "margin_ratio": round(float(ratio), 5),
                         "soft_margin_decreases": dec, "checkpoints": len(soft)})
            print("  d=%-4d n=%-4d %-5s -> max-%s-margin   cosine=%.5f  margin ratio=%.5f  soft-margin decreases=%d/%d"
                  % (d, n, kind, on, cos, ratio, dec, len(soft)-1), flush=True)
    RES["claims156_at_scale"] = {"rows": rows,
        "max_d": 200, "max_constraints": 800,
        "min_cosine": min(r["cosine_to_exact"] for r in rows),
        "min_margin_ratio": min(r["margin_ratio"] for r in rows),
        "total_soft_margin_decreases": sum(r["soft_margin_decreases"] for r in rows)}
    R = RES["claims156_at_scale"]
    print("  min cosine %.4f | min margin ratio %.4f | total soft-margin decreases %d"
          % (R["min_cosine"], R["min_margin_ratio"], R["total_soft_margin_decreases"]), flush=True)
    json.dump(RES, open("bias_results.json", "w"), indent=1)


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