import json, math
from pathlib import Path
import numpy as np

SEED = 2710
rng = np.random.default_rng(SEED)

# Three-state clockwise cycle: i -> i+1 at a and i -> i-1 at b.
# Its stationary distribution is uniform, sigma=(a-b) log(a/b),
# and the nonconstant output mode has tau=1/(1.5*(a+b)).
def exact_stats(a, b):
    sigma = (a - b) * math.log(a / b)
    tau = 1.0 / (1.5 * (a + b))
    return sigma, tau

def simulate_ctmc(a, b, dt=0.01, steps=300000, burn=30000):
    """Uniformized discrete simulation; estimates p, q, sigma and output tau."""
    z = 0
    states = np.empty(steps-burn, dtype=np.int8)
    # y is a centered observable with nonzero projection on both complex modes.
    yy = np.array([1.0, -0.5, -0.5])
    counts = np.zeros((3,3), dtype=np.int64)
    occup = np.zeros(3, dtype=np.int64)
    # Exact one-step transition probabilities to avoid multiple jumps.
    # dt is small for all rates used below.
    for t in range(steps):
        u = rng.random()
        if u < a*dt:
            zn = (z + 1) % 3
        elif u < (a+b)*dt:
            zn = (z - 1) % 3
        else:
            zn = z
        if t >= burn:
            states[t-burn] = z
            occup[z] += 1
            if zn != z:
                counts[z, zn] += 1
        z = zn
    p = occup / occup.sum()
    qhat = counts / (occup[:, None] * dt + 1e-30)
    sig = 0.0
    for i in range(3):
        for j in range(i+1, 3):
            f = p[i]*qhat[i,j]
            g = p[j]*qhat[j,i]
            if f > 0 and g > 0:
                sig += (f-g)*math.log(f/g)
    # Correlation integral via FFT, truncated where correlation is negligible.
    y = yy[states].astype(float)
    y -= y.mean()
    n = len(y)
    f = np.fft.rfft(y, 2*n)
    ac = np.fft.irfft(f*np.conjugate(f), 2*n)[:n]
    ac /= np.arange(n, 0, -1)
    ac /= ac[0]
    cutoff = np.where(ac < 0.02)[0]
    m = int(cutoff[0]) if len(cutoff) else min(n-1, int(10/(1.5*(a+b)*dt)))
    tauhat = max(0.0, np.trapz(ac[:m], dx=dt))
    return sig, tauhat, p

def tracking_mse(rate, freq, amp=1.0, dt=0.002, duration=80.0):
    """Mean-field response of the stochastic mode to a changing target.
    This is the exact mean response of a symmetric reset-like Markov decoder.
    """
    n = int(duration/dt)
    t = np.arange(n)*dt
    target = amp*np.sin(2*np.pi*freq*t)
    out = 0.0
    err = np.empty(n)
    for k in range(n):
        out += dt*rate*(target[k]-out)
        err[k] = out-target[k]
    burn = n//4
    return float(np.mean(err[burn:]**2)), float(np.mean((2*np.pi*freq*amp*np.cos(2*np.pi*freq*t[burn:]))**2))

def main():
    # Cheap sanity check of the exact entropy-production expression on random rates.
    random_checks = []
    for _ in range(1000):
        a, b = 10**rng.uniform(-1, 1, 2)
        s, _ = exact_stats(a,b)
        random_checks.append(s >= -1e-12)

    # Prediction 1: at fixed bias ratio, sigma is linear in global activity.
    ratio = 3.0
    rates = np.array([0.25, 0.5, 1.0, 2.0, 4.0])
    sigmas = np.array([exact_stats(r*ratio/(ratio+1), r/(ratio+1))[0] for r in rates])
    slope_sigma = np.polyfit(np.log(rates), np.log(sigmas), 1)[0]

    # Prediction 2: the output memory time is inverse in activity.
    taus = np.array([exact_stats(r*ratio/(ratio+1), r/(ratio+1))[1] for r in rates])
    slope_tau = np.polyfit(np.log(rates), np.log(taus), 1)[0]

    # Estimate both expressions once, checking numerical estimators against exact values.
    a, b = ratio/(ratio+1), 1/(ratio+1)
    exact_s, exact_tau = exact_stats(a,b)
    est_s, est_tau, p = simulate_ctmc(a,b)

    # Prediction 3: claimed budget says epsilon*sigma >= kappa*tau*v^2.
    # Test the dimensionless ratio with kappa=1 over rate and frequency sweeps.
    budget_rows = []
    for r in [0.5, 1.0, 2.0, 4.0]:
        aa, bb = r*ratio/(ratio+1), r/(ratio+1)
        s, tau = exact_stats(aa,bb)
        for f in [0.01, 0.03, 0.1, 0.3]:
            eps, v2 = tracking_mse(r, f)
            budget_rows.append({"rate":r, "freq":f, "sigma":s, "tau":tau,
                                "epsilon":eps, "v2":v2,
                                "budget_ratio":eps*s/(tau*v2)})

    # Tiny baseline comparison: fixed slow deterministic low-pass vs budget-selected activity.
    # The proposed controller selects the smallest tested rate whose estimated budget ratio
    # exceeds one; this is deliberately reported as a test, not a fitted claim.
    base = []
    idea = []
    for f in [0.03, 0.1, 0.3]:
        eb, _ = tracking_mse(1.0, f)
        choices = [(r, tracking_mse(r,f)[0]) for r in [0.5,1,2,4]]
        chosen = next((x for x in choices if next(q["budget_ratio"] for q in budget_rows if q["rate"]==x[0] and q["freq"]==f) >= 1), choices[-1])
        base.append(eb); idea.append(chosen[1])

    result = {
      "seed": SEED,
      "math_sanity": {"all_random_entropy_rates_nonnegative": all(random_checks),
                       "exact_sigma": exact_s, "estimated_sigma": est_s,
                       "exact_tau": exact_tau, "estimated_tau": est_tau,
                       "stationary_p": p.tolist()},
      "predictions": {
        "P1_sigma_activity_exponent_predicted": 1.0,
        "P1_sigma_activity_exponent_observed": float(slope_sigma),
        "P2_tau_activity_exponent_predicted": -1.0,
        "P2_tau_activity_exponent_observed": float(slope_tau),
        "P3_budget_ratio_should_be_ge_kappa_1": True,
        "P3_budget_ratio_min": float(min(x["budget_ratio"] for x in budget_rows)),
        "P3_budget_ratio_median": float(np.median([x["budget_ratio"] for x in budget_rows])),
        "P3_budget_ratio_max": float(max(x["budget_ratio"] for x in budget_rows))
      },
      "budget_sweep": budget_rows,
      "mini_comparison": {"frequencies":[0.03,0.1,0.3], "baseline_fixed_rate1_mse":base,
                          "budget_controller_mse":idea}
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()