# Figure layout revised 2026-09-23 so the title, subtitle and panel labels no
# longer overlap; the data, the seed and every computed number are unchanged.
# Set OUT to choose the output folder (default: the current folder).
import os
import numpy as np, json, matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
rng=np.random.default_rng(108)
NB=108
def H(h1,h27,h4):
    M=np.zeros((NB,NB))
    for step,t in ((1,h1),(27,h27),(4,h4)):
        if t:
            for k in range(NB): M[k,(k+step)%NB]-=t; M[k,(k-step)%NB]-=t
    return M
def packet(center,width,mom):
    k=np.arange(NB); d=np.abs(k-center); d=np.minimum(d,NB-d)
    return np.exp(-d*d/(2*width*width))*np.exp(1j*mom*k)
def norm(p): return p/np.linalg.norm(p)
psi_two=norm(packet(0,3,0)+packet(54,3,0)); psi_a=norm(packet(0,3,0)); psi_b=norm(packet(54,3,0))
configs=[('Thread only',(1,0,0)),('27 squares only',(0,1,0)),('4 rings only',(0,0,1)),('Thread + squares',(1,1,0)),('All three',(1,1,1))]
times=[1,5,20,100,500]; NS=1000
def evolve(M,psi0,ts):
    w,V=np.linalg.eigh(M); c=V.conj().T@psi0
    return [V@(np.exp(-1j*w*t)*c) for t in ts]
results={}; grid_t=np.linspace(0,100,300)
for name,h in configs:
    M=H(*h)
    P=[np.abs(p)**2 for p in evolve(M,psi_two,grid_t)]
    rows=[]
    for t in times:
        pq=np.abs(evolve(M,psi_two,[t])[0])**2
        pa=np.abs(evolve(M,psi_a,[t])[0])**2; pb=np.abs(evolve(M,psi_b,[t])[0])**2
        pmix=(pa+pb)/2
        samples=rng.choice(NB,size=NS,p=pq/pq.sum())
        near=lambda s,c: np.minimum(np.abs(s-c),NB-np.abs(s-c))<=6
        rows.append(dict(t=t,interference=float(np.abs(pq-pmix).sum()/2),spread=float(1/np.sum(pq**2)),
            near_start_A=float(near(samples,0).mean()),near_start_B=float(near(samples,54).mean()),
            elsewhere=float(1-near(samples,0).mean()-near(samples,54).mean()),
            top_bead=int(np.argmax(np.bincount(samples,minlength=NB))+1),samples=samples.tolist(),p=pq.tolist()))
    results[name]=dict(hops=h,rows=rows,heat=np.array(P).tolist())
json.dump({k:{kk:vv for kk,vv in v.items() if kk!='heat'} for k,v in results.items()},open(os.path.join(os.environ.get('OUT','.'),'quantum-experiment-log.json'),'w'))
# ---------- figure ----------
S='#fcfcfb'; T1='#0b0b0b'; T2='#52514e'; cat=['#2a78d6','#eb6834','#1baf7a','#eda100','#e87ba4']
cmap=LinearSegmentedColormap.from_list('blues',['#fcfcfb','#c9dcf3','#2a78d6','#0d3d7a'])
fig=plt.figure(figsize=(16,19),facecolor=S)
gs=fig.add_gridspec(len(configs)+1,3,width_ratios=[2.2,1.4,1.4],hspace=.62,wspace=.25,top=.905,left=.07,right=.97)
fig.text(.07,.985,'Two places at once: where a measurement finds the particle',ha='left',va='top',fontsize=17,color=T1,weight='bold')
fig.text(.07,.962,'One particle starts as two lumps at bead 1 and bead 55. Five hopping rules, five time scales, 1,000 simulated measurements each.\nGeometry of the thread never enters the equations.',ha='left',va='top',fontsize=10.5,color=T2,linespacing=1.5)
for i,(name,h) in enumerate(configs):
    R=results[name]; ax=fig.add_subplot(gs[i,0]); ax.set_facecolor(S)
    ax.imshow(np.array(R['heat']),aspect='auto',origin='lower',cmap=cmap,extent=[0.5,108.5,0,100],vmin=0,vmax=0.12)
    ax.set_title(f'{name}   (hops thread {h[0]}, squares {h[1]}, rings {h[2]})',loc='left',fontsize=11.5,color=T1,weight='bold')
    ax.set_xlabel('bead',color=T2,fontsize=9); ax.set_ylabel('time',color=T2,fontsize=9); ax.tick_params(colors=T2,labelsize=8)
    for sp in ax.spines.values(): sp.set_color('#dcdbd6')
    ax.set_xticks([1,28,55,82,108])

    # histogram at t=20
    row=[r for r in R['rows'] if r['t']==20][0]; ax2=fig.add_subplot(gs[i,1]); ax2.set_facecolor(S)
    counts=np.bincount(np.array(row['samples']),minlength=NB)
    ax2.bar(np.arange(1,NB+1),counts,width=1,color='#2a78d6',linewidth=0)
    ax2.set_title('1,000 measurements at time 20',loc='left',fontsize=10,color=T1)
    ax2.set_xticks([1,28,55,82,108]); ax2.tick_params(colors=T2,labelsize=8); ax2.set_xlabel('bead found',color=T2,fontsize=9)
    for sp in ('top','right'): ax2.spines[sp].set_visible(False)
    for sp in ('left','bottom'): ax2.spines[sp].set_color('#dcdbd6')
    ax2.set_ylim(0,counts.max()*1.3)
    ax2.text(.02,.97,f"near bead 1: {row['near_start_A']*100:.0f}%   near bead 55: {row['near_start_B']*100:.0f}%   elsewhere: {row['elsewhere']*100:.0f}%",transform=ax2.transAxes,va='top',fontsize=8.5,color=T2)
    # interference vs time
    ax3=fig.add_subplot(gs[i,2]); ax3.set_facecolor(S)
    ts=[r['t'] for r in R['rows']]; inter=[r['interference'] for r in R['rows']]; spread=[r['spread'] for r in R['rows']]
    ax3.plot(ts,inter,color='#2a78d6',lw=2,marker='o',ms=6); ax3.set_xscale('log'); ax3.set_ylim(0,0.62); ax3.set_xlim(.55,950)
    ax3.set_title('interference share of the odds (0 = two lumps never talk)',loc='left',fontsize=10,color=T1)
    ax3.set_xlabel('time (log)',color=T2,fontsize=9); ax3.tick_params(colors=T2,labelsize=8)
    for sp in ('top','right'): ax3.spines[sp].set_visible(False)
    for sp in ('left','bottom'): ax3.spines[sp].set_color('#dcdbd6')
    for ix,(t,v,s) in enumerate(zip(ts,inter,spread)):
        # a point lower than its neighbours takes its label underneath, so the rising line does not run through it
        nb=[inter[j] for j in (ix-1,ix+1) if 0<=j<len(inter)]
        below=v<sum(nb)/len(nb) and v>0.08
        ax3.annotate(f'spread {s:.0f}',(t,v),textcoords='offset points',xytext=(0,-11 if below else 9),ha='center',va='top' if below else 'baseline',fontsize=7.5,color=T2,annotation_clip=False)
ax=fig.add_subplot(gs[len(configs),:]); ax.axis('off')
ax.text(0,0.9,'How to read it',fontsize=11.5,color=T1,weight='bold')
ax.text(0,0.66,'Left: each row is a film strip of the odds. Middle: what 1,000 measurements at time 20 actually returned. Right: how much of the odds come from the two lumps interfering,\n'
 'measured as the gap between the quantum odds and the odds you would get if the particle were simply at bead 1 OR bead 55 with a coin flip (a classical mixture). Zero means no interference.\n'
 '"spread" is the number of beads the wave is effectively spread over. Thread shape, weave and the drawn lines change nothing here: the equations use only which beads are joined.',fontsize=9.5,color=T2,va='top')
plt.savefig(os.path.join(os.environ.get('OUT','.'),'quantum-experiment-report.png'),dpi=110,facecolor=S,bbox_inches='tight')
# console summary
for name,R in results.items():
    print(name); 
    for r in R['rows']: print(f"  t={r['t']:>4}  interference={r['interference']:.3f}  spread={r['spread']:5.1f}  nearA={r['near_start_A']:.2f} nearB={r['near_start_B']:.2f} elsewhere={r['elsewhere']:.2f} top={r['top_bead']}")
