"""Fading-memory habituation gate: math checks and a tiny controlled experiment."""
import json, math, random
from pathlib import Path
import numpy as np

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


def simulate(rho, beta, stim, n=None, a0=0.0):
    """Update fading memory then apply nonlinear gain."""
    if n is not None:
        stim = [stim] * n
    a = float(a0)
    states, gains = [], []
    for s in stim:
        a = rho * a + (1.0 - rho) * float(s)
        states.append(a)
        gains.append(1.0 / (1.0 + beta * a))
    return np.asarray(states), np.asarray(gains)


def half_life(rho):
    return math.log(0.5) / math.log(rho)


def main():
    # Prediction 1: constant stimulation converges to a=s and gain 1/(1+beta*s).
    steady_rows = []
    for rho in [0.5, 0.8, 0.9, 0.97, 0.99]:
        # Enough steps for the slowest pole to settle accurately.
        n = 3000
        for beta in [0.5, 2.0, 5.0]:
            s = 1.7
            a, g = simulate(rho, beta, s, n=n)
            pred = 1.0 / (1.0 + beta * s)
            steady_rows.append({
                "rho": rho, "beta": beta,
                "a_error": abs(float(a[-1]) - s),
                "gain_observed": float(g[-1]),
                "gain_predicted": pred,
                "gain_abs_error": abs(float(g[-1]) - pred),
            })

    # Prediction 2: after withdrawal a_k/a_0=rho^k; recovery half-life is ln(.5)/ln(rho).
    recovery_rows = []
    for rho in [0.5, 0.8, 0.9, 0.97, 0.99]:
        a0 = 2.3
        kmax = 3000
        states, _ = simulate(rho, 0.0, [0.0] * kmax, a0=a0)
        ratios = states / a0
        observed = (int(np.flatnonzero(ratios <= 0.5)[0]) + 1) if np.any(ratios <= 0.5) else None
        predicted = half_life(rho)
        predicted_index = int(math.ceil(predicted))
        # Fit log ratio versus k; exact prediction is slope log(rho).
        ks = np.arange(1, min(kmax, max(20, int(predicted * 8))))
        slope = float(np.polyfit(ks, np.log(ratios[ks - 1]), 1)[0])
        recovery_rows.append({
            "rho": rho, "half_life_predicted": predicted,
            "half_index_observed": observed,
            "half_index_predicted": predicted_index,
            "half_index_relative_error": abs(observed - predicted_index) / max(1, predicted_index),
            "log_slope_observed": slope,
            "log_slope_predicted": math.log(rho),
        })

    # Prediction 3: rho<1 is stable/bounded for nonnegative bounded stimulation;
    # rho=1 freezes the state and rho>1 grows after withdrawal.
    stability_rows = []
    for rho in [0.5, 0.9, 0.99, 1.0, 1.01, 1.1]:
        states, _ = simulate(rho, 1.0, [0.0] * 100, a0=1.0)
        # Direct withdrawal exposes the pole: rho=1 does not decay and rho>1 grows.
        post = states
        stability_rows.append({
            "rho": rho,
            "constant_final_state": float(states[-1]),
            "withdrawal_final_state": float(post[-1]),
            "predicted_pole_behavior": "decay" if rho < 1 else ("constant" if rho == 1 else "growth"),
            "bounded_over_run": bool(np.all(np.isfinite(np.r_[states, post])) and np.max(np.abs(np.r_[states, post])) < 1e6),
        })

    # Controlled mini-experiment: repeated familiar pulses versus novel pulses.
    # The baseline passes amplitude unchanged. The proposed gate suppresses repeated
    # stimulation while its state is reset between independent sequences.
    rho, beta, s = 0.9, 2.0, 1.0
    familiar = [s] * 20 + [0.0] * 20 + [s]
    novel = [0.0] * 20 + [s]
    a_f, g_f = simulate(rho, beta, familiar)
    a_n, g_n = simulate(rho, beta, novel)
    baseline_f = np.ones(len(familiar))
    baseline_n = np.ones(len(novel))
    # Matched linear-in-state gate: choose c so its steady gain equals the
    # nonlinear gate at s. This is not a true LTI input-output system because
    # it multiplies v by a state-dependent factor, but is a useful ablation.
    c = beta / (1.0 + beta * s)
    a_lin, _ = simulate(rho, 0.0, familiar)
    lin_g_f = 1.0 - c * a_lin
    a_lin_n, _ = simulate(rho, 0.0, novel)
    lin_g_n = 1.0 - c * a_lin_n
    comparison = {
        "rho": rho, "beta": beta, "matched_linear_c": c,
        "baseline_repeated_pulse_gain": float(baseline_f[19]),
        "idea_repeated_pulse_gain": float(g_f[19]),
        "linear_repeated_pulse_gain": float(lin_g_f[19]),
        "baseline_post_withdrawal_gain": float(baseline_f[39]),
        "idea_post_withdrawal_gain": float(g_f[39]),
        "linear_post_withdrawal_gain": float(lin_g_f[39]),
        "baseline_novel_gain": float(baseline_n[-1]),
        "idea_novel_gain_after_reset": float(g_n[-1]),
        "linear_novel_gain_after_reset": float(lin_g_n[-1]),
        "repeated_suppression_ratio": float(g_f[19] / baseline_f[19]),
        "novel_retention_ratio": float(g_n[-1] / baseline_n[-1]),
    }

    result = {
        "seed": SEED,
        "predictions": {
            "steady_state": steady_rows,
            "recovery": recovery_rows,
            "stability": stability_rows,
        },
        "comparison": comparison,
    }
    out = Path("results.json")
    out.write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()