"""CPU study of TurboQuant's MSE-stage principle. No model or GPU benchmark.
Run: OPENBLAS_NUM_THREADS=1 python experiment.py
Dependencies: Python, NumPy, SciPy, Matplotlib. No downloads or network.
"""
import os
os.environ.setdefault('MPLCONFIGDIR','/tmp/engineering-blog-mpl')
import json, csv, platform, sys
from pathlib import Path
import numpy as np
from scipy.special import betainc, gammaln
from scipy.stats import norm
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

ROOT=Path(__file__).resolve().parent if '__file__' in globals() else Path.cwd()
D=128; N=2048; NQ=128; BITS=(2,3,4); SEEDS=(11,22,33,44,55)
COLORS={'Per-vector min-max':'#b8b8b8','Rotated sphere codebook':'#536b7a','Unrotated sphere codebook':'#333333'}

def unit(x):
    return x/np.linalg.norm(x,axis=1,keepdims=True)

def sphere_cdf(x,d):
    x=np.clip(np.asarray(x,dtype=float),-1,1)
    return .5+.5*np.sign(x)*betainc(.5,(d-1)/2,x*x)

def sphere_codebook(d,b):
    """Lloyd updates for the exact marginal density of a uniform unit sphere.
    The interval's first moment has a closed form. No training data is used.
    """
    k=2**b
    c=norm.ppf((np.arange(k)+.5)/k)/np.sqrt(d)
    C=np.exp(gammaln(d/2)-.5*np.log(np.pi)-gammaln((d-1)/2))
    for i in range(10000):
        edges=np.r_[-1,(c[:-1]+c[1:])/2,1]
        masses=np.diff(sphere_cdf(edges,d))
        primitive=-C/(d-1)*np.maximum(0,1-edges*edges)**((d-1)/2)
        new=np.diff(primitive)/masses
        if np.max(np.abs(new-c))<1e-13:
            return new.astype(np.float32),i+1
        c=new
    raise RuntimeError('Lloyd iterations did not converge')

def orthogonal(d,seed):
    q,r=np.linalg.qr(np.random.default_rng(seed).normal(size=(d,d)))
    # Correct QR signs to obtain a Haar-distributed orthogonal matrix.
    q=q*np.where(np.diag(r)>=0,1.,-1.)[None,:]
    return q.astype(np.float32)

def fixed_quant(x,c,rotation=None):
    z=x if rotation is None else x@rotation
    idx=np.searchsorted((c[:-1]+c[1:])/2,z).astype(np.uint8)
    zhat=c[idx]
    return (zhat if rotation is None else zhat@rotation.T),idx

def minmax_quant(x,b):
    lo=x.min(axis=1,keepdims=True)
    scale=(x.max(axis=1,keepdims=True)-lo)/(2**b-1)
    scale=np.where(scale==0,1,scale)
    idx=np.clip(np.rint((x-lo)/scale),0,2**b-1).astype(np.uint8)
    return lo+scale*idx,idx,lo.astype(np.float32),scale.astype(np.float32)

def pack_codes(codes,b):
    bits=((codes.ravel()[:,None]>>np.arange(b))&1).astype(np.uint8)
    return np.packbits(bits.ravel(),bitorder='little')

def unpack_codes(data,b,shape):
    bits=np.unpackbits(data,bitorder='little')[:np.prod(shape)*b].reshape(-1,b)
    return (bits*(1<<np.arange(b))).sum(axis=1).astype(np.uint8).reshape(shape)

def datasets():
    rng=np.random.default_rng(20261106)
    out={}
    for name,scale in [('Isotropic',np.ones(D)),('Uneven axes',np.geomspace(1,.02,D)),('One strong axis',np.r_[20.,np.ones(D-1)])]:
        x=unit(rng.normal(size=(N,D))*scale)
        q=unit(rng.normal(size=(NQ,D))*scale)
        out[name]=(x,q)
    return out

def metrics(x,q,xhat):
    truth=q@x.T
    approx=q@xhat.T
    exact_top=np.argpartition(truth,-10,axis=1)[:,-10:]
    approx_top=np.argpartition(approx,-10,axis=1)[:,-10:]
    recall=np.mean([len(set(a)&set(b))/10 for a,b in zip(exact_top,approx_top)])
    return {'relative_squared_error':float(np.mean(np.sum((x-xhat)**2,axis=1))),
            'score_rmse':float(np.sqrt(np.mean((truth-approx)**2))),
            'recall_at_10':float(recall)}

def run_experiment():
    rows=[]; codebooks={}; data=datasets(); matrices={s:orthogonal(D,s) for s in SEEDS}
    for b in BITS:
        c,it=sphere_codebook(D,b); codebooks[b]=c
        for name,(x,q) in data.items():
            for method in ['Per-vector min-max','Unrotated sphere codebook']:
                xhat=minmax_quant(x,b)[0] if method=='Per-vector min-max' else fixed_quant(x,c)[0]
                rows.append(dict(dataset=name,bits=b,method=method,rotation_seed=None,**metrics(x,q,xhat)))
            for seed,r in matrices.items():
                xhat,_=fixed_quant(x,c,r)
                rows.append(dict(dataset=name,bits=b,method='Rotated sphere codebook',rotation_seed=seed,**metrics(x,q,xhat)))
    return rows,data,codebooks,matrices

def summarize(rows):
    out=[]
    for dataset in ['Isotropic','Uneven axes','One strong axis']:
        for b in BITS:
            for method in COLORS:
                values=[r for r in rows if r['dataset']==dataset and r['bits']==b and r['method']==method]
                s=dict(dataset=dataset,bits=b,method=method,n=len(values))
                for metric in ['relative_squared_error','score_rmse','recall_at_10']:
                    v=[r[metric] for r in values]
                    s[metric]=float(np.mean(v));s[metric+'_min']=min(v);s[metric+'_max']=max(v)
                out.append(s)
    return out

def make_figures(summary,out=ROOT):
    plt.rcParams.update({'font.family':'DejaVu Sans','font.size':11,'axes.spines.top':False,'axes.spines.right':False,'axes.edgecolor':'#555555','axes.labelcolor':'#333333','text.color':'#222222','xtick.color':'#555555','ytick.color':'#555555','axes.titleweight':'normal','axes.titlesize':12,'svg.fonttype':'none','figure.facecolor':'white','axes.facecolor':'white','grid.color':'#dddddd','grid.linewidth':.6})
    names=['Isotropic','Uneven axes','One strong axis'];labels=['Isotropic','Uneven\naxes','One strong\naxis']
    methods=['Per-vector min-max','Rotated sphere codebook'];palette={'Per-vector min-max':'#b8b8b8','Rotated sphere codebook':'#536b7a','Unrotated sphere codebook':'#333333'}
    fig,axs=plt.subplots(1,2,figsize=(9.4,4.7))
    for i,m in enumerate(methods):
        rr=[next(r for r in summary if r['dataset']==name and r['bits']==3 and r['method']==m) for name in names]
        x=np.arange(3)+(i-.5)*.32
        for ax,key in zip(axs,['relative_squared_error','recall_at_10']):
            y=np.array([r[key]*100 for r in rr])
            ax.bar(x,y,width=.29,color=palette[m],label=m,zorder=3)
            if m=='Rotated sphere codebook':
                lo=np.array([r[key+'_min']*100 for r in rr]);hi=np.array([r[key+'_max']*100 for r in rr])
                ax.errorbar(x,y,yerr=[y-lo,hi-y],fmt='none',color='#333333',ecolor='#333333',capsize=3,lw=1,zorder=4)
            ax.set_xticks(range(3),labels);ax.set_axisbelow(True);ax.grid(axis='y');ax.tick_params(axis='x',length=0,pad=8)
    axs[0].set_title('Reconstruction error',loc='left',pad=14);axs[0].set_ylabel('Mean squared vector error (%)');axs[0].set_ylim(0,22);axs[0].set_yticks([0,5,10,15,20])
    axs[1].set_title('Nearest-neighbor recall',loc='left',pad=14);axs[1].set_ylabel('Recall@10 (%)');axs[1].set_ylim(0,100);axs[1].set_yticks([0,20,40,60,80,100])
    handles,labels2=axs[0].get_legend_handles_labels();fig.legend(handles,labels2,loc='upper center',ncol=2,frameon=False,bbox_to_anchor=(.5,1.04),fontsize=10)
    fig.text(.065,.015,'Three code bits · 128 dimensions · synthetic vectors · whiskers: range across five rotations',fontsize=9,color='#666666')
    fig.tight_layout(rect=(0,.06,1,.93),w_pad=3)
    for name in ['summary','cover']:
        fig.savefig(out/(name+'.svg'),bbox_inches='tight');fig.savefig(out/(name+'.png'),dpi=180,bbox_inches='tight')
    plt.close(fig)
    # Supporting plots remain available in the notebook and downloadable source.
    fig,axs=plt.subplots(1,3,figsize=(12,4),sharey=True)
    styles={'Per-vector min-max':('--','s'),'Rotated sphere codebook':('-','o'),'Unrotated sphere codebook':(':','^')}
    for ax,name in zip(axs,names):
        for method in methods+['Unrotated sphere codebook']:
            aa=[r for r in summary if r['dataset']==name and r['method']==method]
            xx=[r['bits'] for r in aa];yy=[r['relative_squared_error'] for r in aa];ls,mk=styles[method]
            ax.plot(xx,yy,color=palette[method],linestyle=ls,marker=mk,label=method,lw=1.6,ms=4)
            if method=='Rotated sphere codebook':ax.fill_between(xx,[r['relative_squared_error_min'] for r in aa],[r['relative_squared_error_max'] for r in aa],color=palette[method],alpha=.1)
        ax.set_title(name,loc='left');ax.set_xticks(BITS);ax.set_xlabel('Code bits per coordinate');ax.grid(axis='y');ax.set_ylim(0,1)
    axs[0].set_ylabel('Mean squared vector error / energy');h,l=axs[0].get_legend_handles_labels();fig.legend(h,l,loc='lower center',ncol=3,frameon=False,fontsize=9)
    fig.tight_layout(rect=(0,.12,1,1));fig.savefig(out/'distortion.svg',bbox_inches='tight');fig.savefig(out/'distortion.png',dpi=180,bbox_inches='tight');plt.close(fig)
    fig,ax=plt.subplots(figsize=(7,4))
    n=np.array([32,128,512,2048,8192,65536,1000000]);packed=3/8*D*n+4*n+4*D*D+4*8
    ax.plot(n,2*D*n/packed,'o-',color='#536b7a',lw=1.6,ms=4,label='Codes + norms + shared state');ax.axhline(16/3,ls='--',color='#777777',lw=1,label='Code-only ratio: 16 / 3');ax.axhline(1,color='#bbbbbb',lw=.8)
    ax.set_xscale('log');ax.set_ylim(0,6);ax.set_xlabel('Stored 128-dimensional vectors');ax.set_ylabel('Size ratio versus FP16');ax.legend(frameon=False,fontsize=9,loc='lower right');ax.grid(axis='y');fig.tight_layout()
    fig.savefig(out/'ranking-storage.svg',bbox_inches='tight');fig.savefig(out/'ranking-storage.png',dpi=180,bbox_inches='tight');plt.close(fig)

def validate(data,codebooks,matrices):
    tests=[]
    def check(label,ok):
        assert ok,label;tests.append(label)
    for b,c in codebooks.items():
        check(f'{b}-bit sorted symmetric codebook',bool(np.all(np.diff(c)>0) and np.max(np.abs(c+c[::-1]))<1e-7))
        _,idx=fixed_quant(data['Isotropic'][0],c,matrices[11]);p=pack_codes(idx,b)
        check(f'{b}-bit lossless code packing',np.array_equal(idx,unpack_codes(p,b,idx.shape)))
        check(f'{b}-bit payload size',p.nbytes==(N*D*b+7)//8)
    for seed,r in matrices.items():check(f'rotation {seed} orthogonality',np.max(np.abs(r.astype(float).T@r-np.eye(D)))<2e-7)
    check('exact sphere CDF endpoints',np.array_equal(sphere_cdf([-1,0,1],D),[0,.5,1]))
    check('all inputs normalized',all(np.max(abs(np.linalg.norm(x,axis=1)-1))<1e-12 for x,q in data.values()))
    # Basic recovery and a zero-error identity control on the metrics.
    m=metrics(*data['Isotropic'],data['Isotropic'][0]);check('identity error zero and recall one',m['relative_squared_error']==0 and m['score_rmse']==0 and m['recall_at_10']==1)
    return tests

def main():
    rows,data,codebooks,matrices=run_experiment();summary=summarize(rows);tests=validate(data,codebooks,matrices)
    (ROOT/'results.json').write_text(json.dumps({'config':{'dimension':D,'database_vectors':N,'queries':NQ,'data_seed':20261106,'rotation_seeds':SEEDS,'bits':BITS},'summary':summary,'rows':rows,'tests':tests,'boundary_gaps':{name: {'median':float(np.median(np.sort(q@x.T,axis=1)[:,-10]-np.sort(q@x.T,axis=1)[:,-11])), 'mean':float(np.mean(np.sort(q@x.T,axis=1)[:,-10]-np.sort(q@x.T,axis=1)[:,-11]))} for name,(x,q) in data.items()}},indent=2))
    with (ROOT/'results.csv').open('w',newline='') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
    x=data['Isotropic'][0];_,idx=fixed_quant(x,codebooks[3],matrices[11])
    payload=pack_codes(idx,3);payload.tofile(ROOT/'example-3bit-codes.bin');np.linalg.norm(x,axis=1).astype(np.float32).tofile(ROOT/'example-norms-f32.bin');matrices[11].tofile(ROOT/'example-rotation-f32.bin');codebooks[3].tofile(ROOT/'example-codebook-f32.bin')
    sizes={p.name:p.stat().st_size for p in ROOT.glob('example-*.bin')}
    (ROOT/'storage.json').write_text(json.dumps({'files_bytes':sizes,'total_bytes':sum(sizes.values()),'fp16_vector_bytes':N*D*2,'fp16_size_ratio':N*D*2/sum(sizes.values()),'note':'Actual packed files. Array reconstruction buffers, Python objects, allocator overhead, indexes, and model state are excluded.'},indent=2))
    make_figures(summary)
    print(json.dumps({'checks_passed':len(tests),'summary_3bit':[s for s in summary if s['bits']==3],'storage_bytes':sizes},indent=2))

if __name__=='__main__':main()