import json
import numpy as np

R = np.array([[0.0, 1.0], [-1.0, 0.0]])
GAIN = 1.04
DAMP = 0.8

def g(x):
    x = np.asarray(x, float)
    q = 1.0 + DAMP * (x @ x)
    return GAIN * (R @ x) / q

def jac(x):
    x = np.asarray(x, float)
    q = 1.0 + DAMP * (x @ x)
    z = R @ x
    # derivative of GAIN*(R x)/(1+DAMP*x^T x)
    return GAIN * (R / q - np.outer(z, 2.0 * DAMP * x) / (q*q))

def nonlinear(x0, H):
    xs = [np.asarray(x0, float)]
    for _ in range(H):
        xs.append(g(xs[-1]))
    return np.asarray(xs)

def frozen(x0, H, radius=None):
    A = jac(x0)
    c = g(x0) - A @ x0
    rawrho = max(abs(np.linalg.eigvals(A)))
    if radius is not None and rawrho > radius:
        A *= radius / rawrho
        c = g(x0) - A @ x0
    rho = max(abs(np.linalg.eigvals(A)))
    xs = [np.asarray(x0, float)]
    for _ in range(H):
        xs.append(A @ xs[-1] + c)
    return np.asarray(xs), rawrho, rho

def main():
    H = 100
    x0 = np.array([0.1, 0.0])
    truth = nonlinear(x0, H)
    raw, rawrho, _ = frozen(x0, H, None)
    clipped, _, cliprho = frozen(x0, H, 0.98)
    d = np.array([0.7, -0.3]); h = 1e-6
    fd = (g(x0 + h*d) - g(x0))/h
    out = {
      'initial_jacobian_spectral_radius': float(rawrho),
      'clipped_spectral_radius': float(cliprho),
      'horizon': H,
      'final_norms': {'nonlinear': float(np.linalg.norm(truth[-1])),
                      'frozen_unclipped': float(np.linalg.norm(raw[-1])),
                      'frozen_clipped': float(np.linalg.norm(clipped[-1]))},
      'max_norms': {'nonlinear': float(np.max(np.linalg.norm(truth, axis=1))),
                    'frozen_unclipped': float(np.max(np.linalg.norm(raw, axis=1))),
                    'frozen_clipped': float(np.max(np.linalg.norm(clipped, axis=1)))},
      'finite_difference_jacobian_error': float(np.linalg.norm(fd - jac(x0) @ d))
    }
    open('unstable_results.json','w').write(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()