"""Delay-Budget Controller MVP and numerical check of its stated DDE math."""
import json, math
import numpy as np
from scipy.optimize import brentq, root


def omega_from_G(G, rates):
    rates = np.asarray(rates, dtype=float)
    dc = float(np.prod(rates))
    if G <= dc:
        return None
    f = lambda w: float(np.prod(rates * rates + w*w)) - G*G
    hi = max(1., math.sqrt(G) + float(np.max(rates)))
    while f(hi) < 0: hi *= 2
    return brentq(f, 0., hi)


def tau_crit(G, rates, q=0):
    w = omega_from_G(G, rates)
    if w is None: return math.inf
    phase = sum(math.atan(w/r) for r in rates)
    return ((2*q+1)*math.pi-phase)/w


def characteristic(z, G, rates, tau):
    return np.prod(z + np.asarray(rates)) + G*np.exp(-z*tau)


def max_root_real_part(G, rates, tau):
    """Find the dominant roots by seeded complex Newton solves."""
    rates=np.asarray(rates,float); scale=max(1.,float(np.max(rates)),G)
    found=[]
    for re in np.linspace(-2*scale, 1.0*scale, 13):
      for im in np.linspace(-8*scale, 8*scale, 25):
        def fun(a):
          z=a[0]+1j*a[1]; h=characteristic(z,G,rates,tau)
          dz=np.prod(z+rates)*sum(1/(z+r) for r in rates)-G*tau*np.exp(-z*tau)
          # complex Newton represented as real residual
          return [h.real,h.imag]
        sol=root(fun,[re,im])
        if sol.success and np.linalg.norm(fun(sol.x))<1e-6:
          z=sol.x[0]+1j*sol.x[1]
          if not any(abs(z-w)<1e-4 for w in found): found.append(z)
    return float(max(z.real for z in found)) if found else float('nan')


def controller(G, rates, tau, threshold=.8):
    rates=np.asarray(rates,float); margin=float(np.prod(rates)); tc=tau_crit(G,rates)
    # This is the literal actionable logic, with an explicit undefined case.
    if G < margin:
        return {'branch':'no-Hopf-from-formula','tau':tau,'rates':rates.tolist(),'tau_crit':tc}
    if tau > threshold*tc:
        return {'branch':'delay','tau':threshold*tc,'rates':rates.tolist(),'tau_crit':tc}
    new=rates.copy(); new[1]*=2.5; new[3]*=2.5
    return {'branch':'filter','tau':tau,'rates':new.tolist(),'tau_crit':tc}


def simulate(G,rates,tau,T=80.,dt=.002,seed=3):
    rng=np.random.default_rng(seed); rates=np.asarray(rates,float)
    muX,kX,muY,kY=rates; n=int(T/dt)+1; z=np.zeros((n,4)); z[0]=rng.normal(0,.02,4)
    lag=max(0,int(round(tau/dt))); coupling=G/(kX*muY*kY)
    for i in range(n-1):
      j=max(0,i-lag); x,p,y,v=z[i]
      d=np.array([-muX*x-coupling*z[j,3], kX*(x-p), muY*(p-y), kY*(y-v)])
      z[i+1]=z[i]+dt*d
      if not np.all(np.isfinite(z[i+1])) or np.max(abs(z[i+1]))>1e8:
        return {'tail_rms':float('inf'),'max':float('inf'),'diverged':True}
    return {'tail_rms':float(np.sqrt(np.mean(z[int(.7*n):]**2))), 'max':float(np.max(abs(z))), 'diverged':False}


def main():
    rates=np.ones(4); G=1.2; tc=tau_crit(G,rates)
    mathcheck={'omega':omega_from_G(G,rates),'tau_crit':tc,
      'magnitude_relative_error':abs(math.sqrt(np.prod(rates*rates+omega_from_G(G,rates)**2))-G)/G,
      'phase_error':abs(omega_from_G(G,rates)*tc+sum(math.atan(omega_from_G(G,rates)/r) for r in rates)-math.pi),
      'margin':float(np.prod(rates))}
    cases=[]
    for tau in [0., .8*tc, 1.2*tc, 2.*tc]:
      c=controller(G,rates,tau)
      cases.append({'G':G,'tau':tau,'controller':c,'dominant_real_root':max_root_real_part(G,rates,tau),
                    'simulation':simulate(G,rates,tau),
                    'controlled_simulation':simulate(G,np.array(c['rates']),c['tau'])})
    # Below-margin case: no positive-frequency crossing predicted, even at large delay.
    Glo=.8; taulo=20.; cases.append({'G':Glo,'tau':taulo,'controller':controller(Glo,rates,taulo),
      'dominant_real_root':max_root_real_part(Glo,rates,taulo),'simulation':simulate(Glo,rates,taulo)})
    out={'math_check':mathcheck,'cases':cases}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()