from pathlib import Path
import csv,json,math
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
P=Path(__file__).resolve().parent

def quantize(x,n):
    if type(n)!=int or not 1<=n<=16 or not math.isfinite(x) or not 0<=x<=1:raise ValueError('Invalid quantizer input')
    m=2**n;code=min(m-1,math.floor(m*x));value=(code+.5)/m
    return code,value,value-x

def error(m,sigma):
    if type(m)!=int or not 2<=m<=256 or not math.isfinite(sigma) or sigma<0:raise ValueError('Invalid readout input')
    return 0. if sigma==0 else (m-1)/m*math.erfc(1/(m-1)/(2*sigma*math.sqrt(2)))

def cdf(x):return .5*(1+math.erf(x/math.sqrt(2)))
def matrix(m,sigma):
    if sigma==0:return np.eye(m)
    levels=np.linspace(0,1,m);bounds=[-math.inf]+[(levels[i]+levels[i+1])/2 for i in range(m-1)]+[math.inf]
    return np.array([[cdf((bounds[j+1]-v)/sigma)-cdf((bounds[j]-v)/sigma) for j in range(m)] for v in levels])

def save(name,rows):
    with (P/name).open('w',newline='',encoding='utf-8') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
qrows=[]
for n in [2,4,8]:
    for k in range(1001):
        x=k/1000;c,v,e=quantize(x,n);qrows.append(dict(bits=n,input=x,code=c,reconstruction=v,error=e,step=2**-n))
rr=[dict(states=m,sigma=k/1000,spacing=1/(m-1),nominal_bits=math.log2(m),symbol_error=error(m,k/1000)) for m in [2,4,8] for k in range(201)]
dr=[dict(states=4,state=i,mean=i/3,sigma=.12,observation=k/1000,density=math.exp(-.5*((k/1000-i/3)/.12)**2)/(.12*math.sqrt(2*math.pi))) for i in range(4) for k in range(-500,1501,2)]
save('quantizer.csv',qrows);save('readout-error.csv',rr);save('readout-density.csv',dr)
for r in qrows:assert abs(r['error'])<=r['step']/2+1e-15
for n in [2,4,8]:
    assert quantize(0,n)[0]==0 and quantize(1,n)[0]==2**n-1
    for k in range(1,2**n):assert quantize(k/2**n,n)[0]==k
maxdiff=0.
for r in rr:
    mat=matrix(r['states'],r['sigma']);assert np.max(abs(mat.sum(axis=1)-1))<1e-14;assert mat.min()>=0
    independent=1-np.trace(mat)/r['states'];maxdiff=max(maxdiff,abs(independent-r['symbol_error']))
assert maxdiff<1e-14
for m in [2,4,8]:
    vals=[error(m,k/1000) for k in range(201)];assert all(a<=b for a,b in zip(vals,vals[1:]));assert error(m,0)==0
    assert abs(error(m,1e12)-(1-1/m))<1e-12
assert math.isclose(error(2,.2),.5*math.erfc(.5/(.2*math.sqrt(2))),rel_tol=1e-15)
invalid=0
for fn,args in [(quantize,(-.1,2)),(quantize,(1.1,2)),(quantize,(.5,0)),(quantize,(math.nan,2)),(error,(1,.1)),(error,(4,-.1)),(error,(4,math.inf))]:
    try:fn(*args)
    except ValueError:invalid+=1
    else:raise AssertionError(args)
plt.rcParams.update({'font.size':10,'axes.spines.top':False,'axes.spines.right':False,'savefig.dpi':180})
fig,ax=plt.subplots(2,1,figsize=(8,6.5),layout='constrained')
for n,c in zip([2,4,8],['#3769a2','#7957a8','#278c89']):
    a=[r for r in qrows if r['bits']==n];ax[0].plot([r['input'] for r in a],[r['reconstruction'] for r in a],label=f'{n} bits',color=c,lw=1)
    ax[1].plot([r['input'] for r in a],[r['error'] for r in a],label=f'{n} bits',color=c,lw=.8)
ax[0].plot([0,1],[0,1],ls='--',color='gray',label='Exact value');ax[0].set(ylabel='Reconstruction (arbitrary units)',title='Ideal midpoint quantizer: finite codes, bounded error');ax[1].set(xlabel='Input (arbitrary units)',ylabel='Reconstruction − input')
for a in ax:a.legend(ncol=4,fontsize=8);a.grid(alpha=.2)
for ext in ['png','svg','pdf']:fig.savefig(P/f'quantization.{ext}')
plt.close(fig)
fig,ax=plt.subplots(1,2,figsize=(11,4.5),layout='constrained')
for i in range(4):
    a=[r for r in dr if r['state']==i];ax[0].plot([r['observation'] for r in a],[r['density'] for r in a],label=f'State {i}')
for x in [1/6,.5,5/6]:ax[0].axvline(x,ls='--',color='gray',lw=.8)
ax[0].set(xlabel='Noisy readout (arbitrary units)',ylabel='Conditional probability density (1/unit)',title='A  Four states; noise σ = 0.12');ax[0].legend(fontsize=8)
for m,c in zip([2,4,8],['#3769a2','#7957a8','#278c89']):
    a=[r for r in rr if r['states']==m];ax[1].plot([r['sigma'] for r in a],[r['symbol_error'] for r in a],label=f'{m} states',color=c)
ax[1].set(xlabel='Read noise standard deviation (arbitrary units)',ylabel='Probability of wrong state',title='B  Equally likely states across the same range');ax[1].legend();ax[1].grid(alpha=.2)
fig.suptitle('Synthetic Gaussian readout: nearest-state decision, no drift or retention model')
for ext in ['png','svg','pdf']:fig.savefig(P/f'memory-readout.{ext}')
plt.close(fig)
(P/'validation.json').write_text(json.dumps(dict(all_pass=True,quantizer_rows=3003,readout_rows=603,density_rows=4004,total_rows=7610,transition_matrix_max_error=maxdiff,endpoint_and_all_threshold_checks=True,monotonicity=True,zero_noise_and_large_noise_limits=True,invalid_cases_rejected=invalid),indent=2)+'\n')
(P/'study.json').write_text(json.dumps(dict(synthetic=True,quantizer=dict(bits=[2,4,8],input_min=0,input_max=1,step=.001,reconstruction='bin midpoint',boundary='x=1 last bin; internal thresholds upper bin',outside='reject'),readout=dict(states=[2,4,8],sigma_min=0,sigma_max=.2,sigma_step=.001,state_range=[0,1],prior='equal',noise='independent additive zero-mean Gaussian single read',decision='nearest state; outer intervals unbounded',density=dict(states=4,sigma=.12,observation_min=-.5,observation_max=1.5,step=.002)),units='arbitrary normalized signal units',not_modeled=['drift','retention','physical device performance','bit error coding','energy','throughput']),indent=2)+'\n')
print('PASS 7610 rows; transition error',maxdiff)