import json
import time
import numpy as np


def soft_threshold(x, t):
    return np.sign(x) * np.maximum(np.abs(x) - t, 0.0)


class Quadratic:
    def __init__(self, eigs, seed=0):
        self.eigs = np.asarray(eigs, dtype=float)
        self.n = len(eigs)
        rng = np.random.default_rng(seed)
        q, _ = np.linalg.qr(rng.normal(size=(self.n, self.n)))
        self.A = q @ np.diag(self.eigs) @ q.T
        self.b = rng.normal(size=self.n)
        self.calls = 0

    def grad(self, x):
        self.calls += 1
        return self.A @ x - self.b

    def value(self, x):
        return 0.5 * x @ self.A @ x - self.b @ x

    def prox_exact(self, z, gamma):
        return np.linalg.solve(np.eye(self.n) + gamma * self.A, z + gamma * self.b)


def residual_from_grad(x, z, gamma, grad):
    return gamma * grad + x - z


def bfgs_inverse(H, s, q):
    ys = float(q @ s)
    if ys <= 1e-12 * max(1.0, np.linalg.norm(q) * np.linalg.norm(s)):
        return H
    rho = 1.0 / ys
    I = np.eye(len(s))
    return (I - rho * np.outer(s, q)) @ H @ (I - rho * np.outer(q, s)) + rho * np.outer(s, s)


def verify_math():
    rng = np.random.default_rng(12)
    f = Quadratic([0.4, 1.0, 2.5, 5.0], seed=4)
    gamma = 0.37
    x = rng.normal(size=f.n)
    z0, z1 = rng.normal(size=(2, f.n))
    grad = f.grad(x)
    r0 = residual_from_grad(x, z0, gamma, grad)
    direct = residual_from_grad(x, z1, gamma, grad)
    transported = r0 + z0 - z1
    transport_err = np.linalg.norm(direct - transported)

    secant_errors = []
    secant_ratios = []
    for _ in range(500):
        u, v = rng.normal(size=(2, f.n))
        s = v - u
        q = gamma * (f.A @ s) + s
        expected = np.dot(s, s) + gamma * s @ f.A @ s
        secant_errors.append(abs(s @ q - expected))
        secant_ratios.append((s @ q) / (s @ s))

    # For exact H=(I+gamma A)^-1 and undamped unit Newton steps, residual
    # contracts to zero in one step. With H=I, the factor is 1-gamma*lambda.
    boundary = []
    for lam in [0.5, 1.0, 1.8, 2.2]:
        for g in [0.2, 0.5, 0.9, 1.1]:
            factor = abs(1.0 - g * lam)
            boundary.append((g * lam, factor, factor < 1.0))
    pred_ok = all((prod < 2.0 and stable) or (prod >= 2.0 and not stable)
                  for prod, _, stable in boundary)
    # Sweep confirms the exact scalar stability interval gamma*lambda in (0,2).
    observed_boundary = max(prod for prod, factor, _ in boundary if factor < 1.0)
    return {
        "transport_l2_error": float(transport_err),
        "max_secant_identity_error": float(max(secant_errors)),
        "secant_ratio_min": float(min(secant_ratios)),
        "secant_ratio_predicted_min": 1.0,
        "stability_sweep": [{"gamma_lambda": p, "observed_factor": a, "contractive": c} for p, a, c in boundary],
        "stability_prediction_pass": bool(pred_ok),
        "stability_interval_prediction": "0 < gamma*lambda < 2 for H=I",
        "stability_sweep_max_contracting_product": float(observed_boundary),
    }


def restarted_pg(f, centers, gamma, inner_steps, l1, step):
    x = centers[0].copy()
    calls = 0
    residuals = []
    for z in centers:
        x = z.copy()  # independently restarted solve
        for _ in range(inner_steps):
            g = f.grad(x); calls += 1
            r = residual_from_grad(x, z, gamma, g)
            x -= step * r
        residuals.append(float(np.linalg.norm(residual_from_grad(x, z, gamma, f.A @ x - f.b))))
        x = soft_threshold(x, gamma * l1)
    return x, calls, residuals


def recycled_bfgs(f, centers, gamma, inner_steps, l1, alpha=0.8):
    # Frozen CR-DRS-style predictor: pseudo-steps use B=H^{-1} only;
    # exactly one expensive gradient is evaluated at the final predictor.
    x = centers[0].copy()
    H = np.eye(f.n)
    old_z = centers[0].copy()
    g = f.grad(x); calls = 1
    r = residual_from_grad(x, old_z, gamma, g)
    residuals = []
    for z in centers:
        r = r + old_z - z  # exact center transport, no gradient call
        x_anchor, r_anchor = x.copy(), r.copy()
        B = np.linalg.inv(H)
        for _ in range(inner_steps):
            d = -H @ r
            Bd = B @ d
            denom = float(d @ Bd)
            numer = float(-(r @ d))
            eta = min(alpha, numer / denom) if denom > 1e-14 and numer > 0 else 0.0
            if eta == 0.0:
                break
            dx = eta * d
            x = x + dx
            r = r + B @ dx  # modeled residual; no expensive evaluation
        # One true gradient at the endpoint, as in the paper.
        gt = f.grad(x); calls += 1
        rt = residual_from_grad(x, z, gamma, gt)
        s, q = x - x_anchor, rt - r_anchor
        H = bfgs_inverse(H, s, q)
        r = rt
        residuals.append(float(np.linalg.norm(r)))
        x = soft_threshold(x, gamma * l1)
        old_z = z.copy()
    return x, calls, residuals


def run_experiment():
    rng = np.random.default_rng(7)
    f1 = Quadratic(np.linspace(0.5, 8.0, 24), seed=9)
    f2 = Quadratic(np.linspace(0.5, 8.0, 24), seed=9)
    centers = []
    z = rng.normal(scale=1.5, size=24)
    for k in range(30):
        centers.append(z.copy())
        z = 0.88 * z + 0.12 * rng.normal(size=24)
    gamma, inner, l1 = 0.08, 3, 0.015
    # conservative fixed residual-gradient step; same three expensive calls/outer update
    step = 1.0 / (1.0 + gamma * 8.0)
    t0 = time.perf_counter(); _, cb, rb = restarted_pg(f1, centers, gamma, inner, l1, step); tb = time.perf_counter()-t0
    t0 = time.perf_counter(); _, ci, ri = recycled_bfgs(f2, centers, gamma, inner, l1); ti = time.perf_counter()-t0
    return {
        "outer_updates": len(centers), "inner_steps": inner,
        "baseline_gradient_calls": cb, "idea_gradient_calls": ci,
        "gradient_call_reduction_percent": 100.0*(cb-ci)/cb,
        "baseline_final_residual": rb[-1], "idea_final_residual": ri[-1],
        "baseline_mean_residual": float(np.mean(rb)), "idea_mean_residual": float(np.mean(ri)),
        "baseline_seconds": tb, "idea_seconds": ti
    }


if __name__ == "__main__":
    out = {"math_verification": verify_math(), "mini_experiment": run_experiment()}
    with open("results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))