import json,sys,pathlib,hashlib
import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import expm_multiply
from physics import make_instance

def hamiltonian(inst):
    n=inst['n_sites']; basis=np.arange(2**n, dtype=np.int64)
    rows=[]; cols=[]; vals=[]
    for coef, ops in inst['terms']:
        dest=basis.copy(); v=np.full(len(basis),complex(coef))
        for site,p in ops.items():
            mask=1 << (n-1-int(site)); sign=1-2*((basis & mask)!=0)
            if p in ('X','Y'): dest ^= mask
            if p=='Z': v *= sign
            if p=='Y': v *= 1j*sign
        rows.append(dest); cols.append(basis); vals.append(v)
    return sp.coo_matrix((np.concatenate(vals),(np.concatenate(rows),np.concatenate(cols))),shape=(2**n,2**n)).tocsr()

def trace(inst, dense=False, refinement=1):
    h=hamiltonian(inst); n=inst['n_sites']; times=np.array(inst['times'])
    psi=np.zeros(2**n,complex); psi[int(''.join(map(str,inst['initial_state']['bits'])),2)]=1
    if dense:
        e,u=np.linalg.eigh(h.toarray()); states=(np.exp(-1j*np.outer(times,e))*(u.conj().T@psi))@u.T
    else:
        count=(len(times)-1)*refinement+1
        # Batch the states to bound memory at the 20-site validation probe.
        states=expm_multiply(-1j*h,psi,start=0,stop=times[-1],num=count,traceA=-1j*h.diagonal().sum())[::refinement]
    probs=np.abs(states)**2; basis=np.arange(2**n)
    out=np.column_stack([probs@(1-2*((basis>>(n-1-i))&1)) for i in range(n)])
    return out


def main():
 fam=json.loads(pathlib.Path(sys.argv[1]).read_text());cases=[]
 for seed in fam['suite']['seeds']:
  inst=make_instance(fam,fam['exact']['size'],seed)
  cases.append({'seed':seed,'instance':inst,'trace':trace(inst).tolist()})
 probe=make_instance(fam,fam['accuracy_probe']['size'],fam['accuracy_probe']['seed'])
 small=make_instance(fam,[2,3] if isinstance(fam['exact']['size'],list) else 6,fam['suite']['seeds'][0])
 error=float(np.max(np.abs(trace(small)-trace(small,dense=True))))
 delta=float(np.max(np.abs(np.asarray(cases[0]['trace'])-trace(cases[0]['instance'],refinement=2))))
 assert error<1e-9 and delta<1e-9
 result={'version':fam['suite']['version'],'cases':cases,'accuracy_probe':{'instance':probe,'trace':trace(probe).tolist()},'validation':{'small_dense_max_difference':error,'refined_grid_max_difference':delta,'reference_tolerance':1e-9},'provenance':{'generator_sha256':hashlib.sha256(pathlib.Path(__file__).read_bytes()).hexdigest(),'numpy':np.__version__}}
 pathlib.Path(sys.argv[2]).write_text(json.dumps(result,separators=(',',':'))+'\n')
if __name__=='__main__':main()