"""Online evaluation of state GCBC on cube-single (5 tasks) with Wilson CIs."""
import argparse, os, json, time, statistics
import numpy as np, torch, gymnasium as gym, ogbench
from common import wilson_ci
from train_gcbc import MLPPolicy

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--env", default="cube-single-v0")
    ap.add_argument("--rollouts", type=int, default=50)   # per task
    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" if torch.cuda.is_available() else "cpu"
    ck = torch.load(a.ckpt, map_location=dev, weights_only=False)
    mean = torch.tensor(ck["mean"], device=dev); std = torch.tensor(ck["std"], device=dev)
    model = MLPPolicy(ck["obs_dim"] * 2, ck["act_dim"], ck["hidden"], ck["layers"]).to(dev)
    model.load_state_dict(ck["model"]); model.eval()
    env = gym.make(a.env)
    def act(o, g):
        o = (torch.tensor(np.asarray(o), dtype=torch.float32, device=dev) - mean) / std
        g = (torch.tensor(np.asarray(g), dtype=torch.float32, device=dev) - mean) / std
        with torch.no_grad():
            return model(o[None], g[None])[0].cpu().numpy()
    per_task = {}; all_k = 0; all_n = 0; steps_success = []
    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=False))
            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_success.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_success) if steps_success else None
    res = dict(env=a.env, rollouts_per_task=a.rollouts, tasks=a.tasks,
               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-to-success=%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()