import json, math, random
from pathlib import Path
import numpy as np
from scipy.optimize import linprog

SEED = 2189
np.random.seed(SEED)
random.seed(SEED)


def certificate(q, S, epsilon, wmax=100.0):
    """Maximize delta subject to w_i >= delta and ||q+S w||inf <= epsilon."""
    q = np.asarray(q, dtype=float)
    S = np.asarray(S, dtype=float)
    d, L = len(q), S.shape[1]
    A, b = [], []
    for i in range(L):
        row = np.zeros(L + 1); row[i] = -1; row[-1] = 1
        A.append(row); b.append(0.)
    for k in range(d):
        row = np.zeros(L + 1); row[:L] = S[k]
        A.append(row); b.append(epsilon - q[k])
        row = np.zeros(L + 1); row[:L] = -S[k]
        A.append(row); b.append(epsilon + q[k])
    result = linprog(np.r_[np.zeros(L), -1.], A_ub=np.array(A), b_ub=np.array(b),
                     bounds=[(0., wmax)] * L + [(0., wmax)], method="highs")
    if not result.success:
        return np.zeros(L), 0., float("inf"), False
    w = result.x[:L]
    return w, float(result.x[-1]), float(np.max(np.abs(q + S @ w))), True


def mechanism_checks():
    # Prediction 1: for S=lambda I, q=(-3,-1), exact weights are (3/lambda,1/lambda),
    # hence delta=1/lambda until the wmax ceiling is active.
    lambdas = np.array([.25, .5, 1., 2., 4.])
    rows1 = []
    for lam in lambdas:
        w, delta, residual, ok = certificate([-3., -1.], lam*np.eye(2), 1e-9, 100.)
        rows1.append(dict(lambda_=float(lam), predicted_delta=float(1/lam),
                          observed_delta=delta, weights=w.tolist(), residual=residual))

    # Prediction 2: feasibility transition for scalar q=-3, S=1: delta can be positive
    # iff epsilon >= 3-wmax; for wmax=2 this predicts epsilon*=1.
    epsilons = np.array([.25, .75, 1., 1.25, 2.])
    rows2 = []
    for eps in epsilons:
        w, delta, residual, ok = certificate([-3.], [[1.]], eps, 2.)
        rows2.append(dict(epsilon=float(eps), predicted_positive=bool(eps >= 1.),
                          observed_positive=bool(delta > 1e-8), delta=delta, weight=w.tolist()))

    # Prediction 3: increasing slack epsilon permits a larger common margin for
    # q=(-2,-1), S=I; delta=1+eps (the upper residual bound on the second coordinate) in this unconstrained-wmax sweep.
    eps3 = np.array([0., .25, .5, .75, 1.])
    rows3 = []
    for eps in eps3:
        w, delta, residual, ok = certificate([-2., -1.], np.eye(2), eps, 100.)
        predicted = 1. + eps
        rows3.append(dict(epsilon=float(eps), predicted_delta=predicted,
                          observed_delta=delta, residual=residual))
    return {"scaling_lambda": rows1, "feasibility_transition": rows2,
            "slack_margin": rows3}


def hinge(v):
    return np.maximum(v, 0.)


def objective(z, target, thresholds, weights):
    # Convex 2-D last-layer problem: performance J=1/2||z-target||2,
    # ordered constraints V_i=[z_i-threshold_i]_+ (each tier acts on one coordinate).
    J = .5 * np.sum((z-target)**2)
    V = hinge(z-np.asarray(thresholds))
    return float(J + np.dot(weights, V)), float(J), V


def train_controller(mode, seed=SEED, steps=160, refresh=8):
    rng=np.random.default_rng(seed)
    target=np.array([1.0, 1.0, 1.0])
    thresholds=np.array([.2, .4, .6])
    z=np.zeros(3)
    weights=np.ones(3)
    history=[]
    for t in range(steps):
        if mode == "fixed": weights=np.array([100.,10.,1.])
        elif mode == "equal": weights=np.ones(3)
        elif mode == "certificate" and t % refresh == 0:
            # q = grad J and columns are subgradients of hinge tiers at current z.
            q=z-target
            S=np.eye(3) * (z-thresholds > 0.)
            neww, delta, residual, ok=certificate(q, S, epsilon=.05, wmax=100.)
            if ok and delta > 1e-5: weights=neww
            history.append((t, delta, residual, weights.copy()))
        # subgradient descent on weighted objective
        grad=z-target + weights*(z-thresholds > 0.)
        z -= .08*grad
        J=.5*np.sum((z-target)**2); V=hinge(z-thresholds)
        history.append((t, J, float(np.sum(V)), float(V[0]), float(np.dot(weights,V))))
    return z, history


def summarize():
    checks=mechanism_checks()
    comparison={}
    for mode in ["equal", "fixed", "certificate"]:
        z,h=train_controller(mode)
        comparison[mode]={"final_z":z.tolist(), "final_J":float(.5*np.sum((z-np.array([1.,1.,1.]))**2)),
                          "final_total_violation":float(np.sum(hinge(z-np.array([.2,.4,.6])))),
                          "final_high_priority_violation":float(hinge(z[0]-.2))}
    # Relative numerical errors for the predicted equations.
    e1=max(abs(r["observed_delta"]-r["predicted_delta"])/max(r["predicted_delta"],1e-12)
           for r in checks["scaling_lambda"])
    e3=max(abs(r["observed_delta"]-r["predicted_delta"]) for r in checks["slack_margin"])
    transition=all(r["observed_positive"] == r["predicted_positive"] for r in checks["feasibility_transition"])
    out={"seed":SEED, "mechanism_checks":checks,
         "check_summary":{"lambda_scaling_max_relative_error":e1,
                           "transition_all_correct":transition,
                           "slack_margin_max_absolute_error":e3},
         "training_comparison":comparison}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    summarize()