"""Ray-matched causal persistence benchmark; no ground-truth policy access."""
from pathlib import Path
import argparse,csv,json,time,platform,sys
ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT))
import numpy as np
from aureole import WorldMemory,freeze,draw,correct,exact_mse
from aureole.core import proposal_from_bound
from aureole.renderer import Scene,VisibilityPrior,receiver_grid,light_grid,unoccluded,physical_table


def cluster_interval(values,seed=992):
    """95% percentile interval over independent scenes, not individual pixels."""
    x=np.asarray(values,float)
    rng=np.random.default_rng(seed)
    means=x[rng.integers(0,len(x),(10000,len(x)))].mean(1)
    return [float(v) for v in np.percentile(means,[2.5,97.5])]


def run(followup=False,output="reproduced_results"):
    config=json.loads((ROOT/("experiments_followup.json" if followup else "experiments.json")).read_text())
    prefix="followup" if followup else "rendering"
    points=receiver_grid(*config["receiver_grid"]);lights=light_grid(config["emitter_grid_side"])
    height,width=config["receiver_grid"]; vh,vw=config["viewport"]
    methods=config["methods"];prior=VisibilityPrior(ROOT/"models/visibility_prior.npz")
    phase_frames=[(phase,i) for phase,n in config["phases"].items() for i in range(n)]
    rows=[];prep=[];snapshots={};start=time.perf_counter();n=config["rays_per_receiver_per_frame"]
    for scene_id in config["scene_ids"]:
        scene=Scene.create(scene_id);changed=scene.changed()
        tic=time.perf_counter();features=scene.features(points[:,None,:],lights[None,:,:]);p_all=prior(features)
        prior_seconds=time.perf_counter()-tic
        bound0=unoccluded(points,lights);bound1=unoccluded(points,lights,True,0.7)
        # Privileged reference is used ONLY after the physical online estimate.
        truth0=physical_table(scene,points,lights,bound0)
        truth1=physical_table(scene,points,lights,bound1)
        truth2=physical_table(changed,points,lights,bound1)
        prep.append({"scene":scene_id,"prior_all_receivers_seconds":prior_seconds,"prior_receiver_emitter_pairs":len(points)*len(lights)})
        for seed in config["replicate_seeds"]:
            memory={method:WorldMemory(len(points),len(lights),f"scene-{scene_id}") for method in methods}
            rngs={method:np.random.default_rng(seed+scene_id*1000) for method in methods}
            for frame,(phase,phase_frame) in enumerate(phase_frames):
                if phase=="revisit" and phase_frame==0:
                    for m in memory.values():m.advance(config["unseen_ticks_before_revisit"])
                    memory["screen_cv"].retain_only(np.array([],dtype=int))
                x0=2+(phase_frame%4);y0=8
                ids=(np.arange(y0,y0+vh)[:,None]*width+np.arange(x0,x0+vw)[None,:]).ravel()
                b=(bound1 if phase in ("relight","hidden_change") else bound0)[ids]
                current_scene=changed if phase=="hidden_change" else scene
                offline_truth=(truth2 if phase=="hidden_change" else truth1 if phase=="relight" else truth0)[ids]
                target=offline_truth.sum(1)
                for method in methods:
                    tic=time.perf_counter();m=memory[method]
                    if method=="screen_cv":m.retain_only(ids)
                    base_prior=np.full((len(ids),len(lights)),0.5) if method=="constant_world_cv" else p_all[ids]
                    if method=="raw_importance":v=np.zeros_like(base_prior)
                    elif method=="neural_cv":v=base_prior
                    else:v=m.predict(ids,base_prior)
                    q=proposal_from_bound(b,v,m.trusted(ids),active=(method in ("world_active_cv","world_guarded_cv")))
                    snapshot=freeze(b*v[...,None],q)
                    j=draw(snapshot,n,rngs[method])
                    observed_visibility=current_scene.visibility(points[ids,None,:],lights[j])
                    physical=b[np.arange(len(ids))[:,None],j]*observed_visibility[...,None]
                    result=correct(snapshot,j,physical)
                    if method=="world_plugin":result=snapshot.integral.copy()
                    conflicts=0
                    if method not in ("raw_importance","neural_cv"):
                        conflicts=m.commit(ids,j,observed_visibility,revise_on_conflict=(method=="world_guarded_cv"))
                    m.advance()
                    runtime=time.perf_counter()-tic
                    # Oracle enumeration is outside the online policy and timing.
                    conditional=(np.mean((snapshot.integral-target)**2,axis=1) if method=="world_plugin"
                                 else exact_mse(offline_truth,snapshot,n))
                    rows.append({"scene":scene_id,"seed":seed,"phase":phase,"phase_frame":phase_frame,
                                 "frame":frame,"method":method,"mse":float(np.mean((result-target)**2)),
                                 "conditional_expected_mse":float(conditional.mean()),
                                 "mean_error":float(np.mean(result-target)),"negative_channel_fraction":float(np.mean(result<0)),
                                 "rays":len(ids)*n,"runtime_seconds":runtime,"memory_bytes":0 if method in ("raw_importance","neural_cv") else m.nbytes,
                                 "observed_conflicting_ray_entries":conflicts,"epoch_after":m.epoch,
                                 "known_fraction_after":float(np.isfinite(m.values[ids]).mean())})
                    if scene_id==config["scene_ids"][0] and seed==config["replicate_seeds"][0] and phase_frame==0 and phase in ("revisit","hidden_change"):
                        snapshots[f"{phase}_{method}"]=result.reshape(vh,vw,3)
                        snapshots[f"{phase}_reference"]=target.reshape(vh,vw,3)
        print(f"completed scene {scene_id}; {len(rows)} frame-method observations",flush=True)
    out=ROOT/output;out.mkdir(parents=True,exist_ok=True)
    with (out/f"{prefix}_raw.csv").open("w",newline="") as handle:
        writer=csv.DictWriter(handle,fieldnames=list(rows[0]));writer.writeheader();writer.writerows(rows)
    summary=[]
    for phase in config["phases"]:
        for method in methods:
            subset=[r for r in rows if r["phase"]==phase and r["method"]==method]
            scene_values=[np.mean([r["conditional_expected_mse"] for r in subset if r["scene"]==scene_id]) for scene_id in config["scene_ids"]]
            summary.append({"phase":phase,"method":method,"expected_mse":float(np.mean(scene_values)),
                            "expected_mse_scene_bootstrap_95":cluster_interval(scene_values),
                            "observed_mse":float(np.mean([r["mse"] for r in subset])),
                            "negative_channel_fraction":float(np.mean([r["negative_channel_fraction"] for r in subset])),
                            "cpu_median_ms":float(np.median([r["runtime_seconds"] for r in subset])*1000),
                            "cpu_p99_ms":float(np.quantile([r["runtime_seconds"] for r in subset],.99)*1000)})
    comparisons=[]
    for phase in config["phases"]:
        pairs=([("world_guarded_cv","world_active_cv"),("world_guarded_cv","world_cv"),("world_guarded_cv","raw_importance"),("world_guarded_cv","screen_cv")]
               if followup else [("world_cv","screen_cv"),("world_cv","raw_importance"),("world_active_cv","world_cv"),("world_cv","constant_world_cv"),("world_cv","world_plugin")])
        for contender,baseline in pairs:
            a=[];b=[]
            for sid in config["scene_ids"]:
                a.append(np.mean([r["conditional_expected_mse"] for r in rows if r["scene"]==sid and r["phase"]==phase and r["method"]==contender]))
                b.append(np.mean([r["conditional_expected_mse"] for r in rows if r["scene"]==sid and r["phase"]==phase and r["method"]==baseline]))
            a=np.array(a);b=np.array(b)
            rng=np.random.default_rng(299);draws=rng.integers(0,len(a),(10000,len(a)))
            ratios=1-a[draws].mean(1)/np.maximum(b[draws].mean(1),1e-30)
            comparisons.append({"phase":phase,"contender":contender,"baseline":baseline,
                                "relative_expected_mse_reduction":float(1-a.mean()/max(b.mean(),1e-30)),
                                "scene_bootstrap_95":[float(v) for v in np.percentile(ratios,[2.5,97.5])]})
    report={"protocol":config,"elapsed_seconds":time.perf_counter()-start,"environment":{"python":platform.python_version(),"numpy":np.__version__,"device":"cpu"},
            "record_count":len(rows),"total_physical_ray_calls":int(sum(r["rays"] for r in rows)),
            "per_persistent_method_memory_bytes":len(points)*len(lights)*8,
            "prior_preparation":prep,"summary":summary,"comparisons":comparisons,
            "limitations":["Only a direct-light floor/three-sphere scene family; no game engine or full path tracing.",
                           "Ray counts match. Compute, bandwidth, and VRAM do not match; report is not an equal-frame-time comparison.",
                           "Clock gap has no observations; it is not 500 fully rendered frames.",
                           "CPU timings exclude shared feature/prior preparation and privileged reference enumeration; not full-renderer frame times.",
                           "Hidden changes deliberately omit invalidation and use the stale prior; controls remain fallible.",
                           "Negative linear-radiance estimates are retained for statistical metrics. Display clipping introduces bias.",
                           "Small scene-cluster intervals characterize this generator only; no cross-game generalization claim."]}
    (out/f"{prefix}_report.json").write_text(json.dumps(report,indent=2)+"\n")
    np.savez_compressed(out/f"{prefix}_example_frames.npz",**snapshots)
    print(json.dumps({"records":len(rows),"elapsed_seconds":report["elapsed_seconds"],"comparisons":comparisons},indent=2))

if __name__=="__main__":
    parser=argparse.ArgumentParser();parser.add_argument("--followup",action="store_true")
    parser.add_argument("--output",default="reproduced_results",help="Keep new runs separate from the recorded release evidence")
    args=parser.parse_args();run(args.followup,args.output)