"""Online evaluation of the fine-tuned OpenVLA policy at 224x224 in the eval env."""
import argparse, os, json, time, statistics
import numpy as np, torch, gymnasium as gym, ogbench, mujoco
from PIL import Image
from transformers import AutoModelForVision2Seq, AutoProcessor
from peft import PeftModel
from common import wilson_ci
from vla_finetune import ActionHead, prompt_for

MID = "openvla/openvla-7b"

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", required=True)                 # runs_vla/vla_b1000_seed0
    ap.add_argument("--env", default="visual-cube-single-v0")
    ap.add_argument("--rollouts", type=int, default=20)     # per task
    ap.add_argument("--tasks", type=int, default=5)
    ap.add_argument("--base_seed", type=int, default=20260825)
    a = ap.parse_args()
    dev = "cuda:0"
    proc = AutoProcessor.from_pretrained(MID, trust_remote_code=True)
    vla = AutoModelForVision2Seq.from_pretrained(
        MID, attn_implementation="sdpa", torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True, trust_remote_code=True).to(dev)
    vla = PeftModel.from_pretrained(vla, os.path.join(a.run, "lora")).eval()
    ck = torch.load(os.path.join(a.run, "head.pt"), map_location=dev, weights_only=False)
    head = ActionHead(act=5).to(dev).float(); head.load_state_dict(ck["head"]); head.eval()

    env = gym.make(a.env); u = env.unwrapped
    env.reset(seed=0, options=dict(task_id=1, render_goal=False))
    u.model.vis.global_.offwidth = 224; u.model.vis.global_.offheight = 224
    u._renderer = mujoco.Renderer(u.model, height=224, width=224)
    u._render_height = 224; u._render_width = 224

    def act_fn(gx, gy):
        img = Image.fromarray(np.asarray(u.get_pixel_observation()))
        inp = proc(prompt_for(gx, gy), img).to(dev, dtype=torch.bfloat16)
        with torch.no_grad():
            out = vla(input_ids=inp["input_ids"], attention_mask=inp["attention_mask"],
                      pixel_values=inp["pixel_values"], output_hidden_states=True)
            h = out.hidden_states[-1][:, -1, :].float()
            return head(h)[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=False))
            gxyz = np.asarray(u.task_infos[task_id - 1]["goal_xyzs"])[0]
            gx, gy = float(gxyz[0]), float(gxyz[1])      # goal cube world xy
            succ, sstep = 0, None
            for t in range(1000):
                ob, rew, term, trunc, info = env.step(np.clip(act_fn(gx, gy), -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}  [{time.time()-t0:.0f}s]", flush=True)
    p, lo, hi = wilson_ci(all_k, all_n)
    med = statistics.median(steps_s) if steps_s else None
    res = dict(model="openvla-7b-oft-style", 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)
    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))
    json.dump(res, open(os.path.join(a.run, "eval.json"), "w"), indent=2)

if __name__ == "__main__":
    main()