#!/usr/bin/env python3
"""Harness-independent compute claim: FLOPs per action for WorldDiT.
Their default = 20 flow steps (per the model card). Ours = 2 steps + text cached.
Separates the FLOP reduction (few-step + text-cache) from latency wins (CUDA graphs)."""
import os, sys, json
os.environ.setdefault("LIBERO_CONFIG_PATH", "/root/.libero_lat")
os.environ.setdefault("MUJOCO_GL", "osmesa"); os.environ.setdefault("PYOPENGL_PLATFORM", "osmesa")
sys.path.insert(0, "/root/LIBERO")
import torch
from pathlib import Path
from eval import load_model
from torch.utils.flop_counter import FlopCounterMode

m = load_model(Path("/workspace/worlddit"), "libero_spatial", torch.device("cuda"))
blob = torch.load("/workspace/worlddit/inputs.pt", map_location="cuda")
primary, wrist, state, text = (blob[k].cuda() for k in ("primary", "wrist", "state", "text"))
head = m.unified_action_world_head

def flops_at(steps):
    _hf = head.forward
    head.forward = (lambda ctx, sampling_steps=None, _hf=_hf, s=steps: _hf(ctx, sampling_steps=s))
    fc = FlopCounterMode(display=False)
    with fc:
        m(primary, wrist, state, text)
    head.forward = _hf
    return fc.get_total_flops()

f20 = flops_at(20)          # their default
f2  = flops_at(2)           # our step count
v   = (f20 - f2) / 18.0     # FLOPs per velocity (flow) step
enc = f2 - 2 * v            # everything that runs once (vision + text + perceiver + projectors)

fc = FlopCounterMode(display=False)
with fc:
    m.clip_model.encode_text(text.flatten(0, 1))
f_text = fc.get_total_flops()

opt = enc - f_text + 2 * v  # our steady-state per replan: text cached away, 2 flow steps
G = 1e9
res = {
    "their_default_20step_GFLOPs": round(f20/G, 2),
    "per_flow_step_GFLOPs": round(v/G, 3),
    "encoders_run_once_GFLOPs": round(enc/G, 2),
    "of_which_text_encoder_GFLOPs": round(f_text/G, 3),
    "ours_2step_textcached_GFLOPs": round(opt/G, 2),
    "FLOP_reduction_x": round(f20/opt, 2),
    "note": "FLOP reduction < latency speedup; the rest of latency is CUDA-graphs (systems, no FLOP change)",
}
print("FLOPS " + json.dumps(res, indent=2), flush=True)
Path("/workspace/worlddit/results/flops.json").write_text(json.dumps(res, indent=2))