"""Online evaluation of the pixel GCBC policy in the visual OGBench env."""
import argparse, os, json, time, statistics
import numpy as np, torch, gymnasium as gym, ogbench
from common import wilson_ci
from train_pix import PixelPolicy

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--env", default="visual-cube-single-v0")
    ap.add_argument("--rollouts", type=int, default=100)
    ap.add_argument("--tasks", type=int, default=5)
    ap.add_argument("--base_seed", type=int, default=20260825)
    ap.add_argument("--out", default=None)
    a = ap.parse_args()
    dev = "cuda"
    ck = torch.load(a.ckpt, map_location=dev, weights_only=False)
    model = PixelPolicy(6, ck["act_dim"], ck["width"]).to(dev)
    model.load_state_dict(ck["model"]); model.eval()
    env = gym.make(a.env)
    def act(ob, goal):
        o = torch.from_numpy(np.asarray(ob)).to(dev).permute(2, 0, 1).float().div_(255.).sub_(0.5)
        g = torch.from_numpy(np.asarray(goal)).to(dev).permute(2, 0, 1).float().div_(255.).sub_(0.5)
        with torch.no_grad():
            return model(torch.cat([o, g], 0)[None])[0].cpu().numpy()
    per_task, all_k, all_n, steps_s = {}, 0, 0, []
    t0 = time.time()
    for task_id in range(1, a.tasks + 1):
        k = 0
        for r in range(a.rollouts):
            ob, info = env.reset(seed=a.base_seed + 1000 * task_id + r,
                                 options=dict(task_id=task_id, render_goal=True))
            goal = info["goal"]; succ = 0; sstep = None
            for t in range(1000):
                ob, rew, term, trunc, info = env.step(np.clip(act(ob, goal), -1, 1))
                if info.get("success"):
                    succ = 1
                    if sstep is None: sstep = t + 1
                if term or trunc: break
            k += succ; all_n += 1
            if succ: steps_s.append(sstep)
        p, lo, hi = wilson_ci(k, a.rollouts)
        per_task[task_id] = dict(k=k, n=a.rollouts, rate=p, ci=[lo, hi])
        all_k += k
        print(f"  task{task_id}: {k}/{a.rollouts} = {p:.3f} 95% CI [{lo:.3f},{hi:.3f}]", flush=True)
    p, lo, hi = wilson_ci(all_k, all_n)
    med = statistics.median(steps_s) if steps_s else None
    res = dict(env=a.env, rollouts_per_task=a.rollouts, overall=dict(k=all_k, n=all_n, rate=p, ci=[lo, hi]),
               per_task=per_task, median_steps_to_success=med, timeout_rate=1 - all_k / all_n,
               eval_seconds=time.time() - t0, ckpt=a.ckpt)
    print("OVERALL: %d/%d = %.3f  95%% Wilson CI [%.3f, %.3f] | median steps=%s | %.0fs"
          % (all_k, all_n, p, lo, hi, med, time.time() - t0))
    out = a.out or os.path.join(os.path.dirname(a.ckpt), "eval.json")
    json.dump(res, open(out, "w"), indent=2); print("saved ->", out)

if __name__ == "__main__":
    main()