import json
import math
import numpy as np


def rhs(x, lam=0.8, q=0.7):
    return -lam * x + q * x * x


def exact_lift(x, d):
    return np.array([x ** k for k in range(1, d + 1)], dtype=float)


def K_matrix(d, lam=0.8, q=0.7):
    # For z_k=x^k, d/dt z_k=-k*lam*z_k + k*q*z_{k+1}.
    K = np.zeros((d, d))
    for k in range(1, d + 1):
        K[k - 1, k - 1] = -k * lam
        if k < d:
            K[k - 1, k] = k * q
    return K


def residual(x, d, lam=0.8, q=0.7):
    z = exact_lift(x, d)
    dz = np.array([k * x ** (k - 1) * rhs(x, lam, q) for k in range(1, d + 1)])
    return dz - K_matrix(d, lam, q) @ z


def scaling_sweep():
    # At fixed shape x=a*x0, ||r_d|| must scale as a^(d+1).
    amps = np.geomspace(0.15, 1.2, 9)
    rows = []
    for d in (2, 3, 4):
        vals = np.array([np.linalg.norm(residual(a, d)) for a in amps])
        slope = np.polyfit(np.log(amps), np.log(vals), 1)[0]
        rows.append({'d': d, 'predicted_exponent': d + 1, 'observed_exponent': float(slope),
                     'max_relative_fit_error': float(np.max(np.abs(vals / (vals[0] * (amps / amps[0]) ** (d + 1)) - 1)))})
    # q enters linearly, including the q=0 vanishing prediction.
    qs = np.geomspace(0.1, 1.6, 8)
    vals = np.array([np.linalg.norm(residual(0.55, 3, q=q)) for q in qs])
    q_slope = np.polyfit(np.log(qs), np.log(vals), 1)[0]
    return {'amplitude': rows, 'q_scaling': {'predicted_exponent': 1, 'observed_exponent': float(q_slope)},
            'zero_q_residual_norm': float(np.linalg.norm(residual(0.55, 3, q=0.0)))}


def rho(x, d, lam=0.8, q=0.7, eps=1e-10):
    z = exact_lift(x, d)
    return np.linalg.norm(residual(x, d, lam, q)) / (np.linalg.norm(K_matrix(d, lam, q) @ z) + eps)


def euler_truth(x, dt, lam, q):
    return x + dt * rhs(x, lam, q)


def lifted_step(z, d, dt, lam, q):
    return z + dt * (K_matrix(d, lam, q) @ z)


def rollout(x0, d, steps, dt, lam, q, gate=False, threshold=0.0, dwell=2):
    x = float(x0)
    active = 2 if gate else d
    z = exact_lift(x, active)
    over = 0
    pred = []
    truth = []
    active_counts = []
    for t in range(steps):
        truth.append(x)
        # Observation-based certificate; in a real deployment x_next is measured.
        x_next = euler_truth(x, dt, lam, q)
        if gate and active == 2:
            r = rho(x, 2, lam, q)
            over = over + 1 if r > threshold else 0
            if over >= dwell:
                active = 3
                z = exact_lift(x, 3)
        if active == 2:
            z = lifted_step(z, 2, dt, lam, q)
        else:
            z = lifted_step(z, 3, dt, lam, q)
        pred.append(float(z[0]))
        active_counts.append(active)
        x = x_next
    return np.array(pred), np.array(truth), np.array(active_counts)


def experiment():
    lam, q, dt, steps, x0 = 0.8, 0.7, 0.04, 50, 0.35
    # Calibrate up threshold as requested from a low-amplitude training window.
    train_x = np.linspace(0.05, 0.35, 100)
    train_rhos = np.array([rho(x, 2, lam, q) for x in train_x])
    tau = float(np.quantile(train_rhos, 0.95))
    results = {}
    for name, d, gate in [('fixed_d2', 2, False), ('fixed_d3', 3, False), ('gated', 2, True)]:
        pred, truth, active = rollout(x0, d, steps, dt, lam, q, gate, tau, 2)
        results[name] = {'one_step_like_final_error': float(abs(pred[0] - truth[1])),
                         'rollout_rmse': float(np.sqrt(np.mean((pred - truth) ** 2))),
                         'final_abs_error': float(abs(pred[-1] - truth[-1])),
                         'mean_active_degree': float(np.mean(active)),
                         'active_feature_count_mean': float(np.mean(active))}
    results['settings'] = {'tau_up_95pct': tau, 'initial_rho': rho(x0, 2, lam, q),
                           'final_rho_d2': rho(0.12, 2, lam, q), 'steps': steps, 'dt': dt}
    # Falsifiable gate transition: predict activation from the certificate sequence,
    # then compare with the controller's actual first activation time.
    transition = []
    for x_init in np.linspace(0.08, 0.55, 10):
        xs = []
        x_tmp = float(x_init)
        for _ in range(50):
            xs.append(x_tmp)
            x_tmp = euler_truth(x_tmp, dt, lam, q)
        cert = np.array([rho(xx, 2, lam, q) for xx in xs])
        pred_t = None
        for t in range(len(cert) - 1):
            if cert[t] > tau and cert[t + 1] > tau:
                pred_t = t
                break
        _, _, act = rollout(x_init, 2, 50, dt, lam, q, True, tau, 2)
        obs = np.where(act >= 3)[0]
        obs_t = int(obs[0]) if len(obs) else None
        transition.append({'x0': float(x_init), 'predicted_activation_step': pred_t,
                           'observed_activation_step': obs_t,
                           'rho_x0': float(cert[0])})
    results['gate_transition_sweep'] = transition
    return results


def finite_difference_check():
    lam, q, dt = 0.8, 0.7, 1e-6
    rows = []
    for d in (2, 3, 4):
        x = 0.41
        fd = (exact_lift(euler_truth(x, dt, lam, q), d) - exact_lift(x, d)) / dt
        analytic = K_matrix(d, lam, q) @ exact_lift(x, d) + residual(x, d, lam, q)
        rows.append({'d': d, 'relative_error': float(np.linalg.norm(fd - analytic) / np.linalg.norm(analytic))})
    return rows


if __name__ == '__main__':
    out = {'math_verification': scaling_sweep(), 'finite_difference_check': finite_difference_check(), 'mini_experiment': experiment()}
    print(json.dumps(out, indent=2))

    out = {'math_verification': scaling_sweep(), 'mini_experiment': experiment()}
    print(json.dumps(out, indent=2))