# Figure layout revised 2026-09-23: title, subtitle and panel titles no longer
# overlap, the fourth panel is numbered 4, the summary fits the figure and its
# minima are shown in scientific notation. Data and computed numbers are
# unchanged. Set OUT to choose the output folder (default: the current folder).
import os
import numpy as np, json, math, matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
S='#fcfcfb'; T1='#0b0b0b'; T2='#52514e'; C=['#2a78d6','#eb6834','#1baf7a','#eda100','#e87ba4']
def ring_H(N, hops):  # hops: {step: t}
    M=np.zeros((N,N))
    for step,t in hops.items():
        if t:
            for k in range(N): M[k,(k+step)%N]-=t; M[k,(k-step)%N]-=t
    return M
def evolve_fn(M):
    w,V=np.linalg.eigh(M)
    return lambda psi0,t: V@(np.exp(-1j*w*t)*(V.conj().T@psi0))
log={}
# ---------- Test 1: dark state ----------
N=108; M=ring_H(N,{27:1}); U=evolve_fn(M); ts=np.linspace(0,4,400)
t1={}
for deg in (0,45,90,135,180):
    psi0=np.zeros(N,complex); psi0[0]=1; psi0[54]=np.exp(1j*math.radians(deg)); psi0/=np.linalg.norm(psi0)
    Pe=[abs(U(psi0,t)[27])**2+abs(U(psi0,t)[81])**2 for t in ts]
    t1[deg]=Pe; log[f'test1_phase_{deg}_maxEmpty']=float(max(Pe))
# ---------- Test 2: Chinese remainder ----------
M2=ring_H(N,{27:1,4:1}); U2=evolve_fn(M2); psi0=np.zeros(N,complex); psi0[0]=1
ts2=np.linspace(0,4*math.pi,1600); F=np.array([abs(np.vdot(psi0,U2(psi0,t)))**2 for t in ts2])
n=np.arange(27); F27=np.array([abs(np.mean(np.exp(2j*t*np.cos(2*math.pi*n/27))))**2 for t in ts2]); F4=np.cos(ts2)**4
prod=F4*F27; log['test2_max_abs_diff_fidelity_vs_product']=float(np.max(np.abs(F-prod)))
pred_peaks=[math.pi*m for m in (1,2,3,4)]; log['test2_predicted_revival_times']=pred_peaks
log['test2_fidelity_at_predicted']=[float(np.interp(p,ts2,F)) for p in pred_peaks]
# ---------- Test 3: quantum vs classical walk ----------
M3=ring_H(N,{1:1}); U3=evolve_fn(M3); psi0=np.zeros(N,complex); psi0[0]=1
k=np.arange(N); d=np.minimum(k,N-k)
ts3=np.linspace(0,25,120)
sq=[math.sqrt(np.sum(abs(U3(psi0,t))**2*d**2)) for t in ts3]
L=np.zeros((N,N)); 
for i in range(N): L[i,(i+1)%N]+=1; L[i,(i-1)%N]+=1; L[i,i]-=2
wl,Vl=np.linalg.eigh(L); p0=np.zeros(N); p0[0]=1
sc=[math.sqrt(max(0,np.sum((Vl@(np.exp(wl*t)*(Vl.T@p0)))*d**2))) for t in ts3]
log['test3_quantum_spread_at_t20']=float(np.interp(20,ts3,sq)); log['test3_classical_spread_at_t20']=float(np.interp(20,ts3,sc))
# ---------- Test 6: other malas ----------
def coprime_splits(Nn):
    out=[]
    for p in range(2,int(math.isqrt(Nn))+1):
        if Nn%p==0 and math.gcd(p,Nn//p)==1: out.append((p,Nn//p))
    return out
rows=[]
for Nn in (27,54,108,1008):
    sp=coprime_splits(Nn)
    if not sp: rows.append(dict(N=Nn,split='none',cycle='-',minStart='-',freqs='-')); continue
    for p,q in sp:
        Mq=ring_H(Nn,{q:1}); Uq=evolve_fn(Mq)
        a=0; b=(p//2)*q % Nn   # opposite (or nearest to opposite) corner
        psi0=np.zeros(Nn,complex); psi0[a]=1
        if b!=a: psi0[b]=1
        psi0/=np.linalg.norm(psi0)
        tt=np.linspace(0,12,1200); Ps=np.array([abs(Uq(psi0,t)[a])**2+(abs(Uq(psi0,t)[b])**2 if b!=a else 0) for t in tt])
        # frequencies present in P_start: differences of eigenvalues of the p-cycle with weight
        w,V=np.linalg.eigh(Mq); c=V.conj().T@psi0; sites=[a]+([b] if b!=a else [])
        # P_start(t) = sum over pairs of modes; collect the rhythms |E_n-E_m| that carry weight
        amp={}
        for i in np.where(abs(c)>1e-9)[0]:
            for j in np.where(abs(c)>1e-9)[0]:
                f=round(abs(w[i]-w[j]),4)
                if f==0: continue
                wgt=abs(sum(c[i]*np.conj(c[j])*V[s_,i]*np.conj(V[s_,j]) for s_ in sites))
                amp[f]=amp.get(f,0)+wgt
        tot=sum(amp.values()) or 1; freqs=sum(1 for f,v in amp.items() if v>0.01*tot)
        rows.append(dict(N=Nn,split=f'{p} × {q}',cycle=p,minStart=float(Ps.min()),freqs=freqs,Ps=Ps.tolist(),tt=tt.tolist()))
log['test6']=[{k:v for k,v in r.items() if k not in('Ps','tt')} for r in rows]
json.dump(log,open(os.path.join(os.environ.get('OUT','.'),'quantum-tests-2-log.json'),'w'),indent=1)
# ---------- figure ----------
fig=plt.figure(figsize=(16,15.5),facecolor=S); gs=fig.add_gridspec(3,2,hspace=.72,wspace=.2,top=.87,bottom=.04,left=.07,right=.97,height_ratios=[1,1,.8])
fig.text(.07,.985,'Four tests on the (4, 27) mala network',ha='left',va='top',fontsize=17,color=T1,weight='bold')
fig.text(.07,.962,'Each panel states a prediction made before the run. Solid = the model; dotted = the prediction where one was computed by hand.',ha='left',va='top',fontsize=10.5,color=T2)
ax=fig.add_subplot(gs[0,0]); ax.set_facecolor(S)
for i,deg in enumerate((0,45,90,135,180)): ax.plot(ts,t1[deg],color=C[i],lw=2,label=f'{deg}°')
ax.plot(ts,np.sin(2*ts)**2,':',color=T2,lw=1.2)
ax.set_title('1. Dark state: phase between the two places decides motion',loc='left',fontsize=11.5,color=T1,weight='bold',pad=24)
ax.set_xlabel('time',color=T2); ax.set_ylabel('probability on the empty corners (28, 82)',color=T2,fontsize=9); ax.legend(title='phase of bead 55',fontsize=8,title_fontsize=8,frameon=False,labelcolor=T2,loc='upper center',bbox_to_anchor=(0.5,-0.2),ncol=5)
ax.text(0,1.025,'prediction: 0° gives sin²(2t) (dotted), 90° tops out at ½, 180° never moves',transform=ax.transAxes,va='bottom',fontsize=8.5,color=T2)
ax=fig.add_subplot(gs[0,1]); ax.set_facecolor(S)
ax.plot(ts2,F,color=C[0],lw=2,label='model (squares + rings, no thread)'); ax.plot(ts2,prod,':',color=T2,lw=1.4,label='4-ring × 27-ring, computed separately')
for p in pred_peaks: ax.axvline(p,color='#dcdbd6',lw=1)
ax.set_title('2. Chinese remainder: 108 beads behave as a 4-ring times a 27-ring',loc='left',fontsize=11.5,color=T1,weight='bold',pad=24)
ax.set_xlabel('time  (grey lines: predicted revival times π, 2π, 3π, 4π)',color=T2); ax.set_ylabel('chance of finding it back at bead 1',color=T2,fontsize=9); ax.legend(fontsize=8,frameon=False,labelcolor=T2)
gap=log['test2_max_abs_diff_fidelity_vs_product']
ax.text(0,1.025,'max gap between model and product: under 1e-14, which is rounding error' if gap<1e-14 else f'max gap between model and product: {gap:.1e}',transform=ax.transAxes,va='bottom',fontsize=8.5,color=T2)
ax=fig.add_subplot(gs[1,0]); ax.set_facecolor(S)
ax.plot(ts3,sq,color=C[0],lw=2,label='quantum wave'); ax.plot(ts3,sc,color=C[1],lw=2,label='classical random walker'); ax.plot(ts3,np.sqrt(2*ts3),':',color=T2,lw=1.2)
ax.set_title('3. Quantum walk vs classical walk along the thread',loc='left',fontsize=11.5,color=T1,weight='bold',pad=24)
ax.set_xlabel('time',color=T2); ax.set_ylabel('spread (beads from the start)',color=T2,fontsize=9); ax.legend(fontsize=8,frameon=False,labelcolor=T2)
ax.text(0,1.025,'prediction: quantum grows in a straight line, classical as a square root (dotted √2t)',transform=ax.transAxes,va='bottom',fontsize=8.5,color=T2)
ax=fig.add_subplot(gs[1,1]); ax.set_facecolor(S)
ci=0
for r in rows:
    if r['split']=='none': continue
    ax.plot(r['tt'],r['Ps'],color=C[ci%5],lw=2 if r['N']==108 else 1.4,label=f"{r['N']} beads, split {r['split']} (cycle of {r['cycle']})"); ci+=1
ax.set_title('4. Other bead counts: 27, 54, 108 and 1008',loc='left',fontsize=11.5,color=T1,weight='bold',pad=24)
ax.text(0,1.025,'question: does the starting pair ever empty completely, and through how many rhythms?',transform=ax.transAxes,va='bottom',fontsize=8.5,color=T2)
ax.set_xlabel('time',color=T2); ax.set_ylabel('probability still on the starting pair',color=T2,fontsize=9); ax.legend(fontsize=7.5,frameon=False,labelcolor=T2,loc='upper center',bbox_to_anchor=(0.5,-0.14),ncol=2)
ax=fig.add_subplot(gs[2,:]); ax.axis('off')
lines=['Summary of test 4','','min   = lowest probability ever left on the starting pair over the sampled times; 0 means it empties completely.','freqs = how many rhythms mix in the motion; 1 means a pure two-level flip.','']
for r in rows:
    if r['split']=='none': lines.append(f"  {r['N']:>5} beads   no coprime split: a plain ring, nothing to test")
    else: lines.append(f"  {r['N']:>5} beads   split {r['split']:>9}   cycle {r['cycle']:>3}   min {r['minStart']:.1e}   freqs {r['freqs']}")
lines+=['', 'Reading: 108 splits into 4-bead cycles, and the 4-cycle is the only cycle with perfect state transfer: its starting pair',
        'empties and refills on one clock. 54 (2 x 27) starts in a standing state and never moves. 1008 also empties its starting',
        'pair (the 16-cycle almost exactly, the 7- and 9-cycles nearly), but through many rhythms at irregular times, and its',
        'cycles have no perfect state transfer. 27 has no split at all. So the earlier guess that only 108 empties completely was wrong.']
ax.text(0,1,'\n'.join(lines),fontsize=9.5,color=T1,va='top',family='monospace',linespacing=1.45)
for a in fig.axes:
    for spn in ('top','right'): a.spines[spn].set_visible(False)
    for spn in ('left','bottom'): a.spines[spn].set_color('#dcdbd6')
    a.tick_params(colors=T2,labelsize=8)
plt.savefig(os.path.join(os.environ.get('OUT','.'),'quantum-tests-2-report.png'),dpi=110,facecolor=S,bbox_inches='tight')
print(json.dumps({k:v for k,v in log.items() if k!='test6'},indent=1)); print(json.dumps(log['test6'],indent=1))