import json
from pathlib import Path
import numpy as np

SEED = 2907


def aug_matrix(A, B, KC, KI, rho):
    return np.array([[A, -B * KC], [KI, rho]], dtype=float)


def spectral_radius(M):
    return float(np.max(np.abs(np.linalg.eigvals(M))))


def unsaturated_decay(A, B, KC, KI, rho, n=100):
    M = aug_matrix(A, B, KC, KI, rho)
    v = np.array([1.0, 0.0])
    norms = []
    for _ in range(n):
        norms.append(float(np.linalg.norm(v)))
        v = M @ v
    # Dominant asymptotic factor estimated from log(norm) after transients.
    fit = np.polyfit(np.arange(25, n), np.log(np.maximum(np.asarray(norms[25:]), 1e-15)), 1)
    return spectral_radius(M), float(np.exp(fit[0]))


def stability_sweep():
    A, B, KC, rho = 0.82, 1.0, 0.32, 0.68
    kis = np.linspace(0.0, 1.8, 181)
    rows = []
    for ki in kis:
        sr = spectral_radius(aug_matrix(A, B, KC, ki, rho))
        M = aug_matrix(A, B, KC, ki, rho)
        v = np.array([1.0, 0.0])
        for _ in range(500):
            v = M @ v
        # A deliberately strict finite-horizon divergence test, not the theory.
        empirical_stable = bool(np.linalg.norm(v) < 1e-3)
        rows.append((float(ki), sr, empirical_stable))
    stable = [x for x in rows if x[2]]
    empirical_boundary = max(x[0] for x in stable) if stable else float('nan')
    crossing = [rows[i] for i in range(1, len(rows))
                if rows[i-1][1] < 1.0 <= rows[i][1]]
    predicted_boundary = crossing[0][0] if crossing else float('nan')
    return {
        'A': A, 'B': B, 'KC': KC, 'rho': rho,
        'predicted_KI_boundary_rho_eq_1': predicted_boundary,
        'observed_KI_boundary_500_step_state_decay_below_1e-3': empirical_boundary,
        'boundary_abs_difference': abs(predicted_boundary - empirical_boundary),
        'spectral_stability_classification_agreement': float(np.mean([x[2] == (x[1] < 1.0) for x in rows])),
    }


def saturation_episode(ki, ka, sat=0.25, n=180, hold_end=95, rho=0.92, kp=0.28, kc=0.55):
    y, c = 0.0, 0.0
    ys, rs, cs, es = [], [], [], []
    for t in range(n):
        yd = 1.0 if 10 <= t < hold_end else 0.0
        e = yd - y
        tilde = kp * e + kc * c  # nominal neural reference z_t is zero
        r = float(np.clip(tilde, -sat, sat))
        d = r - tilde
        y = 0.92 * y + 0.08 * r
        c = rho * c + ki * e + ka * d
        ys.append(y); rs.append(r); cs.append(c); es.append(e)
    es, cs, rs = np.asarray(es), np.asarray(cs), np.asarray(rs)
    recovery = None
    for t in range(hold_end, n):
        if np.max(np.abs(es[t:min(t + 8, n)])) < 0.1:
            recovery = t - hold_end
            break
    return {
        'rms_error': float(np.sqrt(np.mean(es ** 2))),
        'post_release_rms_error': float(np.sqrt(np.mean(es[hold_end:] ** 2))),
        'max_abs_c': float(np.max(np.abs(cs))),
        'release_abs_c': float(abs(cs[hold_end])),
        'recovery_steps': recovery if recovery is not None else n - hold_end,
        'max_constraint_violation': float(np.max(np.maximum(np.abs(rs) - sat, 0.0))),
        'terminal_abs_error': float(abs(es[-1])),
    }


def saturation_sweep():
    durations = np.array([10, 25, 40, 55, 70])
    ordinary = [saturation_episode(.075, 0.0, hold_end=10 + int(d)) for d in durations]
    antiwindup = [saturation_episode(.075, .9, hold_end=10 + int(d)) for d in durations]
    ordinary_c = np.array([x['release_abs_c'] for x in ordinary])
    anti_c = np.array([x['release_abs_c'] for x in antiwindup])
    # Prediction: residual feedback lowers the slope of windup versus saturation duration.
    ordinary_slope = float(np.polyfit(durations, ordinary_c, 1)[0])
    anti_slope = float(np.polyfit(durations, anti_c, 1)[0])
    return {
        'durations': durations.tolist(), 'ordinary_integral': ordinary,
        'residual_antiwindup': antiwindup,
        'predicted_ordering': 'antiwindup slope < ordinary slope',
        'observed_release_c_slope_ordinary': ordinary_slope,
        'observed_release_c_slope_antiwindup': anti_slope,
        'observed_slope_ratio_anti_over_ordinary': anti_slope / ordinary_slope,
    }


def unsaturated_sweep():
    A, B, KC, rho = .80, 1.0, .30, .70
    kis = [.02, .08, .16, .24]
    rows = []
    for ki in kis:
        predicted, observed = unsaturated_decay(A, B, KC, ki, rho)
        rows.append({'KI': ki, 'predicted_spectral_radius': predicted,
                     'observed_asymptotic_decay_factor': observed,
                     'relative_error': abs(observed-predicted)/predicted})
    return rows


def main():
    sat = saturation_sweep()
    result = {
        'seed': SEED,
        'math_prediction_1_unsaturated_decay': unsaturated_sweep(),
        'math_prediction_2_stability_boundary': stability_sweep(),
        'math_prediction_3_saturation_windup_scaling': sat,
        'mini_experiment_baseline_no_compensation': saturation_episode(0.0, 0.0),
        'mini_experiment_ordinary_integral': saturation_episode(.075, 0.0),
        'mini_experiment_idea_residual_antiwindup': saturation_episode(.075, .9),
    }
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()