import json, math
from pathlib import Path
import numpy as np

# Reproducible finite-lattice test of the coherent-to-diffusive claim.
SEED = 1273
J = 1.0
N = 81
center = N // 2

def rhs(rho, gamma):
    # H=-J(sum |x><x+1| + h.c.), and pure dephasing of off-diagonal entries.
    H = np.zeros((N, N), dtype=np.complex128)
    q = np.arange(N - 1)
    H[q, q + 1] = H[q + 1, q] = -J
    return -1j * (H @ rho - rho @ H) - gamma * (rho - np.diag(np.diag(rho)))

def run(gamma, T=10.0, dt=0.01):
    rho = np.zeros((N, N), dtype=np.complex128)
    rho[center, center] = 1.0
    times, msd, corr, hf = [], [], [], []
    steps = int(round(T / dt))
    for s in range(steps + 1):
        if s % 5 == 0:
            p = np.real(np.diag(rho)); p = np.maximum(p, 0)
            x = np.arange(N) - center
            msd.append(float(np.sum(x*x*p) / max(np.sum(p), 1e-12)))
            diag_norm = np.linalg.norm(np.real(np.diag(rho)))
            corr.append(float(np.linalg.norm(rho - np.diag(np.diag(rho))) / (diag_norm + 1e-12)))
            # high spatial frequencies of the population (proxy for oscillations)
            fft = np.abs(np.fft.rfft(p - p.mean())) ** 2
            hf.append(float(fft[len(fft)//2:].sum() / (fft.sum() + 1e-12)))
            times.append(s * dt)
        if s == steps: break
        # RK4 keeps the integration explicit and makes the mechanism auditable.
        k1 = rhs(rho, gamma)
        k2 = rhs(rho + .5*dt*k1, gamma)
        k3 = rhs(rho + .5*dt*k2, gamma)
        k4 = rhs(rho + dt*k3, gamma)
        rho = rho + dt*(k1 + 2*k2 + 2*k3 + k4)/6
        rho = (rho + rho.conj().T) / 2
        rho /= np.trace(rho).real
    return np.array(times), np.array(msd), np.array(corr), np.array(hf)

def fit_power(t, y, lo=3.0):
    z = (t >= lo) & (y > 1e-10)
    return float(np.polyfit(np.log(t[z]), np.log(y[z]), 1)[0])

def fit_slope(t, y, lo=5.0):
    z = t >= lo
    return float(np.polyfit(t[z], y[z], 1)[0])

def feedback(T=10.0, dt=.01, gamma0=.5, target=.10, alpha=.8, gmin=.02, gmax=32.):
    rho = np.zeros((N,N), complex); rho[center,center] = 1
    x=np.arange(N)-center; t=[]; m=[]; gs=[]; rs=[]
    gamma=gamma0; steps=int(T/dt)
    for s in range(steps+1):
        if s % 5 == 0:
            p=np.maximum(np.real(np.diag(rho)),0)
            r=np.linalg.norm(rho-np.diag(np.diag(rho)))/(np.linalg.norm(np.real(np.diag(rho)))+1e-12)
            t.append(s*dt); m.append(np.sum(x*x*p)/max(p.sum(),1e-12)); gs.append(gamma); rs.append(r)
        if s==steps: break
        k1=rhs(rho,gamma); k2=rhs(rho+.5*dt*k1,gamma); k3=rhs(rho+.5*dt*k2,gamma); k4=rhs(rho+dt*k3,gamma)
        rho=(rho+dt*(k1+2*k2+2*k3+k4)/6); rho=(rho+rho.conj().T)/2; rho/=np.trace(rho).real
        # controller is stopped-gradient by construction: r is a scalar observation.
        if s % 5 == 4:
            gamma=float(np.clip(gamma*np.exp(alpha*(r-target)),gmin,gmax))
    return np.array(t),np.array(m),np.array(gs),np.array(rs)

def main():
    np.random.seed(SEED)
    gammas=[0., .5, 1., 2., 4., 8., 16.]
    rows=[]
    for g in gammas:
        t,m,r,h=run(g)
        slope=fit_slope(t,m); beta=fit_power(t,m)
        # D prediction says MSD slope = 2D = 4J^2/gamma.
        pred=None if g==0 else 4*J*J/g
        rows.append(dict(gamma=g, msd_final=float(m[-1]), late_msd_slope=slope,
                         predicted_slope=pred, slope_ratio=None if pred is None else slope/pred,
                         power_exponent=beta, corr_final=float(r[-1]), hf_final=float(h[-1])))
    # Selected-frequency crossover: D k^2 = J |sin k|, gamma=2J k^2/|sin k|.
    kstar=np.pi/2; predicted_cross=2*J*kstar*kstar/abs(np.sin(kstar))
    # empirical crossover: nearest beta=1.5, with gamma~4J prediction separately.
    empirical=min(rows[1:], key=lambda z: abs(z['power_exponent']-1.5))['gamma']
    tc,mm,gg,rr=feedback()
    out={
      'parameters': {'J':J,'N':N,'T':10.,'dt':.01,'seed':SEED},
      'predictions': {
        'diffusion_slope': 'MSD/t = 4 J^2/gamma',
        'crossover_bandwidth': 'gamma approximately 4J',
        'selected_k_crossover': predicted_cross,
        'ballistic_exponent': 2.0, 'diffusive_exponent': 1.0
      },
      'sweep': rows,
      'observed': {'beta_1p5_crossover_gamma': empirical,
                   'feedback_initial_gamma':.5,'feedback_final_gamma':float(gg[-1]),
                   'feedback_final_msd':float(mm[-1]),'feedback_max_gamma':float(gg.max())}
    }
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__=='__main__': main()