"""Small reproducible MVP for a truncated causal Volterra stabilizer."""
import json
from itertools import combinations_with_replacement
from pathlib import Path
import numpy as np

SEED = 264


def lag_tuples(L, order):
    """Ordered simplex: 0 <= lag_n <= ... <= lag_1 < L."""
    return list(combinations_with_replacement(range(L), order))


class TruncatedVolterra:
    """Scalar finite-order causal compensator summing kernels of orders 2..N."""
    def __init__(self, L=1, order=2, weights=None):
        self.L, self.order = L, order
        self.weights = {}
        weights = {} if weights is None else weights
        for n in range(2, order + 1):
            self.weights[n] = np.asarray(weights.get(n, np.zeros(len(lag_tuples(L, n)))))

    def __call__(self, history):
        total = 0.0
        for n in range(2, self.order + 1):
            tuples = lag_tuples(self.L, n)
            phi = np.array([np.prod([history[i] for i in t]) for t in tuples])
            total += float(phi @ self.weights[n])
        return total


def verify_math():
    q, c = 0.12, 0.08
    amplitudes = np.geomspace(1e-3, 0.4, 40)
    ratio = abs(c * amplitudes**3) / abs(q * amplitudes**2)
    slope = np.polyfit(np.log(amplitudes), np.log(ratio), 1)[0]
    return {
        "ordered_tuple_counts_L4": {str(n): len(lag_tuples(4, n)) for n in (2, 3)},
        "remainder_over_quadratic_loglog_slope": float(slope),
        "ratio_small_amplitude": float(ratio[0]),
        "ratio_large_amplitude": float(ratio[-1]),
        "stable_linear_transition": bool(abs(0.72) < 1),
    }


def fit_coefficients(seed=SEED):
    rng = np.random.default_rng(seed)
    A, q, c = 0.72, 0.55, 0.22
    x = rng.uniform(-0.35, 0.35, 500)
    y = A*x + q*x*x + c*x**3 + rng.normal(0, 0.001, len(x))
    residual = y - A*x
    q2 = float(np.linalg.lstsq((x*x)[:, None], residual, rcond=None)[0][0])
    q3, c3 = np.linalg.lstsq(np.column_stack([x*x, x**3]), residual, rcond=None)[0]
    return A, q, c, q2, float(q3), float(c3)


def rollout(A, q, c, compensator, x0, steps=1000):
    history = [float(x0)] * compensator.L
    trajectory = []
    for _ in range(steps):
        h = history[0]
        nxt = A*h + q*h*h + c*h**3 - compensator(history)
        history = [nxt] + history[:-1]
        trajectory.append(nxt)
        if not np.isfinite(nxt) or abs(nxt) > 1e12:
            return np.asarray(trajectory + [np.inf] * (steps-len(trajectory)))
    return np.asarray(trajectory)


def sweep(A, q, c, q2, q3, c3, amplitudes):
    controllers = {
        "baseline": TruncatedVolterra(L=1, order=2, weights={2: [0.0]}),
        "quadratic": TruncatedVolterra(L=1, order=2, weights={2: [q2]}),
        "cubic": TruncatedVolterra(L=1, order=3, weights={2: [q3], 3: [c3]}),
    }
    result = {}
    for name, ctrl in controllers.items():
        final, maxima, divergent = [], [], 0
        for x0 in amplitudes:
            z = rollout(A, q, c, ctrl, x0)
            bad = (not np.all(np.isfinite(z))) or np.max(np.abs(z)) > 1e6
            divergent += int(bad)
            final.append(np.inf if bad else abs(float(z[-1])))
            maxima.append(np.inf if bad else float(np.max(np.abs(z))))
        finite = np.isfinite(final)
        result[name] = {
            "mean_final_abs_state": float(np.mean(np.asarray(final)[finite])) if finite.any() else float("inf"),
            "max_abs_state": float(np.max(np.asarray(maxima)[finite])) if finite.any() else float("inf"),
            "divergent_rollouts": divergent,
            "bounded_initial_amplitude_max": float(amplitudes[np.where(finite)[0][-1]]) if finite.any() else 0.0,
            "parameter_count": 0 if name == "baseline" else (1 if name == "quadratic" else 2),
        }
    return result


def run():
    A, q, c, q2, q3, c3 = fit_coefficients()
    amplitudes = np.geomspace(0.02, 3.0, 60)
    return {
        "seed": SEED,
        "plant": {"A": A, "quadratic": q, "cubic": c},
        "learned_coefficients": {"qhat_quadratic": q2, "qhat_cubic": q3, "chat_cubic": c3},
        "math_check": verify_math(),
        "rollout_steps": 1000,
        "amplitude_sweep": [float(x) for x in amplitudes],
        "rollout_metrics": sweep(A, q, c, q2, q3, c3, amplitudes),
    }


if __name__ == "__main__":
    result = run()
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))