import json
from pathlib import Path
import numpy as np
from scipy.linalg import expm, eigvals

OUT = Path('floquet_results.json')
M1 = np.array([[-5.83182011, 6.67253161], [4.98376025, -2.86979823]])
M2 = np.array([[4.15640208, -3.67581620], [5.62849205, -1.08915620]])
I = np.eye(2)

def edge(A):
    return float(np.max(np.real(eigvals(A))))

def gamma(M1, M2, damping, period):
    h = period / 2.0
    P = expm((-damping*I + M2)*h) @ expm((-damping*I + M1)*h)
    return float(np.log(np.max(np.abs(eigvals(P)))) / period)

def avg_edge(d):
    return .5*(edge(M1-d*I) + edge(M2-d*I))

def crossing(xs, ys):
    for x0,x1,y0,y1 in zip(xs[:-1],xs[1:],ys[:-1],ys[1:]):
        if y0*y1 <= 0:
            return float(x0-y0*(x1-x0)/(y1-y0))
    return None

def main():
    e1,e2=edge(M1),edge(M2)
    comm=float(np.linalg.norm(M2@M1-M1@M2))
    periods=[.01,.03,.1,.2,.5,1.,2.,5.,10.,20.,50.,100.]
    # P1: scalar damping shifts the exact Floquet exponent by -d.
    shifts=[]
    for T in periods:
        g0,g2=gamma(M1,M2,0,T),gamma(M1,M2,2,T)
        shifts.append({'period':T,'measured_shift':g0-g2,'predicted_shift':2.,
                       'abs_error':abs(g0-g2-2.)})
    # P2: at long periods the boundary approaches the mean instantaneous edge.
    ds=np.linspace(0,3,1201)
    adiabatic=.5*(e1+e2)
    boundaries=[]
    for T in periods:
        ys=np.array([gamma(M1,M2,d,T) for d in ds])
        boundaries.append({'period':T,'floquet_boundary':crossing(ds,ys),
                           'adiabatic_boundary':adiabatic})
    # P3: noncommuting finite-period correction can stabilize even when the
    # phasewise rightmost eigenvalues and constant-mean system are unstable.
    mean_edge=edge((M1+M2)/2)
    correction=[]
    for T in periods:
        g=gamma(M1,M2,0,T)
        correction.append({'period':T,'frequency':2*np.pi/T,'floquet_gamma':g,
                           'constant_mean_gamma':mean_edge,'correction':g-mean_edge})
    d=1.4; T=1.0
    rescue={'damping':d,'period':T,'phase_edges':[e1-d,e2-d],
            'floquet_gamma':gamma(M1,M2,d,T),'adiabatic_estimate':avg_edge(d),
            'constant_mean_baseline':mean_edge-d}
    rescue['instantaneous_edges_positive']=min(rescue['phase_edges'])>0
    rescue['floquet_stable']=rescue['floquet_gamma']<0
    rescue['baseline_unstable']=rescue['constant_mean_baseline']>0
    # Linearized-training proxy: norm after repeated periods, same mean B and d.
    x0=np.array([1.,-1.]); n=20
    h=T/2
    P=expm((-d*I+M2)*h)@expm((-d*I+M1)*h)
    periodic=np.linalg.matrix_power(P,n)@x0
    baseline=expm((-d*I+(M1+M2)/2)*T*n)@x0
    proxy={'periods':n,'initial_norm':float(np.linalg.norm(x0)),
           'constant_mean_final_norm':float(np.linalg.norm(baseline)),
           'periodic_final_norm':float(np.linalg.norm(periodic))}
    result={'matrices':{'M1':M1.tolist(),'M2':M2.tolist(),'edge_M1':e1,
                         'edge_M2':e2,'commutator_norm':comm},
      'predictions':{'P1':'gamma(d)=gamma(0)-d exactly',
       'P2':'slow-period boundary approaches average instantaneous edge',
       'P3':'noncommuting finite-period forcing can stabilize beyond mean-matrix dynamics'},
      'damping_shift':shifts,'boundaries':boundaries,'mean_matrix_edge':mean_edge,
      'correction':correction,'rescue':rescue,'proxy_comparison':proxy}
    OUT.write_text(json.dumps(result,indent=2)); print(json.dumps(result,indent=2))
if __name__=='__main__': main()