"""Make publication figures from recorded measurements; never simulated gains."""
from pathlib import Path
import argparse,csv,json
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
ROOT=Path(__file__).resolve().parents[1]
plt.rcParams.update({"font.family":"DejaVu Sans","font.size":10,"axes.spines.top":False,"axes.spines.right":False})
COLORS={"raw_importance":"#788995","screen_cv":"#ca8845","world_cv":"#116e83","world_active_cv":"#9a4966","world_guarded_cv":"#298653","constant_world_cv":"#555e62","neural_cv":"#6d6dc4"}
LABELS={"raw_importance":"Raw importance","screen_cv":"Screen memory","world_cv":"World memory","world_active_cv":"Active world","world_guarded_cv":"Guarded active","constant_world_cv":"Constant prior","neural_cv":"Neural prior"}
FIGURES=ROOT/"figures_reproduced"


def save(fig,name):
    FIGURES.mkdir(parents=True,exist_ok=True)
    for ext in ("png","pdf"):fig.savefig(FIGURES/f"{name}.{ext}",dpi=180,bbox_inches="tight")
    plt.close(fig)


def run(results="results",output="figures_reproduced"):
    global FIGURES
    FIGURES=ROOT/output;result_dir=ROOT/results
    initial=json.loads((result_dir/"rendering_report.json").read_text())
    follow=json.loads((result_dir/"followup_report.json").read_text())
    fig,axes=plt.subplots(1,2,figsize=(11,4.3))
    for ax,report,phase,methods,title in [
        (axes[0],initial,"revisit",["raw_importance","screen_cv","world_cv","world_active_cv"],"Initial study: revisit"),
        (axes[1],follow,"hidden_change",["raw_importance","screen_cv","world_cv","world_active_cv","world_guarded_cv"],"Independent scenes: hidden change")]:
        values=[next(s for s in report["summary"] if s["phase"]==phase and s["method"]==m) for m in methods]
        y=np.array([v["expected_mse"] for v in values]);ci=np.array([v["expected_mse_scene_bootstrap_95"] for v in values])
        ax.bar(np.arange(len(methods)),y,color=[COLORS[m] for m in methods],width=.68)
        ax.errorbar(np.arange(len(methods)),y,yerr=np.stack([y-ci[:,0],ci[:,1]-y]),fmt="none",color="#243746",capsize=3)
        ax.set_xticks(np.arange(len(methods)),[LABELS[m] for m in methods],rotation=25,ha="right",fontsize=9)
        ax.set_title(title,fontweight="bold");ax.set_ylabel("Expected linear RGB MSE")
        ax.ticklabel_format(axis="y",style="sci",scilimits=(0,0));ax.grid(axis="y",alpha=.18)
    fig.suptitle("Two shadow rays per receiver per frame; 95% scene-bootstrap intervals",y=1.01,fontsize=12)
    fig.tight_layout();save(fig,"physical_results")
    rows=list(csv.DictReader((result_dir/"followup_raw.csv").open()))
    fig,ax=plt.subplots(figsize=(8,4))
    for method in ["raw_importance","screen_cv","world_cv","world_active_cv","world_guarded_cv"]:
        y=[np.mean([float(r["conditional_expected_mse"]) for r in rows if r["phase"]=="hidden_change" and int(r["phase_frame"])==i and r["method"]==method]) for i in range(8)]
        ax.plot(range(8),y,marker="o",label=LABELS[method],color=COLORS[method],lw=2)
    ax.set(xlabel="Frame after unannounced geometry change",ylabel="Expected linear RGB MSE",title="Trust revocation cannot prevent the first surprise frame")
    ax.set_xticks(range(8));ax.grid(alpha=.2)
    ax.legend(fontsize=9,ncol=3,loc="upper center",bbox_to_anchor=(.5,-.2))
    fig.tight_layout();save(fig,"change_recovery")
    snapshots=np.load(result_dir/"rendering_example_frames.npz",allow_pickle=False)
    fig,axes=plt.subplots(2,4,figsize=(10,5.6))
    for row,phase in enumerate(("revisit","hidden_change")):
        for col,method in enumerate(("reference","raw_importance","screen_cv","world_cv")):
            im=snapshots[f"{phase}_{method}"]
            axes[row,col].imshow(np.clip(im,0,1)**(1/2.2),origin="lower",interpolation="nearest")
            axes[row,col].set_title("Exact finite-light" if method=="reference" else LABELS[method],fontsize=10)
            axes[row,col].set_xticks([]);axes[row,col].set_yticks([])
        axes[row,0].set_ylabel("Revisit" if row==0 else "First change frame",fontsize=10)
    fig.suptitle("One recorded scene/seed; preview clipping is excluded from reported error metrics",fontsize=11)
    fig.tight_layout();save(fig,"physical_frames")
    training=json.loads((ROOT/"results/training.json").read_text())
    fig,ax=plt.subplots(figsize=(6,3.1))
    ax.plot([v["epoch"] for v in training["history"]],[v["validation_bce"] for v in training["history"]],color="#116e83",lw=2)
    ax.axvline(training["selected_epoch"],color="#ca8845",linestyle="--",label=f"Selected epoch {training['selected_epoch']}")
    ax.set(xlabel="Epoch",ylabel="Validation binary cross-entropy",title="Later training overfits this small scene collection")
    ax.legend();ax.grid(alpha=.2);fig.tight_layout();save(fig,"training_curve")

if __name__=="__main__":
    parser=argparse.ArgumentParser();parser.add_argument("--results",default="results")
    parser.add_argument("--output",default="figures_reproduced")
    args=parser.parse_args();run(args.results,args.output)