#!/usr/bin/env python3
"""GO-1 e2e inference smoke test on Thor.

Validates:
  * model load on Thor sm_110 GPU
  * one forward pass via model(**inputs) with the full dataset-style input
    dict, mirroring evaluate/deploy.py:multi_image_get_item + predict_action

Run:
  export TMPDIR=/data/agi/tmp
  /data/agi/venvs/go1_torch/bin/python /data/agi/check_model_load.py
"""
import os
import sys
import time
import traceback

import go1_env  # 路径来自 go1_env(GO1_* 环境变量, 仓库相对默认)

# Stub `decord` — no aarch64 wheel and only used for video reading in training,
# but dataset.py imports it at module level. Inference doesn't need it.
import types as _types
if "decord" not in sys.modules:
    _fake = _types.ModuleType("decord")
    _fake.VideoReader = type("VideoReader", (), {})
    _fake.cpu = lambda *a, **k: None
    _fake.bridge = _types.ModuleType("decord.bridge")
    _fake.bridge.set_bridge = lambda *a, **k: None
    sys.modules["decord"] = _fake
    sys.modules["decord.bridge"] = _fake.bridge

import torch
import numpy as np
from PIL import Image

from go1_env import CKPT

# ── Header ────────────────────────────────────────────────────────────────
print("=" * 60)
print(f"torch         : {torch.__version__}")
print(f"cuda available: {torch.cuda.is_available()}")
if not torch.cuda.is_available():
    print("✗ CUDA unavailable")
    raise SystemExit(1)
print(f"device        : {torch.cuda.get_device_name(0)} (cap {torch.cuda.get_device_capability(0)})")
free_before, total = torch.cuda.mem_get_info()
print(f"gpu mem free  : {free_before / 1e9:.2f} / {total / 1e9:.2f} GB")
print("=" * 60)

# ── Stage 1: import + load model ──────────────────────────────────────────
print("\n[1/3] Load model")
t0 = time.perf_counter()
from go1.internvl.model.go1 import GO1Model, GO1ModelConfig
from go1.internvl.train.constants import IMG_END_TOKEN
from go1.internvl.train.dataset import build_transform, dynamic_preprocess, preprocess_internvl2_5
from transformers import AutoTokenizer
print(f"   imports: {time.perf_counter()-t0:.2f}s")

t0 = time.perf_counter()
config = GO1ModelConfig.from_pretrained(CKPT, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True)
print(f"   config:  {time.perf_counter()-t0:.2f}s  (action_chunk={config.action_chunk_size}, "
      f"img_size={config.force_image_size}, dynamic={config.dynamic_image_size}, "
      f"max_patch={config.max_dynamic_patch}, use_thumbnail={config.use_thumbnail})")

t0 = time.perf_counter()
model = GO1Model.from_pretrained(CKPT, config=config).to(torch.bfloat16).to("cuda").eval()
torch.cuda.synchronize()
print(f"   weights+cuda: {time.perf_counter()-t0:.2f}s")

n_params = sum(p.numel() for p in model.parameters())
gpu_used = (free_before - torch.cuda.mem_get_info()[0]) / 1e9
print(f"   ✓ {n_params/1e9:.2f}B params, GPU mem occupied: {gpu_used:.2f} GB")

tokenizer = AutoTokenizer.from_pretrained(CKPT, trust_remote_code=True, use_fast=False, add_eos_token=False)
img_transform = build_transform(is_train=False, input_size=config.force_image_size, pad2square=config.pad2square)
num_image_token = int((config.force_image_size // config.vision_config.patch_size) ** 2 * (config.downsample_ratio ** 2))
print(f"   tokenizer + transforms ready (num_image_token={num_image_token})")

# ── Stage 2: build dataset-style sample (3 random images + prompt + state) ─
print("\n[2/3] Build input sample")
def make_sample(prompt, images_pil, state, ctrl_freq):
    """Replicates evaluate/deploy.py:multi_image_get_item, no server deps."""
    images, num_tiles = [], []
    num_image = 0
    for img in images_pil:
        num_image += 1
        if config.dynamic_image_size:
            tiles = dynamic_preprocess(img,
                min_num=config.min_dynamic_patch, max_num=config.max_dynamic_patch,
                image_size=config.force_image_size, use_thumbnail=config.use_thumbnail)
        else:
            tiles = [img]
        images += tiles
        num_tiles.append(len(tiles))

    pixel_values = torch.stack([img_transform(im) for im in images])
    num_patches = pixel_values.size(0)
    num_image_tokens = [num_image_token * n for n in num_tiles]

    conversation = [
        {"from": "human", "value": f"{'<image>'*num_image}{prompt}"},
        {"from": "gpt", "value": ""},
    ]
    ret = preprocess_internvl2_5(
        "internvl2_5", [conversation], tokenizer, num_image_tokens,
        num_image=num_image, group_by_length=True,
    )
    position_ids = ret["attention_mask"].long().cumsum(-1) - 1
    position_ids.masked_fill_(ret["attention_mask"] == 0, 1)
    image_end_token_id = tokenizer.convert_tokens_to_ids(IMG_END_TOKEN)
    assert (ret["input_ids"][0] == image_end_token_id).sum() == num_image, "image tokens truncated"

    return dict(
        input_ids=ret["input_ids"][0],
        attention_mask=ret["attention_mask"][0],
        position_ids=position_ids[0],
        pixel_values=pixel_values,
        image_flags=torch.tensor([1] * num_patches, dtype=torch.long),
        # state needs an extra dim — model's state_adaptor wants (B, 1, state_dim)
        # and predict_action unsqueezes once, so sample state should be (1, state_dim).
        state=torch.from_numpy(state.astype(np.float32)).unsqueeze(0),
        ctrl_freqs=torch.from_numpy(np.array([ctrl_freq], dtype=np.float32)),
    )

# 3 random color images (cam_head + cam_hand_left + cam_hand_right)
images = [Image.fromarray((np.random.rand(480, 640, 3) * 255).astype(np.uint8)) for _ in range(3)]
state = np.zeros(22, dtype=np.float32)
prompt = "What action should the robot take to pick up the black block and place into bin?"
sample = make_sample(prompt, images, state, ctrl_freq=30.0)
print(f"   pixel_values  : {tuple(sample['pixel_values'].shape)} {sample['pixel_values'].dtype}")
print(f"   input_ids     : {tuple(sample['input_ids'].shape)}")
print(f"   image_flags   : {tuple(sample['image_flags'].shape)}  sum={int(sample['image_flags'].sum())}")
print(f"   state         : {tuple(sample['state'].shape)}")
print(f"   ctrl_freqs    : {tuple(sample['ctrl_freqs'].shape)}")

# ── Stage 3: forward pass on GPU ──────────────────────────────────────────
print("\n[3/3] Forward pass (model.forward) — measuring inference latency")

device = "cuda"
def to_dev(s):
    """Mirror predict_action's tensor placement + unsqueeze."""
    return dict(
        pixel_values=s["pixel_values"].to(torch.bfloat16).to(device),
        input_ids=s["input_ids"].to(device).unsqueeze(0),
        attention_mask=s["attention_mask"].to(device).unsqueeze(0),
        position_ids=s["position_ids"].to(device).unsqueeze(0),
        image_flags=s["image_flags"].to(device),
        state=s["state"].to(torch.bfloat16).to(device).unsqueeze(0),
        ctrl_freqs=s["ctrl_freqs"].to(torch.bfloat16).to(device).unsqueeze(0),
    )
inputs = to_dev(sample)

# warm-up (first call always slow — cuBLAS/cuDNN heuristics + JIT compile)
print("   warmup (first call always slow due to autotune)...")
t0 = time.perf_counter()
try:
    with torch.no_grad():
        out = model(**inputs)
    torch.cuda.synchronize()
    print(f"   warmup done in {time.perf_counter()-t0:.1f}s")
except Exception as e:
    print(f"   ✗ FAILED: {type(e).__name__}: {e}")
    traceback.print_exc()
    raise SystemExit(1)

# 3 timed runs
timings = []
for i in range(3):
    t0 = time.perf_counter()
    with torch.no_grad():
        out = model(**inputs)
    torch.cuda.synchronize()
    timings.append(time.perf_counter() - t0)

print(f"   steady-state inference: median={np.median(timings):.3f}s  "
      f"min={min(timings):.3f}s  max={max(timings):.3f}s")

# Output shape + sanity
if isinstance(out, torch.Tensor):
    print(f"   output: tensor shape={tuple(out.shape)}")
elif isinstance(out, (tuple, list)):
    print(f"   output: {type(out).__name__} of len {len(out)}")
    for i, x in enumerate(out):
        if isinstance(x, torch.Tensor):
            print(f"     [{i}] tensor shape={tuple(x.shape)} dtype={x.dtype}")
        elif isinstance(x, (tuple, list)):
            print(f"     [{i}] {type(x).__name__} of len {len(x)}")
            for j, y in enumerate(x):
                if isinstance(y, torch.Tensor):
                    print(f"       [{j}] tensor shape={tuple(y.shape)} "
                          f"dtype={y.dtype} min={y.float().min().item():.3f} max={y.float().max().item():.3f}")
elif isinstance(out, dict):
    for k, v in out.items():
        if isinstance(v, torch.Tensor):
            print(f"     {k}: shape={tuple(v.shape)} dtype={v.dtype}")

# Final GPU memory state
free_final = torch.cuda.mem_get_info()[0]
gpu_used = (free_before - free_final) / 1e9
print(f"\n   final GPU mem occupied: {gpu_used:.2f} GB")

print()
print("=" * 60)
print("✓ GO-1 INFERENCE WORKS ON THOR")
print(f"  median latency: {np.median(timings)*1000:.0f} ms")
print(f"  → action chunk = {config.action_chunk_size} → can re-plan at "
      f"{1.0/np.median(timings):.1f} Hz (chunk lasts {config.action_chunk_size/30.0:.1f}s @ 30Hz)")
print("=" * 60)