import json
from pathlib import Path
import numpy as np

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

class HystereticRouter:
    def __init__(self, eta_on=0.6, eta_off=0.4, active=0):
        self.on, self.off, self.active = eta_on, eta_off, active
        self.switches = []
    def step(self, h, t):
        # A switch is permitted only after the current attractor has exited.
        if h[self.active] <= self.off:
            candidates = [i for i in range(len(h)) if i != self.active and h[i] >= self.on]
            if candidates:
                old = self.active
                self.active = int(candidates[np.argmax(h[candidates])])
                self.switches.append((t, old, self.active))
        return self.active

def dwell_sweep():
    # Scores are deliberately generated with a known bounded derivative.
    # Candidate rises from off to on, then the old score falls; traversal takes Delta/L.
    L = 0.017
    rows = []
    for delta in [0.04, 0.08, 0.14, 0.22]:
        off, on = 0.40, 0.40 + delta
        dt = 0.01
        n = 30000
        t = np.arange(n) * dt
        # candidate 1 rises while active score 0 falls, each slope <= L
        h0 = np.clip(0.75 - L*t, 0, 1)
        h1 = np.clip(0.20 + L*t, 0, 1)
        router = HystereticRouter(on, off)
        for k in range(n): router.step(np.array([h0[k], h1[k]]), t[k])
        assert router.switches
        first = router.switches[0][0]
        # Now explicitly measure the candidate's time from off to on.
        expected = delta / L
        observed = expected # continuous ramp identity, independently verified below
        # finite-difference derivative and actual threshold crossing times
        rise_on = np.where(h1 >= on)[0][0] * dt
        rise_off = np.where(h1 >= off)[0][0] * dt
        observed = rise_on - rise_off
        max_deriv = np.max(np.abs(np.diff(h1) / dt))
        rows.append({'delta': delta, 'predicted_dwell': expected,
                     'observed_dwell': observed, 'bound_ratio': observed / expected,
                     'max_fd_derivative': float(max_deriv), 'switch_time': float(first)})
    return rows

def chatter_sweep():
    # A sinusoidal score of amplitude A around the threshold. Without hysteresis,
    # crossings occur repeatedly. If Delta > 2A, the score cannot traverse both thresholds.
    dt, duration, A, period = 0.01, 40.0, 0.08, 2.0
    t = np.arange(0, duration, dt)
    s = 0.5 + A*np.sin(2*np.pi*t/period)
    rows = []
    for delta in [0.00, 0.04, 0.10, 0.16, 0.20]:
        off, on = 0.5-delta/2, 0.5+delta/2
        # use complementary scores, which makes each threshold crossing meaningful
        r = HystereticRouter(on, off)
        for k, v in enumerate(s): r.step(np.array([1-v, v]), t[k])
        observed = len(r.switches)
        predicted_zero = delta > 2*A
        rows.append({'delta': delta, 'noise_amplitude': A,
                     'predicted_zero_switch': predicted_zero,
                     'switches': observed})
    return rows

def tau_sweep():
    # First-order h update h'=(s-h)/tau. For sinusoidal score, larger tau attenuates
    # fluctuations; measured RMS gain should follow 1/sqrt(1+(omega*tau)^2).
    dt, duration, A, period = 0.001, 20.0, 0.1, 1.0
    t = np.arange(0, duration, dt)
    omega = 2*np.pi/period
    s = 0.5 + A*np.sin(omega*t)
    rows=[]
    for tau in [0.005, 0.02, 0.1, 0.5]:
        h=0.5; vals=[]
        for v in s:
            h += dt/tau*(v-h)
            vals.append(h)
        vals=np.asarray(vals); trim=t > 5
        gain=(np.std(vals[trim])/A)
        predicted=1/np.sqrt(2*(1+(omega*tau)**2))
        rows.append({'tau_h': tau, 'predicted_rms_gain': predicted,
                     'observed_gain': float(gain), 'relative_error': float(abs(gain-predicted)/predicted)})
    return rows

def regime_prediction_demo():
    # Same observed score is used by both methods. Hysteresis suppresses rapid regime
    # changes caused by score noise; soft routing averages incompatible dynamics.
    rng=np.random.default_rng(SEED+7); T=6000
    true=(np.arange(T)//30)%2
    # regimes have opposite slopes, noisy observations make scores ambiguous at boundaries
    x=np.zeros(T); x[0]=0.2
    for t in range(1,T): x[t]=0.92*x[t-1] + (0.18 if true[t] else -0.18) + 0.04*rng.normal()
    score=np.clip(true + 0.22*rng.normal(size=T),0,1)
    # one-step model uses known regime-specific intercepts, but route is noisy
    pred_soft=0.92*x[:-1] + 0.18*(2*score[1:]-1)
    soft_mse=float(np.mean((pred_soft-x[1:])**2))
    r=HystereticRouter(0.68,0.32); chosen=[]
    for t in range(T): chosen.append(r.step(np.array([1-score[t],score[t]]),t))
    chosen=np.asarray(chosen)
    pred_h=0.92*x[:-1] + 0.18*(2*chosen[1:]-1)
    hyst_mse=float(np.mean((pred_h-x[1:])**2))
    return {'soft_mse':soft_mse, 'hysteretic_mse':hyst_mse,
            'soft_switch_proxy':int(np.sum(np.diff((score>0.5).astype(int))!=0)),
            'hysteretic_switches':len(r.switches), 'improvement_fraction':(soft_mse-hyst_mse)/soft_mse}

def main():
    dwell=dwell_sweep(); chatter=chatter_sweep(); tau=tau_sweep(); demo=regime_prediction_demo()
    result={'seed':SEED, 'dwell_bound':dwell, 'chatter_threshold':chatter,
            'timescale_lowpass':tau, 'regime_demo':demo}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))
    # core acceptance checks
    assert min(x['bound_ratio'] for x in dwell) >= 0.99
    assert chatter[-1]['switches'] == 0
    assert max(x['relative_error'] for x in tau) < 0.02
if __name__ == '__main__': main()