"""EEML-PT: finite Hamiltonian model; synthetic checks, no experimental data."""
from pathlib import Path
import itertools
import json
import math
import numpy as np

edges=[(0,1),(0,2),(1,2)]
graphs=[]
for bits in itertools.product([0,1],repeat=3):
    A=np.zeros((3,3))
    for b,(i,j) in zip(bits,edges):
        A[i,j]=A[j,i]=b
    if sum(bits)>=2: # exactly the connected simple graphs on three nodes
        graphs.append((bits,A))
assert len(graphs)==4

def full_hamiltonian(Delta,J,g):
    """4 bond states x 8 hard-core boson occupation states."""
    H=np.zeros((32,32))
    for k,(bits,A) in enumerate(graphs):
        for occ in range(8):
            row=8*k+occ
            H[row,row]=Delta*sum(bits)
            for e,(i,j) in enumerate(edges):
                if bits[e] and ((occ>>i)&1)!=((occ>>j)&1):
                    H[row,8*k+(occ^(1<<i)^(1<<j))]=-J
        for l,(other,B) in enumerate(graphs):
            if sum(a!=b for a,b in zip(bits,other))==1:
                # The endpoints of the flipped edge must have a common neighbour.
                e=next(e for e in range(3) if bits[e]!=other[e])
                i,j=edges[e]
                fac=sum(A[i,k2]*A[k2,j] for k2 in range(3))
                assert fac==1
                for occ in range(8):H[8*k+occ,8*l+occ]=-g
    return H

def thermal_loop(H,T):
    eig,vec=np.linalg.eigh(H)
    w=np.exp(-(eig-eig.min())/T)
    rho=(vec*(w/w.sum()))@vec.T
    tri=next(k for k,(bits,A) in enumerate(graphs) if sum(bits)==3)
    return float(np.trace(rho[3*tri:3*tri+3,3*tri:3*tri+3]))

indices=[8*k+occ for k in range(4) for occ in [1,2,4]]
N=np.diag([occ.bit_count() for k in range(4) for occ in range(8)])
H=full_hamiltonian(1.,2.,.03)
hermitian_error=float(np.max(abs(H-H.T)))
number_commutator=float(np.max(abs(N@H-H@N)))
assert hermitian_error==0 and number_commutator==0

# Closed, time-independent dynamics conserves norm and mean energy.
E,V=np.linalg.eigh(H)
psi=np.zeros(32,complex);psi[indices[0]]=1
psi2=V@(np.exp(-1j*E*2.7)*(V.T@psi))
energy_error=float(abs(np.vdot(psi2,H@psi2)-np.vdot(psi,H@psi)))
norm_error=float(abs(np.vdot(psi2,psi2)-1))
assert energy_error<1e-12 and norm_error<1e-12

# Fixed-graph spectra and ground-state crossing, Delta=1.
spectral_errors=[]
for bits,A in graphs:
    expected=([-math.sqrt(2),0,math.sqrt(2)] if sum(bits)==2 else [-2,1,1])
    spectral_errors.append(float(np.max(abs(np.linalg.eigvalsh(-A)-expected))))
assert max(spectral_errors)<1e-12
critical=1/(2-math.sqrt(2))
assert abs((2-math.sqrt(2)*critical)-(3-2*critical))<1e-14

scans=[]
for J in [1.,critical,2.]:
    T=.05
    # Numerically stable partition weights over all four labelled graphs.
    ground_energies=np.array([2-math.sqrt(2)*J,2,2+math.sqrt(2)*J,
                             3-2*J,3+J,3+J])
    shift=ground_energies.min()
    z=np.exp(-(ground_energies-shift)/T)
    analytic=float(z[3:].sum()/(3*z[:3].sum()+z[3:].sum()))
    probabilities=[]
    for g in [0.,.0001,.03]:
        hf=full_hamiltonian(1,J,g)
        h1=hf[np.ix_(indices,indices)]
        probabilities.append(thermal_loop(h1,T))
    assert abs(probabilities[0]-analytic)<1e-12
    scans.append({"J_over_Delta":J,"kBT_over_Delta":T,
                  "P_loop_classical_sector_formula":analytic,
                  "P_loop_g0":probabilities[0],
                  "P_loop_g0001":probabilities[1],
                  "P_loop_g003":probabilities[2]})

# Countercontrol: without mixing, a line stays a line, even above the crossing.
h1=full_hamiltonian(1,2,0)[np.ix_(indices,indices)]
v0=np.zeros(12,complex);v0[0]=1
ev,ve=np.linalg.eigh(h1)
vt=ve@(np.exp(-1j*ev*3.7)*(ve.T@v0))
tri=next(k for k,(bits,A) in enumerate(graphs) if sum(bits)==3)
frozen_loop_prob=float(np.sum(abs(vt[3*tri:3*tri+3])**2))
assert frozen_loop_prob==0

result={
 "status":"exact finite model and numerical checks; no physical measurements",
 "connected_graphs":len(graphs),"Hilbert_dimension":32,"one_excitation_dimension":12,
 "hermiticity_error":hermitian_error,"number_commutator":number_commutator,
 "closed_energy_error":energy_error,"closed_norm_error":norm_error,
 "fixed_graph_spectral_error":max(spectral_errors),
 "critical_J_over_Delta":critical,
 "labelled_graph_thermal_scans":scans,
 "frozen_topology_above_crossing_loop_probability":frozen_loop_prob
}
Path(__file__).with_suffix(".json").write_text(json.dumps(result,indent=2))
print(json.dumps(result,indent=2))