"""One frozen state, new spatial points, known times and three readouts."""
from pathlib import Path
import argparse,csv,json,time,sys,hashlib
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
import numpy as np
from aureole.renderer import Scene,receiver_grid,light_grid,unoccluded
from aureole.certificates import CertificateMemory,visibility_certificate


def run(output):
    protocol=ROOT/'experiments_queries.json';cfg=json.loads(protocol.read_text())
    out=ROOT/output;out.mkdir(parents=True,exist_ok=True)
    anchors=receiver_grid(*cfg['anchor_grid']);queries=receiver_grid(*cfg['query_grid']);lights=light_grid(cfg['emitter_grid_side'])
    ah,aw=cfg['anchor_grid'];nearest=(np.rint((queries[:,1]+1)*(ah-1)/2).astype(int)*aw
                                   +np.rint((queries[:,0]+1)*(aw-1)/2).astype(int))
    rows=[];images={};warm_count=0;init_seconds=0
    for sid in cfg['scene_ids']:
        base=Scene.create(sid);memory=CertificateMemory(anchors,lights,base.spheres)
        ai=np.repeat(np.arange(len(anchors)),len(lights));aj=np.tile(np.arange(len(lights)),len(anchors))
        tic=time.perf_counter();v,m=visibility_certificate(base,anchors[ai],lights[aj]);memory.commit(ai,aj,v,m)
        init_seconds+=time.perf_counter()-tic;warm_count+=len(ai)
        for tau in cfg['times']:
            g=base.spheres.copy();g[:,0]+=tau*np.array([.035,-.020,.015]);g[:,1]+=tau*np.array([.010,.018,-.012])
            scene=Scene(sid,g)
            # Known affine trajectory; this endpoint difference also bounds its whole prefix.
            movement=float(np.max(np.linalg.norm(g[:,:3]-base.spheres[:,:3],axis=1)))
            tic=time.perf_counter();v,known=memory.lookup(nearest,queries,extra_motion=movement)
            rr,j=np.where(~known)
            fresh,_=visibility_certificate(scene,queries[rr],lights[j]);v[rr,j]=fresh
            lookup_trace_seconds=time.perf_counter()-tic
            # Three linear appearance readouts of identical reconstructed visibility.
            outputs=[];tic=time.perf_counter()
            for appearance in range(cfg['appearance_readouts']):
                b=unoccluded(queries,lights,appearance==1,appearance*.7)
                outputs.append((b*v[...,None]).sum(1))
            readout_seconds=time.perf_counter()-tic
            # Independent quadratic oracle audits current exact visibility after output.
            truth=scene.visibility(queries[:,None,:],lights[None,:,:])
            wrong=int(np.count_nonzero(known & (v!=truth)));error=0
            for a,y in enumerate(outputs):
                b=unoccluded(queries,lights,a==1,a*.7);target=(b*truth[...,None]).sum(1)
                error=max(error,float(np.max(np.abs(y-target))))
                if sid==cfg['scene_ids'][0] and tau in (0.,.5,1.):
                    images[f't{tau}_appearance{a}']=y.reshape(*cfg['query_grid'],3)
            rows.append({'scene':sid,'time':tau,'terms':int(v.size),'certified_terms':int(known.sum()),
                         'fresh_queries':len(rr),'shared_fresh_baseline_queries':int(v.size),
                         'separate_readout_baseline_queries':int(v.size*cfg['appearance_readouts']),
                         'false_certificates':wrong,'max_linear_rgb_error':error,
                         'lookup_trace_seconds':lookup_trace_seconds,'three_readout_seconds':readout_seconds})
        print(f'completed query scene {sid}',flush=True)
    with (out/'queries_raw.csv').open('w',newline='') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
    fresh=sum(r['fresh_queries'] for r in rows);shared=sum(r['shared_fresh_baseline_queries'] for r in rows)
    separate=sum(r['separate_readout_baseline_queries'] for r in rows)
    report={'protocol':cfg,'protocol_sha256':hashlib.sha256(protocol.read_bytes()).hexdigest(),'records':len(rows),
            'initialization_queries':warm_count,'residual_queries':fresh,'total_queries_including_initialization':warm_count+fresh,
            'shared_fresh_baseline_queries':shared,'independent_readout_baseline_queries':separate,
            'amortized_query_reduction_vs_shared':1-(warm_count+fresh)/shared,
            'query_reduction_after_warmup_vs_shared':1-fresh/shared,
            'false_certificates':sum(r['false_certificates'] for r in rows),
            'max_linear_rgb_error':max(r['max_linear_rgb_error'] for r in rows),
            'initialization_seconds':init_seconds,'memory_bytes':memory.nbytes,
            'audit_queries_not_in_policy_budget':shared,
            'scope':'Eight new scenes, five prescribed times, three known appearance readouts. No neural SR/FG claims. Warmup charged. Baseline comparison counts physical queries, not total time.'}
    (out/'queries_report.json').write_text(json.dumps(report,indent=2)+'\n')
    np.savez_compressed(out/'queries_frames.npz',**images)
    print(json.dumps(report,indent=2))


if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('--output',default='queries_reproduced');run(p.parse_args().output)