import json, math, random
import numpy as np
from scipy.optimize import least_squares, minimize
from scipy.signal import residue

SEED=7
np.random.seed(SEED); random.seed(SEED)

def tf_step(t, T, z):
    # Exact inverse Laplace response of (1+z s)/(s prod(1+T_i s)).
    den=np.poly([-1.0/x for x in T])
    num=np.array([z,1.0])
    # append integrator for the step; residue gives y=sum(r exp(p t)).
    p,r,k=residue(num, np.polymul(den,[1.,0.]))
    y=np.zeros_like(t,dtype=float)
    for pi,ri in zip(p,r): y += np.real(ri*np.exp(np.real(pi)*t))
    return y+float(np.real(k[0])) if len(k) else y

def moment_rho(T,z):
    B=float(sum(T)-z)
    C=float(T[0]*T[1]+T[0]*T[2]+T[1]*T[2]-z*sum(T)+z*z)
    return B*B/C

def twostate_response(t,p):
    Tf,Ts,z=np.exp(p)
    den=np.poly([-1/Tf,-1/Ts])
    q,r,k=residue(np.array([z,1.]),np.polymul(den,[1.,0.]))
    y=np.zeros_like(t,dtype=float)
    for qi,ri in zip(q,r): y += np.real(ri*np.exp(np.real(qi)*t))
    return y+(float(np.real(k[0])) if len(k) else 0.)

def fit_window(t,y,init=None):
    if init is None: init=np.log([.3,1.5,.1])
    r=least_squares(lambda p:twostate_response(t,p)-y,init,max_nfev=70)
    Tf,Ts,z=np.exp(r.x); a1=1/Tf+1/Ts; a0=1/(Tf*Ts)
    return float(a1*a1/a0),r.x,float(np.mean(r.fun*r.fun))

def find_case():
    t=np.linspace(.01,10,121); best=None
    for _ in range(180):
        T=np.exp(np.random.uniform(np.log(.08),np.log(4),3)); z=np.random.uniform(-1,1)
        rm=moment_rho(T,z)
        if not 4.35<rm<4.65: continue
        y=tf_step(t,T,z); rw,_,err=fit_window(t,y)
        score=abs(rm-4.5)+abs(rw-3.31)
        if best is None or score<best[0]: best=(score,T,z,rm,rw,err)
    if best is None: raise RuntimeError('no case; increase search')
    return t,best[1],best[2],best[3],best[4]

def reduction_penalty(T,z,lam=.35,margin=.15):
    ts=np.linspace(.01,10,81); rm=moment_rho(T,z); rw,_,_=fit_window(ts,tf_step(ts,T,z))
    return abs(rm-rw)+lam*max(0,margin-abs(rm-4))+lam*max(0,margin-abs(rw-4)),rm,rw

def optimize_model(t,y,T0,z0,regularized):
    x0=np.r_[np.log(T0),z0]
    def obj(x):
        T=np.exp(x[:3]); z=x[3]; task=np.mean((tf_step(t,T,z)-y)**2)
        if not regularized:return task
        return task+.002*reduction_penalty(T,z)[0]
    r=minimize(obj,x0,method='Nelder-Mead',options={'maxiter':70,'xatol':2e-5,'fatol':1e-9})
    T=np.exp(r.x[:3]);z=float(r.x[3]); pen,rm,rw=reduction_penalty(T,z)
    return {'task_mse':float(np.mean((tf_step(t,T,z)-y)**2)),'rho_moment':float(rm),'rho_window':float(rw),'disagreement':float(abs(rm-rw)),'objective':float(r.fun),'T':T.tolist(),'z':z}

def main():
    t,T,z,rm,rw=find_case(); y=tf_step(t,T,z)
    check=[]
    for rho in [2.,3.31,4.,4.5,7.]:
        D=rho-4;check.append({'rho':rho,'Delta':D,'class':'complex' if D<0 else 'real'})
    out={'seed':SEED,'source_system':{'T':T.tolist(),'z':float(z)},'initial_reductions':{'rho_moment':rm,'rho_window':rw,'disagreement':abs(rm-rw),'window_fit_mse':float(np.mean((tf_step(t,T,z)-tf_step(t,T,z))**2))},'discriminant_check':check,'baseline':optimize_model(t,y,T,z,False),'idea':optimize_model(t,y,T,z,True)}
    with open('results.json','w') as f:json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__':main()