import json, time
import numpy as np
from scipy.optimize import minimize
from pdas_filter import solve_active_set, kkt_residual

np.set_printoptions(precision=5, suppress=True)

# Strictly convex QP: min 1/2 ||v-v_nom||^2, A v >= b.
def make_problem(v_nom, alpha=1.0, h=None, A=None):
    n = len(v_nom)
    if A is None:
        A = np.vstack([np.eye(n), -np.eye(n)])
    if h is None:
        h = -np.ones(A.shape[0])
    return np.eye(n), -np.asarray(v_nom), np.asarray(A, float), -alpha*np.asarray(h)


def projected_gradient(v_nom, A, b, steps=250, lr=.12):
    # Generic approximate baseline: gradient step followed by sequential projections.
    v = np.zeros_like(v_nom)
    for _ in range(steps):
        v -= lr * (v - v_nom)
        for a, bj in zip(A, b):
            gap = bj - a @ v
            if gap > 0:
                v += gap * a / (a @ a)
    return v

# Prediction 1: for one constraint v >= alpha*q, activity changes at alpha=v_nom/q.
v_nom = np.array([0.37]); q = 0.8
alphas = np.linspace(0., 1.2, 25)
active = []
for alpha in alphas:
    H,c,A,b = make_problem(v_nom, alpha, np.array([-q]), np.array([[1.]]))
    _,_,act,_,_ = solve_active_set(H,c,A,b)
    active.append(bool(act))
pred_boundary = v_nom[0]/q
changes = np.where(np.diff(np.array(active, dtype=int)) != 0)[0]
obs_boundary = alphas[changes[0]+1] if len(changes) else np.nan

# Prediction 2: after activation, mu = alpha*q-v_nom, slope dmu/dalpha=q.
alpha_fit = np.linspace(pred_boundary + .08, 1.6, 12)
mus = []
for alpha in alpha_fit:
    H,c,A,b = make_problem(v_nom, alpha, np.array([-q]), np.array([[1.]]))
    _,mu,_,_,_ = solve_active_set(H,c,A,b)
    mus.append(mu[0])
slope = np.polyfit(alpha_fit, mus, 1)[0]

# Prediction 3: smoothly moving nominal velocities preserve active sets and warm starts
# converge in one solve iteration after the first transition.
A = np.array([[1., 0.], [0., 1.], [-1., 0.]])
h = np.array([-0.42, -0.25, 1.0])
trajectory = np.column_stack([np.linspace(.1, .8, 40), np.linspace(.1, .35, 40)])
prev = None; iterations = []; changes_count = 0; prior_act = None; max_kkt = 0.
for vn in trajectory:
    H,c,A2,b = make_problem(vn, 1., h, A)
    v,mu,act,it,_ = solve_active_set(H,c,A2,b,active_init=prev)
    iterations.append(it)
    changes_count += int(prior_act is not None and act != prior_act)
    prior_act = act; prev = act
    max_kkt = max(max_kkt, kkt_residual(H,c,A2,b,v,mu))

# Same small collection of changing QPs: exact PDAS versus projected-gradient approximation.
rng = np.random.default_rng(2265)
A4 = np.eye(4)
h4 = -np.ones(4)*.55
pdas_t = []; pg_t = []; pdas_err = []; pg_err = []; violations = []
for vn in rng.normal(0, .5, size=(60,4)):
    H,c,A5,b = make_problem(vn, 1., h4, A4)
    t=time.perf_counter(); v,mu,act,it,_=solve_active_set(H,c,A5,b); pdas_t.append(time.perf_counter()-t)
    t=time.perf_counter(); vp=projected_gradient(vn,A5,b); pg_t.append(time.perf_counter()-t)
    pdas_err.append(np.linalg.norm(v-vn)**2/2); pg_err.append(np.linalg.norm(vp-vn)**2/2)
    violations.append(max(0., np.max(b-A5@vp)))

result = {
 "prediction_boundary": {"predicted_alpha": float(pred_boundary), "observed_grid_alpha": float(obs_boundary), "grid_error": float(abs(obs_boundary-pred_boundary))},
 "prediction_dual_slope": {"predicted_dmu_dalpha": float(q), "observed_slope": float(slope), "relative_error": float(abs(slope-q)/q)},
 "prediction_warm_start": {"mean_iterations": float(np.mean(iterations)), "median_iterations": float(np.median(iterations)), "fraction_one_iteration": float(np.mean(np.array(iterations)==1)), "active_set_changes": int(changes_count), "max_kkt_residual": float(max_kkt)},
 "baseline_comparison": {"pdas_mean_objective": float(np.mean(pdas_err)), "projected_gradient_mean_objective": float(np.mean(pg_err)), "pdas_mean_ms": float(1000*np.mean(pdas_t)), "projected_gradient_mean_ms": float(1000*np.mean(pg_t)), "projected_gradient_max_violation": float(max(violations))},
 "all_iterations": iterations
}
print(json.dumps(result, indent=2))
with open("results.json", "w") as f: json.dump(result, f, indent=2)