"""
deploy.py — FrozenVLA 本地部署脚本
====================================
加载 Qwen3-VL-2B-Instruct 基座 + 训练好的 checkpoint → 本地推理

用法:
    python deploy.py                          # 默认加载 epoch_3.pt
    python deploy.py --checkpoint epoch_1.pt  # 指定 checkpoint
    python deploy.py --test-image cat.jpg     # 用真实图片测试

环境: Windows + RTX 4060 Laptop 8GB, 使用 sdpa 注意力
"""

import sys
import argparse
from pathlib import Path

# 把 train 目录加入 sys.path 以导入 model 模块
sys.path.insert(0, str(Path(__file__).parent.parent / "train"))

import torch
from PIL import Image
from model import FrozenVLA


# === 路径配置 ===
BASE_MODEL_PATH = r"G:\LingArm\models\pretrained\Qwen3-VL-2B-Instruct"
CHECKPOINT_DIR = r"G:\LingArm\models\train\frozen_preview"
DEFAULT_CHECKPOINT = "epoch_4.pt"


def load_model(checkpoint_name: str = DEFAULT_CHECKPOINT, verbose: bool = True):
    """
    加载完整部署模型: 基座 + checkpoint head

    Returns:
        model (FrozenVLA): 已加载权重、设为 eval 模式、在 GPU 上的模型
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if verbose:
        print("=" * 60)
        print("  FrozenVLA 本地部署加载")
        print("=" * 60)
        print(f"  Base model : {BASE_MODEL_PATH}")
        print(f"  Checkpoint : {CHECKPOINT_DIR}/{checkpoint_name}")
        print(f"  Device     : {device}")
        print(f"  Attention  : sdpa (Windows/4060)")
        print("-" * 60)

    # Step 1: 加载基座 Qwen3-VL-2B-Instruct（冻结）+ 初始化 MLP/Head
    if verbose:
        print("\n[1/3] Loading Qwen3-VL-2B-Instruct base model...")
    model = FrozenVLA(
        llm_name=BASE_MODEL_PATH,
        mlp_hidden_dim=512,
        mlp_depth=2,
        action_dim=7,
        attn_implementation="sdpa",
    )

    # Step 2: 加载训练好的 head checkpoint
    if verbose:
        print(f"\n[2/3] Loading checkpoint: {checkpoint_name}")
    ckpt_path = Path(CHECKPOINT_DIR) / checkpoint_name
    model.load_trainable(ckpt_path)

    # Step 3: 移到 GPU + 统一 dtype（Qwen3-VL 输出 bf16，head 必须同 dtype）
    if verbose:
        print(f"\n[3/3] Moving to {device} + casting head to bfloat16...")
    model = model.to(device)
    # 将 MLP + ActionHead 转为 bf16 以匹配 Qwen3-VL 输出
    model.mlp_projector = model.mlp_projector.to(dtype=torch.bfloat16)
    model.action_head = model.action_head.to(dtype=torch.bfloat16)
    model.translation_scale.data = model.translation_scale.data.to(torch.bfloat16)
    model.rotation_scale.data = model.rotation_scale.data.to(torch.bfloat16)
    model.eval()

    if verbose:
        trainable = model.count_trainable_params()
        frozen = model.count_frozen_params()
        print(f"\n  Trainable params : {trainable:,}")
        print(f"  Frozen params    : {frozen:,}")
        print(f"  Total params     : {trainable + frozen:,}")
        print(f"  VRAM used        : {torch.cuda.memory_allocated()/1024**3:.2f} GB")
        print("=" * 60)
        print("  [OK] Model loaded successfully!")
        print("=" * 60)

    return model


def predict(model: FrozenVLA, images, instructions):
    """
    推理接口

    Args:
        model: 已加载的 FrozenVLA 模型
        images: PIL Image 或 list of PIL Images
        instructions: str 或 list of str（任务描述）

    Returns:
        action: numpy array, shape (N, 7), EEF delta [dx,dy,dz,ax,ay,az,gripper]
    """
    # 统一转为 list
    if isinstance(images, Image.Image):
        images = [images]
    if isinstance(instructions, str):
        instructions = [instructions]

    with torch.no_grad():
        action = model(images, instructions)
    return action.float().cpu().numpy()  # bf16 -> fp32 -> numpy


def basic_test(model: FrozenVLA):
    """
    基础冒烟测试：用两张假图片推理，验证 pipeline 通畅
    """
    print("\n" + "=" * 60)
    print("  Basic Smoke Test")
    print("=" * 60)

    dummy_images = [
        Image.new("RGB", (448, 448), color=(128, 128, 128)),
        Image.new("RGB", (448, 448), color=(64, 200, 64)),
    ]
    instructions = [
        "pick up the red block",
        "place the spoon into the bowl",
    ]

    print(f"  Input: {len(dummy_images)} images, {len(instructions)} instructions")
    action = predict(model, dummy_images, instructions)

    print(f"\n  Action shape : {action.shape}")
    print(f"  Action range : [{action.min():.4f}, {action.max():.4f}]")
    print(f"  Action mean  : {action.mean():.4f}")
    for i in range(len(action)):
        print(f"  Sample {i}: {action[i]}")
    print("\n  [OK] Inference pipeline OK!")


def main():
    parser = argparse.ArgumentParser(description="FrozenVLA Local Deployment")
    parser.add_argument("--checkpoint", default=DEFAULT_CHECKPOINT,
                        help=f"Checkpoint filename (default: {DEFAULT_CHECKPOINT})")
    parser.add_argument("--test-image", type=str, default=None,
                        help="Path to a test image for real inference")
    parser.add_argument("--skip-test", action="store_true",
                        help="Skip the basic smoke test")
    args = parser.parse_args()

    # 加载模型
    model = load_model(args.checkpoint)

    # 冒烟测试
    if not args.skip_test:
        basic_test(model)

    # 真实图片测试（如果提供）
    if args.test_image:
        img_path = Path(args.test_image)
        if not img_path.exists():
            print(f"\n  [ERROR] Image not found: {img_path}")
            sys.exit(1)
        print(f"\n  Testing with real image: {img_path.name}")
        image = Image.open(img_path).convert("RGB")
        action = predict(model, image, "manipulate the object in front of the robot")
        print(f"  Predicted action (EEF delta 7D): {action[0]}")
        print(f"    dx={action[0,0]:.6f}  dy={action[0,1]:.6f}  dz={action[0,2]:.6f}")
        print(f"    ax={action[0,3]:.6f}  ay={action[0,4]:.6f}  az={action[0,5]:.6f}")
        print(f"    gripper={action[0,6]:.6f}")

    print("\n  [OK] Deployment complete!")


if __name__ == "__main__":
    main()