import json
import numpy as np

SX=np.array([[0,1],[1,0]],complex)
SY=np.array([[0,-1j],[1j,0]],complex)
SZ=np.array([[1,0],[0,-1]],complex)
I2=np.eye(2,dtype=complex)

def H(kx,ky,m,q=1.):
    return q*np.sin(kx)*SX + q*np.sin(ky)*SY + (m+np.cos(kx)+np.cos(ky))*SZ

def drift(kx,ky,m,gamma,q=1.):
    return -gamma*I2-1j*H(kx,ky,m,q)

def spectrum(kx,ky,m,gamma,omega=0.8,noise=1.,q=1.):
    A=drift(kx,ky,m,gamma,q)
    G=np.linalg.inv(-1j*omega*I2-A)
    return G@(noise*I2)@G.conj().T

def link(a,b):
    z=np.vdot(a,b)
    return z/(abs(z)+1e-30)

def fukui_chern(m,gamma=.7,omega=.8,n=41,band=0):
    us=np.empty((n,n,2),complex)
    for ix in range(n):
      for iy in range(n):
        w,v=np.linalg.eigh(spectrum(2*np.pi*ix/n,2*np.pi*iy/n,m,gamma,omega))
        us[ix,iy]=v[:,::-1][:,band]
    total=0.
    for ix in range(n):
      for iy in range(n):
        u=us[ix,iy]; ux=us[(ix+1)%n,iy]; uy=us[ix,(iy+1)%n]
        uxy=us[(ix+1)%n,(iy+1)%n]
        plaquette=link(u,ux)*link(ux,uxy)*link(uxy,uy)*link(uy,u)
        total += np.angle(plaquette)
    return total/(2*np.pi)

def min_gap(m,gamma=.7,omega=.8,n=81,noise=1.):
    gaps=[]
    for kx in np.linspace(-np.pi,np.pi,n,endpoint=False):
      for ky in np.linspace(-np.pi,np.pi,n,endpoint=False):
        w=np.linalg.eigvalsh(spectrum(kx,ky,m,gamma,omega,noise))
        gaps.append(w[1]-w[0])
    return float(np.min(gaps))

def max_h2(m,n=81):
    return max(np.linalg.eigvalsh(H(x,y,m)).max()**2 for x in np.linspace(-np.pi,np.pi,n,endpoint=False) for y in np.linspace(-np.pi,np.pi,n,endpoint=False))

def euler_radius(dt,m,gamma,n=41):
    r=0.
    for x in np.linspace(-np.pi,np.pi,n,endpoint=False):
      for y in np.linspace(-np.pi,np.pi,n,endpoint=False):
        ev=np.linalg.eigvals(I2+dt*drift(x,y,m,gamma))
        r=max(r,float(np.max(np.abs(ev))))
    return r

def strip_edge(m,gamma=.18,nx=40,ny=20,q=1.):
    best=[]
    chi=q
    for kx in 2*np.pi*np.arange(nx)/nx:
      M=np.zeros((2*ny,2*ny),complex)
      for y in range(ny):
        M[2*y:2*y+2,2*y:2*y+2]=(m+np.cos(kx))*SZ+chi*np.sin(kx)*SX
        if y+1<ny:
          T=(-1j*chi*SY+SZ)/2
          M[2*y:2*y+2,2*(y+1):2*(y+1)+2]=T
          M[2*(y+1):2*(y+1)+2,2*y:2*y+2]=T.conj().T
      w,v=np.linalg.eigh(M)
      for j,e in enumerate(w):
        if abs(e)<0.35:
          p=np.sum(abs(v[:,j].reshape(ny,2))**2,axis=1)
          edge=(p[0]+p[1]+p[-1]+p[-2])/(p.sum()+1e-15)
          ipr=float(np.sum(p*p))
          best.append((edge,ipr,float(e)))
    return max(best) if best else (0.,0.,999.)

def main():
  np.random.seed(0); gamma=.7; omega=.8
  lam=max_h2(-1.); dtcrit=2*gamma/(gamma*gamma+lam)
  dts=[.8*dtcrit,.99*dtcrit,1.01*dtcrit,1.2*dtcrit]
  stability=[(dt,euler_radius(dt,-1.,gamma)) for dt in dts]
  masses=[-2.5,-1.5,-.5,.5,1.5,2.5]
  phase=[(m,round(fukui_chern(m,gamma,omega,n=31)),min_gap(m,gamma,omega,n=51)) for m in masses]
  near=[(m,round(fukui_chern(m,gamma,omega,n=31),3),min_gap(m,gamma,omega,n=61)) for m in [-.2,-.1,-.05,.05,.1,.2]]
  edge_top=strip_edge(-1.,q=1.); edge_triv=strip_edge(-1.,q=0.)
  chirality=[(q,strip_edge(-1.,nx=32,ny=16,q=q)[:2]) for q in [0.,.1,.25,.5,1.]]
  exact_gap=[(m,min_gap(m,gamma,omega,n=81)) for m in [-2.,0.,2.]]
  out={'predictions':{
    'stability_bound':{'predicted_dt_critical':dtcrit,'observed_dt_radius_1_crossing':stability,'pass': bool(all((r<1)==(dt<dtcrit) for dt,r in stability))},
    'chern_gap_transition':{'predicted_boundaries':[-2,0,2],'phase_sweep':phase,'near_m0':near,'exact_boundary_gaps':exact_gap,'pass': bool(abs(phase[1][1])==1 and abs(phase[-2][1])==1 and min(x[2] for x in near)<max(x[2] for x in phase))},
    'boundary_localization':{'topological_m_minus1_edge_ratio_ipr':edge_top,'deterministic_nonchiral_q0_baseline':edge_triv,'chirality_sweep':chirality,'pass':edge_top[0]>edge_triv[0]*2 and edge_top[1]>edge_triv[1]*1.3}
  },'parameters':{'gamma':gamma,'omega':omega,'grid':31}}
  print(json.dumps(out,indent=2))
if __name__=='__main__': main()