import json
import math
import random
from pathlib import Path

import numpy as np

SEED = 982
np.random.seed(SEED)
random.seed(SEED)


def q_weights(q, alpha, K):
    """Finite truncation of (q^alpha;q)_k/(q;q)_k from the paper."""
    w = np.empty(K, dtype=float)
    w[0] = 1.0
    for k in range(1, K):
        w[k] = w[k - 1] * (1.0 - q ** (alpha + k - 1)) / (1.0 - q ** k)
    return w


def companion(g, weights):
    K = len(weights)
    M = np.zeros((K, K), dtype=float)
    M[0, :] = g * weights
    M[1:, :-1] = np.eye(K - 1)
    return M


def rho(g, weights):
    return float(np.max(np.abs(np.linalg.eigvals(companion(g, weights)))))


def controller(g, rho_hat, delta=0.05, eps=1e-8, g_max=10.0, eta_gain=0.01):
    target = 1.0 - delta
    # Conservative branch from the proposal; optional slow growth is only used
    # when the estimate is comfortably below the target.
    if rho_hat > target:
        return max(0.0, g * min(1.0, target / (rho_hat + eps)))
    return min(g_max, g * (1.0 + eta_gain))


def threshold_sweep(weights):
    predicted = 1.0 / float(np.sum(weights))
    gs = predicted * np.linspace(0.70, 1.30, 25)
    rs = np.array([rho(float(g), weights) for g in gs])
    # Linear interpolation of the closest sign change in rho-1.
    j = int(np.argmin(np.abs(rs - 1.0)))
    if j == 0 or j == len(gs) - 1:
        observed = float(gs[j])
    else:
        observed = float(gs[j - 1] + (1-rs[j-1]) * (gs[j]-gs[j-1])/(rs[j]-rs[j-1]))
    return predicted, observed, gs, rs


def gradient_slope(weights, g, n=160):
    M = companion(g, weights)
    # A generic perturbation avoids selecting a special eigenvector.
    v = np.linspace(1.0, 0.3, len(weights))
    vals = []
    for _ in range(n):
        vals.append(np.linalg.norm(v))
        v = M @ v
    vals = np.maximum(np.asarray(vals), 1e-300)
    # Ignore the transient and fit the asymptotic log slope.
    slope = float(np.polyfit(np.arange(40, n), np.log(vals[40:]), 1)[0])
    return slope


def run_controller(weights, g0, delta=0.05, steps=20):
    g = float(g0)
    rows = []
    for t in range(steps):
        r = rho(g, weights)
        rows.append((t, g, r))
        g = controller(g, r, delta=delta)
    return rows


def main():
    q, alpha, K = 0.7, 0.5, 30
    w = q_weights(q, alpha, K)
    predicted, observed, gs, rs = threshold_sweep(w)

    # Prediction 1: positive-memory boundary is g*S=1 (z=1).
    boundary_rel_error = abs(observed - predicted) / predicted
    below = rho(0.90 * predicted, w)
    above = rho(1.10 * predicted, w)

    # Prediction 2: near the boundary, rho crosses one monotonically.
    monotone = bool(np.all(np.diff(rs) > -1e-9))
    crossing = bool(below < 1.0 and above > 1.0)

    # Prediction 3: gradient norm has asymptotic slope log(rho).
    g_stable = 0.80 * predicted
    g_unstable = 1.20 * predicted
    r_stable, r_unstable = rho(g_stable, w), rho(g_unstable, w)
    slope_stable = gradient_slope(w, g_stable)
    slope_unstable = gradient_slope(w, g_unstable)
    slope_err = abs(slope_stable - math.log(r_stable))

    # Controller starts 35% above the predicted boundary and should move to margin.
    rows = run_controller(w, 1.35 * predicted, delta=0.05, steps=50)
    final_g, final_r = rows[-1][1], rows[-1][2]
    controller_target = 0.95
    controller_ok = final_r <= controller_target + 0.005

    # A tiny nonlinear rollout provides an interpretable baseline comparison.
    # Both systems use the same initial state; controller rescales g when rho is high.
    x0 = np.zeros(K); x0[0] = 1.0
    def rollout(controlled):
        state, g, maxabs = x0.copy(), 1.35 * predicted, 1.0
        losses = []
        for t in range(80):
            r = rho(g, w)
            if controlled:
                g = controller(g, r, delta=0.05, eta_gain=0.0)
            state = companion(g, w) @ state
            maxabs = max(maxabs, float(np.max(np.abs(state))))
            losses.append(float(np.sum(state * state)))
        return maxabs, losses[-1]
    base_max, base_loss = rollout(False)
    ctrl_max, ctrl_loss = rollout(True)

    report = {
        "setup": {"q": q, "alpha": alpha, "K": K, "weight_sum": float(np.sum(w))},
        "prediction_1_boundary": {
            "predicted_gcrit_1_over_sum_w": predicted,
            "observed_interpolated_gcrit": observed,
            "relative_error": boundary_rel_error,
            "pass_within_20pct": boundary_rel_error <= 0.20,
        },
        "prediction_2_transition": {
            "rho_at_0.90_gcrit": below,
            "rho_at_1.10_gcrit": above,
            "crosses_unit_circle": crossing,
            "rho_sweep_monotone": monotone,
        },
        "prediction_3_gradient_scaling": {
            "stable_g": g_stable, "stable_rho": r_stable,
            "measured_log_slope": slope_stable,
            "predicted_log_slope": math.log(r_stable),
            "absolute_slope_error": slope_err,
            "unstable_g": g_unstable, "unstable_rho": r_unstable,
            "unstable_measured_log_slope": slope_unstable,
        },
        "controller": {
            "initial_g": rows[0][1], "initial_rho": rows[0][2],
            "final_g": final_g, "final_rho": final_r,
            "target_rho": controller_target, "pass": controller_ok,
            "trajectory": [{"step": t, "g": g, "rho": r} for t, g, r in rows],
        },
        "mini_comparison": {
            "uncontrolled_max_state": base_max, "uncontrolled_final_squared_norm": base_loss,
            "controlled_max_state": ctrl_max, "controlled_final_squared_norm": ctrl_loss,
        },
        "mechanism_manifested": bool(boundary_rel_error <= .20 and crossing and monotone and slope_err <= .03 and controller_ok),
    }
    Path("results.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()