import json, math
import numpy as np

# Scalar quadratic: f(x)=x^2/2. A worker applies a gradient from d steps ago.
# The IMM observes x with Gaussian noise and controls the learning rate after alarm.
def detect(obs, q=.002, r=.01, D=4, persistence=.985, threshold=.01):
    m=D+1; T=np.full((m,m),(1-persistence)/(m-1)); np.fill_diagonal(T,persistence)
    pi=np.ones(m)/m; hist=[0.]*(D+1); alarms=0
    out=[]
    for y in obs:
        hist.append(hist[-1]); hist=hist[-D-1:]
        c=T.T@pi; means=np.asarray(hist[::-1])
        ll=-.5*((y-means)**2/r+math.log(2*math.pi*r))
        u=c*np.exp(ll-ll.max()); pi=u/u.sum(); alarms=alarms+1 if pi[0]<=threshold else 0
        out.append((pi.copy(), alarms>=3))
    return out

def run(lr, delay, controlled, seed):
    rng=np.random.default_rng(seed); n=250; x=[2.0]*(delay+1); ys=[]
    # First 30 clean observations calibrate the delay monitor; attack thereafter.
    for t in range(n):
        d=0 if t<30 else delay
        g=x[-1-d]
        x.append(x[-1]-lr*g if not controlled else x[-1])
        ys.append(x[-1]+rng.normal(0,.1))
    decisions=detect(ys)
    # Re-run dynamics with detector decisions causally.
    x=[2.0]*(delay+1); lr_now=lr; alarm_step=None
    for t in range(n):
        if controlled and decisions[t][1] and alarm_step is None:
            alarm_step=t+1; lr_now=lr/4
        d=0 if t<30 else delay
        x.append(x[-1]-lr_now*x[-1-d])
    return float(.5*x[-1]**2), alarm_step, max(abs(v) for v in x)

def main():
    rows=[]
    for lr in (.8,1.0,1.2):
        b=[run(lr,2,False,s) for s in range(20)]
        i=[run(lr,2,True,s) for s in range(20)]
        rows.append({'lr':lr,'baseline_final_loss':float(np.median([z[0] for z in b])),
                     'imm_final_loss':float(np.median([z[0] for z in i])),
                     'baseline_max_abs_x':float(np.median([z[2] for z in b])),
                     'imm_max_abs_x':float(np.median([z[2] for z in i])),
                     'imm_alarm_median':float(np.median([z[1] for z in i if z[1] is not None]))})
    print(json.dumps({'quadratic_delayed_feedback':rows,'steps':250,'attack_starts':30,'replicates':20},indent=2))
if __name__=='__main__': main()