"""Lyapunov-sign-preserving stepping MVP and quantitative verification."""
import json, math, os
import numpy as np
import torch

SEED = 7
np.random.seed(SEED)
torch.manual_seed(SEED)


def euler_step(z, drift, h):
    return z + h * drift(z)


def semi_implicit_linear_step(z, A, h):
    return np.linalg.solve(np.eye(A.shape[0]) - h * A, z)


def exponent_from_jacobians(jacobians, h):
    q = np.ones(jacobians[0].shape[0] if np.ndim(jacobians[0]) > 0 else 1)
    q /= np.linalg.norm(q)
    total = 0.0
    for J in jacobians:
        v = (float(J) * q) if np.ndim(J) == 0 else np.asarray(J) @ q
        norm = np.linalg.norm(v)
        total += math.log(max(norm, 1e-15))
        q = v / max(norm, 1e-15)
    return total / (len(jacobians) * h)


def estimate_with_torch(drift, z, h, method="euler", A=None):
    """One exact JVP of a discrete map, suitable for a neural drift."""
    z = z.detach().clone().requires_grad_(True)
    q = torch.randn_like(z)
    q = q / q.norm()
    if method == "euler":
        step = lambda x: x + h * drift(x)
    elif method == "semi_implicit" and A is not None:
        eye = torch.eye(len(z), dtype=z.dtype, device=z.device)
        step = lambda x: torch.linalg.solve(eye - h * A, x)
    else:
        raise ValueError(method)
    out, jvp = torch.autograd.functional.jvp(step, z, q)
    return out.detach(), float(torch.log(jvp.norm().clamp_min(1e-15)).item()), q.detach()


class SignPreservingController:
    """Halve h when a nominally stable continuous exponent is estimated positive."""
    def __init__(self, h, stable=True, tolerance=0.0, min_h=1e-5):
        self.h = float(h)
        self.stable = stable
        self.tolerance = tolerance
        self.min_h = min_h
        self.history = []

    def decide(self, lambda_hat):
        self.history.append((self.h, float(lambda_hat)))
        if self.stable and lambda_hat > self.tolerance and self.h > self.min_h:
            self.h *= 0.5
            return "halve"
        return "keep"


def scalar_results(a=1.0, hs=None, n=200000):
    """Exact scalar contraction dz=-a z; continuous lambda=-a.

    Euler lambda_h=log(abs(1-a h))/h and implicit lambda_h=-log(1+a h)/h.
    """
    if hs is None:
        hs = np.array([0.25, 0.5, 1.0, 1.5, 1.9, 2.01, 2.5, 3.0])
    rows = []
    for h in hs:
        le = math.log(abs(1-a*h))/h if abs(1-a*h) > 0 else -math.inf
        li = -math.log(1+a*h)/h
        rows.append({"h": float(h), "euler": le, "implicit": li,
                     "euler_sign": int(np.sign(le)), "implicit_sign": int(np.sign(li)),
                     "euler_error": abs(le+a), "implicit_error": abs(li+a)})
    return rows


def matrix_results(hs=None):
    # Stable damped oscillator: continuous top exponent is real part eigenvalue=-0.2.
    A = np.array([[-0.2, -2.0], [2.0, -0.2]])
    true = -0.2
    if hs is None: hs = np.array([0.05, 0.1, 0.2, 0.4, 0.8])
    rows=[]
    for h in hs:
        Je = np.eye(2)+h*A
        Ji = np.linalg.inv(np.eye(2)-h*A)
        # Long-run exponent is log spectral radius / h; agrees with QR estimate.
        le = math.log(max(abs(np.linalg.eigvals(Je))))/h
        li = math.log(max(abs(np.linalg.eigvals(Ji))))/h
        rows.append({"h":float(h), "euler":float(le), "implicit":float(li),
                     "euler_error":abs(le-true), "implicit_error":abs(li-true),
                     "euler_sign":int(np.sign(le)), "implicit_sign":int(np.sign(li))})
    return true, rows


def main():
    true, matrix = matrix_results()
    scalar = scalar_results()
    # Prediction 1: scalar Euler changes sign exactly at h*=2/a.
    crossing = next(r["h"] for r in scalar if r["euler"] > 0)
    boundary_sweep = []
    for a in [0.5, 1.0, 2.0, 4.0]:
        predicted = 2.0/a
        grid = np.linspace(0.8*predicted, 1.2*predicted, 401)
        observed = next(float(x) for x in grid if math.log(abs(1-a*x))/x > 0)
        boundary_sweep.append({"a":a, "predicted_h_star":predicted,
                               "observed_grid_h_star":observed,
                               "relative_grid_error":abs(observed-predicted)/predicted})
    # Empirical order from small scalar h: error ratio approximately 2 for h -> h/2.
    small = scalar_results(hs=np.array([0.4, 0.2, 0.1, 0.05]))
    ratios = [small[i]["euler_error"] / small[i+1]["euler_error"] for i in range(3)]
    controller = SignPreservingController(2.5, stable=True)
    decisions=[]
    h=controller.h
    while True:
        lam=math.log(abs(1-h))/h
        action=controller.decide(lam)
        decisions.append({"h":h,"lambda_hat":lam,"action":action})
        if action == "keep": break
        h=controller.h
    # Prediction 3: exact JVP should equal the discrete Jacobian action.
    z = torch.tensor([0.7, -0.2], dtype=torch.float64)
    A_t = torch.tensor([[-1.0, 0.0], [0.0, -2.0]], dtype=torch.float64)
    drift = lambda x: A_t @ x
    _, jvp_log, q_used = estimate_with_torch(drift, z, 0.3, method="euler")
    exact_jvp = (torch.eye(2, dtype=torch.float64) + 0.3*A_t) @ q_used
    expected_jvp_log = float(torch.log(exact_jvp.norm()).item())
    out={"continuous_scalar":-1.0, "predicted_euler_boundary":2.0,
         "observed_first_positive_grid_h":crossing, "boundary_parameter_sweep":boundary_sweep,
         "scalar":scalar,
         "jvp_log_error_for_known_linear_map":abs(jvp_log-expected_jvp_log),
         "error_ratios_h_to_h2":ratios, "oscillator_true":true,
         "oscillator":matrix, "controller":decisions}
    with open("results.json","w") as f: json.dump(out,f,indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__": main()