import json
from pathlib import Path
import numpy as np


def cg(A, b, tol=1e-11, max_iter=1000):
    x = np.zeros_like(b, dtype=float)
    r = b - A(x)
    p = r.copy()
    rr = float(r @ r)
    r0 = np.sqrt(rr)
    if r0 == 0:
        return x, 0, 0.0
    for k in range(1, max_iter + 1):
        Ap = A(p)
        den = float(p @ Ap)
        if den <= 0 or not np.isfinite(den):
            raise RuntimeError("CG operator was not positive definite")
        alpha = rr / den
        x += alpha * p
        r -= alpha * Ap
        new_rr = float(r @ r)
        if np.sqrt(new_rr) <= tol * max(1.0, r0):
            return x, k, np.sqrt(new_rr)
        p = r + (new_rr / rr) * p
        rr = new_rr
    return x, max_iter, np.linalg.norm(r)


def rel(a, b):
    return float(np.linalg.norm(a - b) / max(1e-15, np.linalg.norm(b)))


def main():
    rng = np.random.default_rng(2796)
    q, _ = np.linalg.qr(rng.normal(size=(4, 4)))
    eig = np.array([0.0, 0.1, 1.0, 10.0])
    H = q @ np.diag(eig) @ q.T
    # B is in Range(H), so the pseudoinverse derivative is finite.
    B = q @ np.array([0.0, 1.0, -0.7, 0.4])
    v_true = np.linalg.pinv(H) @ B
    eps_values = np.logspace(-1, -8, 8)

    stable = []
    for eps in eps_values:
        A = H + eps * np.eye(4)
        v, it, residual = cg(lambda z: A @ z, B)
        stable.append({
            "eps": float(eps), "relative_error": rel(v, v_true),
            "iterations": int(it), "residual": float(residual),
        })

    # Prediction 1: stable-range Tikhonov derivative converges to -H^+B.
    # Prediction 2: for each positive eigencomponent, bias is exactly
    # eps/(lambda+eps), hence asymptotically linear in eps.
    component_errors = []
    for eps in eps_values:
        v = np.linalg.solve(H + eps * np.eye(4), B)
        coeff = q.T @ v
        true_coeff = q.T @ v_true
        component_errors.append({
            "eps": float(eps),
            "lambda_0_error": float(abs(coeff[0] - true_coeff[0])),
            "lambda_0p1_relative_error": float(abs(coeff[1] - true_coeff[1]) / abs(true_coeff[1])),
            "predicted_lambda_0p1_relative_error": float(eps / (0.1 + eps)),
        })

    # Prediction 3: if B has a nullspace component, the regularized solve
    # diverges as (B_null)/eps. We sweep eps and estimate the log-log slope.
    B_bad = B + 0.8 * q[:, 0]
    unstable = []
    bad_norms = []
    for eps in eps_values:
        v = np.linalg.solve(H + eps * np.eye(4), B_bad)
        null_coeff = float(q[:, 0] @ v)
        bad_norms.append(np.linalg.norm(v))
        unstable.append({"eps": float(eps), "norm": float(np.linalg.norm(v)), "null_component": null_coeff,
                         "eps_times_null_component": float(eps * null_coeff)})
    slope = float(np.polyfit(np.log(eps_values[-5:]), np.log(bad_norms[-5:]), 1)[0])

    # Tiny bilevel quadratic: f(x,theta)=1/2 x^T Hx - theta B^T x.
    # The minimum-norm inner solution is x*=H^+B theta. Outer loss is
    # 1/2||x-c||^2. Exact hypergradient is (x-c)^T H^+B.
    c = q @ np.array([0.3, -0.2, 0.4, 0.1])
    direction = v_true
    theta0 = 2.0
    exact_x = direction * theta0
    exact_loss = 0.5 * np.sum((exact_x - c) ** 2)
    exact_grad = float((exact_x - c) @ direction)

    def damped_grad(theta, eps):
        x = np.linalg.solve(H + eps * np.eye(4), B) * theta
        dx = np.linalg.solve(H + eps * np.eye(4), B)
        return float((x - c) @ dx), float(0.5 * np.sum((x - c) ** 2))

    # Compare fixed damping and continuation (eps halves toward a floor).
    fixed_eps = 0.1
    theta_fixed = theta0
    theta_cont = theta0
    fixed_trace, cont_trace = [], []
    lr = 0.015
    for step in range(40):
        gf, lf = damped_grad(theta_fixed, fixed_eps)
        gc, lc = damped_grad(theta_cont, max(1e-7, 0.1 * 0.7 ** step))
        theta_fixed -= lr * gf
        theta_cont -= lr * gc
        fixed_trace.append(lf)
        cont_trace.append(lc)

    out = {
        "seed": 2796,
        "eigenvalues": eig.tolist(),
        "predictions": {
            "stable_range_limit": "relative error -> 0 for B in Range(H)",
            "positive_eigen_bias": "relative bias on lambda=0.1 component = eps/(0.1+eps)",
            "nullspace_violation": "norm scales as eps^-1 when B has null component",
        },
        "stable_sweep": stable,
        "component_sweep": component_errors,
        "unstable_sweep": unstable,
        "unstable_loglog_slope": slope,
        "bilevel": {
            "exact_minimum_norm_loss_at_theta0": exact_loss,
            "exact_hypergradient_at_theta0": exact_grad,
            "fixed_eps": fixed_eps,
            "fixed_final_loss": fixed_trace[-1],
            "continuation_final_loss": cont_trace[-1],
            "fixed_final_theta": theta_fixed,
            "continuation_final_theta": theta_cont,
            "fixed_loss_trace": fixed_trace,
            "continuation_loss_trace": cont_trace,
        },
    }
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps({
        "stable_last_relative_error": stable[-1]["relative_error"],
        "lambda01_observed_vs_predicted": [component_errors[-1]["lambda_0p1_relative_error"], component_errors[-1]["predicted_lambda_0p1_relative_error"]],
        "unstable_loglog_slope": slope,
        "exact_grad": exact_grad,
        "fixed_final_loss": fixed_trace[-1],
        "continuation_final_loss": cont_trace[-1],
    }, indent=2))


if __name__ == "__main__":
    main()