import json, math
from pathlib import Path
import numpy as np


def rot(theta):
    return np.array([[math.cos(theta), -math.sin(theta)],
                     [math.sin(theta),  math.cos(theta)]], dtype=float)


def exact(v, R):
    n = np.linalg.norm(v)
    return v if n <= R else (R / n) * v


def smooth(v, R):
    # Differentiable radial surrogate from the idea.
    return v / math.sqrt(1.0 + (np.linalg.norm(v) / R) ** 2)


def map_x(x, gain, R, K, theta, limiter):
    # Centered affine map: x0 is an exact fixed point, so perturbation
    # dynamics are governed by the displayed Jacobian at x0.
    Q = rot(theta) @ np.diag([K, 1.0])
    x0 = np.array([1.2, 0.0])
    return x0 + gain * Q @ (limiter(x, R) - limiter(x0, R))


def fd_jacobian(fun, x, h=1e-6):
    J = np.zeros((2, 2))
    for j in range(2):
        xp, xm = x.copy(), x.copy()
        xp[j] += h
        xm[j] -= h
        J[:, j] = (fun(xp) - fun(xm)) / (2.0 * h)
    return J


def analytic_jacobian(R, r, K, theta, kind):
    t = r / R
    Q = rot(theta) @ np.diag([K, 1.0])
    if kind == "exact":
        # radial/tangential derivatives of projection outside the ball
        L = np.diag([0.0, 1.0 / t])
    else:
        # S(v)=v/sqrt(1+(||v||/R)^2)
        radial = (1.0 + t * t) ** (-1.5)
        tangential = (1.0 + t * t) ** (-0.5)
        L = np.diag([radial, tangential])
    return Q @ L


def rho(J):
    return float(np.max(np.abs(np.linalg.eigvals(J))))


def complex_pair(J):
    ev = np.linalg.eigvals(J)
    return bool(abs(ev[0].imag) > 1e-7 and abs(ev[1].imag) > 1e-7)


def predicted_exact_boundary(R, r, theta):
    # For Q=Rot(theta)diag(K,1), exact projection is rank one and
    # rho(g Q DP)=g*|cos(theta)|/(r/R), independent of K.
    return (r / R) / abs(math.cos(theta))


def predicted_smooth_boundary(R, r, K):
    t = r / R
    a = (1.0 + t * t) ** (-1.5)
    s = (1.0 + t * t) ** (-0.5)
    # In the complex-eigenvalue regime, rho(g Q DS)=g*sqrt(K*a*s).
    return 1.0 / math.sqrt(K * a * s)


def trajectory(gain, R, K, theta, kind, n=80, perturb=1e-5):
    limiter = exact if kind == "exact" else smooth
    x = np.array([1.2 + perturb, 0.0])
    norms = []
    for _ in range(n):
        norms.append(float(np.linalg.norm(x)))
        x = map_x(x, gain, R, K, theta, limiter)
        if not np.all(np.isfinite(x)) or np.linalg.norm(x) > 1e12:
            norms.extend([float("inf")] * (n - len(norms)))
            break
    return norms


def main():
    np.random.seed(7)
    R, r, theta = 1.0, 1.2, 1.2
    x0 = np.array([r, 0.0])
    Ks = [1.0, 2.0, 4.0, 8.0]
    rows = []

    # Core math sanity check: analytic Jacobians versus finite differences.
    fd_errors = {}
    for kind, limiter in [("exact", exact), ("smooth", smooth)]:
        for K in [2.0, 4.0]:
            f = lambda x: map_x(x, 1.0, R, K, theta, limiter)
            Ja = analytic_jacobian(R, r, K, theta, kind)
            Jfd = fd_jacobian(f, x0)
            fd_errors[f"{kind}_K{int(K)}"] = float(np.max(np.abs(Ja - Jfd)))

    # Prediction 1: exact threshold independent of anisotropy K.
    exact_pred = predicted_exact_boundary(R, r, theta)
    exact_obs = []
    for K in Ks:
        J = analytic_jacobian(R, r, K, theta, "exact")
        exact_obs.append({"K": K, "predicted_gain": exact_pred,
                          "observed_gain_from_rho": 1.0 / rho(J)})

    # Prediction 2: smooth crossing scales as K^(-1/2), and is complex.
    smooth_rows = []
    for K in Ks:
        pred = predicted_smooth_boundary(R, r, K)
        Junit = analytic_jacobian(R, r, K, theta, "smooth")
        observed = 1.0 / rho(Junit)
        ev = np.linalg.eigvals(pred * Junit)
        smooth_rows.append({"K": K, "predicted_gain": pred,
                            "observed_gain_from_rho": observed,
                            "complex_at_crossing": complex_pair(pred * Junit),
                            "eigenvalues_at_crossing": [[float(z.real), float(z.imag)] for z in ev]})

    # Prediction 3: increasing radius moves both boundaries linearly in r/R
    # for exact projection, while the smooth boundary follows the analytic
    # nonlinear formula; compare a radius sweep directly.
    radius_rows = []
    for rr in [0.8, 1.0, 1.2, 1.6, 2.0]:
        ep = predicted_exact_boundary(R, rr, theta)
        sp = predicted_smooth_boundary(R, rr, 4.0)
        radius_rows.append({"r_over_R": rr / R, "exact_predicted": ep,
                            "smooth_predicted_K4": sp})

    # A direct continuation-style grid demonstrates the guard interval.
    K = 4.0
    ep, sp = exact_pred, predicted_smooth_boundary(R, r, K)
    grid = np.linspace(0.8 * sp, 1.05 * ep, 25)
    for g in grid:
        Jp = g * analytic_jacobian(R, r, K, theta, "exact")
        Js = g * analytic_jacobian(R, r, K, theta, "smooth")
        rows.append({"gain": float(g), "rho_exact": rho(Jp),
                     "rho_smooth": rho(Js), "smooth_complex": complex_pair(Js)})

    # Perturbed trajectories: choose a gain strictly between the smooth and
    # exact crossings, where the predicted artificial instability should occur.
    middle = 0.5 * (sp + ep)
    traj = {kind: trajectory(middle, R, K, theta, kind) for kind in ["exact", "smooth"]}
    result = {
        "setup": {"R": R, "r": r, "theta": theta, "Ks": Ks},
        "finite_difference_max_errors": fd_errors,
        "prediction_1_exact_boundary": exact_obs,
        "prediction_2_smooth_complex_crossing": smooth_rows,
        "prediction_3_radius_sweep": radius_rows,
        "continuation_grid_K4": rows,
        "trajectory_gain_between_boundaries": middle,
        "trajectory_final_norms": {k: v[-1] for k, v in traj.items()},
        "trajectory_max_norms": {k: float(np.max(v)) for k, v in traj.items()},
        "trajectory_norms": traj,
        "guard_interval": {"smooth_crossing": sp, "exact_crossing": ep,
                           "width": ep - sp}
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps({"fd_errors": fd_errors, "exact_boundary": exact_pred,
                      "smooth_K4_boundary": sp, "smooth_K4_complex": smooth_rows[2]["complex_at_crossing"],
                      "guard_width": ep-sp, "trajectory_final_norms": result["trajectory_final_norms"]}, indent=2))


if __name__ == "__main__":
    main()