"""Claim 6: evaluate Algorithms 3 (DoWS) and 4 (T-DoWS) against a primal-dual
baseline on SVM classification with real datasets.

The logbook stated it does not reproduce the real-data SVM comparison and ran no
SVM experiment. Two of the three datasets are run here; MNIST 3-vs-5 could not be
retrieved (OpenML returned HTTP 504).

SVM as constrained optimisation, which is the form the algorithms address:
    minimise  f(w,b) = 0.5||w||^2
    subject to  y_i (w.x_i + b) >= 1   (one constraint per sample)

Randomised feasibility: sample a violated constraint and project onto its
halfspace, rather than forming the full constraint set.
"""
import json, pickle, numpy as np

RES = {}


def load():
    D = pickle.load(open("svm_data.pkl", "rb"))
    out = {}
    for k, (X, y) in D.items():
        X = (X-X.mean(0))/(X.std(0)+1e-9)
        out[k] = (X, y.astype(float))
    return out


C_SLACK = 1.0


def obj(w, b, xi=None):
    """SOFT-margin objective. Hard margin is INFEASIBLE on banknote (min margin
    -0.767 under an exact LinearSVC fit), so violations there could never reach 0
    and the feasibility comparison was measuring an infeasible program."""
    v = 0.5*float(w @ w)
    return v+C_SLACK*float(np.sum(xi)) if xi is not None else v


def viol(X, y, w, b, xi):
    return float(np.maximum(0.0, 1.0-xi-y*(X @ w+b)).max())


def project_one(X, y, w, b, xi, i):
    """Project (w,b,xi_i) onto y_i(w.x_i+b) + xi_i >= 1, xi_i >= 0."""
    a = y[i]*X[i]; c = y[i]
    g = a @ w+c*b+xi[i]-1.0
    if g < 0:
        n2 = a @ a+c*c+1.0
        lam = -g/max(n2, 1e-12)
        w = w+lam*a; b = b+lam*c; xi[i] = xi[i]+lam
    xi[i] = max(0.0, xi[i])
    return w, b


def run(X, y, alg, T=4000, seed=0, nproj=None):
    """nproj scales with the CONSTRAINT COUNT. Claim 4 of this same paper (already
    verified) says infeasibility decays geometrically in the NUMBER of feasibility
    updates; with a fixed 5 projections/step over 569-1372 constraints most
    constraints are almost never visited and violation plateaus."""
    rng = np.random.default_rng(seed)
    n, d = X.shape
    if nproj is None: nproj = max(5, n//4)
    w = np.zeros(d); b = 0.0; w0 = w.copy(); xi = np.ones(n)
    G2 = 1e-12; rmax = 1e-8
    lam_d = 0.0
    hist = []
    for t in range(1, T+1):
        g = w.copy()                                   # grad of 0.5||w||^2
        if alg == "primal_dual":
            i = int(rng.integers(n))
            s = 1.0-y[i]*(X[i] @ w+b)
            gv = -y[i]*X[i] if s > 0 else np.zeros(d)
            gb = -y[i] if s > 0 else 0.0
            eta = 0.5/np.sqrt(t)
            w = w-eta*(g+lam_d*gv); b = b-eta*(lam_d*gb)
            xi[i] = max(0.0, xi[i]-eta*(C_SLACK-lam_d))
            lam_d = max(0.0, lam_d+eta*(s-xi[i]))
        else:
            G2 += g @ g
            rmax = max(rmax, float(np.linalg.norm(w-w0)))
            if alg == "dows":
                eta = rmax/np.sqrt(G2)
            else:                                       # t_dows: tamed, no bounded-Y assumption
                eta = rmax/np.sqrt(G2*np.log(np.e+t))
            w = w-eta*g
            xi = np.maximum(0.0, xi-eta*C_SLACK)
            for _ in range(nproj):                      # randomised feasibility updates
                i = int(rng.integers(n))
                w, b = project_one(X, y, w, b, xi, i)
        if t % 200 == 0:
            hist.append({"t": t, "obj": obj(w, b, xi), "max_violation": viol(X, y, w, b, xi)})
    acc = float(np.mean(np.sign(X @ w+b) == y))
    return w, b, hist, acc


def main():
    data = load()
    rows = []
    for name, (X, y) in data.items():
        for alg in ("dows", "t_dows", "primal_dual"):
            o, v, a = [], [], []
            for s in range(3):
                w, b, h, acc = run(X, y, alg, seed=s)
                o.append(h[-1]["obj"]); v.append(h[-1]["max_violation"]); a.append(acc)
            rows.append({"dataset": name, "n": int(X.shape[0]), "d": int(X.shape[1]),
                         "algorithm": alg, "final_objective": round(float(np.mean(o)), 5),
                         "final_max_violation": round(float(np.mean(v)), 6),
                         "train_accuracy": round(float(np.mean(a)), 4)})
            print("  %-14s %-12s  objective=%9.4f   max violation=%9.5f   accuracy=%.4f"
                  % (name, alg, np.mean(o), np.mean(v), np.mean(a)), flush=True)
    RES["claim6_svm"] = {"rows": rows, "seeds": 3, "iterations": 4000,
        "datasets_run": sorted(data.keys()),
        "mnist_note": "MNIST 3-vs-5 unavailable: OpenML returned HTTP 504"}
    # comparison summary
    summ = []
    for name in data:
        sub = {r["algorithm"]: r for r in rows if r["dataset"] == name}
        summ.append({"dataset": name,
                     "dows_beats_pd_on_violation": bool(sub["dows"]["final_max_violation"] <= sub["primal_dual"]["final_max_violation"]),
                     "t_dows_beats_pd_on_violation": bool(sub["t_dows"]["final_max_violation"] <= sub["primal_dual"]["final_max_violation"]),
                     "dows_acc_minus_pd": round(sub["dows"]["train_accuracy"]-sub["primal_dual"]["train_accuracy"], 4)})
        print("  [%s] DoWS beats PD on feasibility: %s | T-DoWS: %s | accuracy delta %+.4f"
              % (name, summ[-1]["dows_beats_pd_on_violation"], summ[-1]["t_dows_beats_pd_on_violation"],
                 summ[-1]["dows_acc_minus_pd"]), flush=True)
    RES["claim6_svm"]["summary"] = summ
    json.dump(RES, open("svm_results.json", "w"), indent=1)


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