"""
VLA 策略评估脚本（简化版）

该脚本使用 agent 内置的队列管理来评估训练好的 VLA 策略。
无需单独的评估器类 - agent 处理一切！

使用方法:
    python roboimi/demos/eval_vla_simple.py
    python roboimi/demos/eval_vla_simple.py eval.ckpt_path=checkpoints/vla_model_final.pt
    python roboimi/demos/eval_vla_simple.py eval.ckpt_path=checkpoints/vla_model_best.pt
"""

import sys
import os
import json
import logging
import time
import torch
import numpy as np
import hydra
from pathlib import Path
from typing import Dict
from tqdm import tqdm
from omegaconf import DictConfig, OmegaConf
from hydra.utils import instantiate
from einops import rearrange

from roboimi.envs.double_pos_ctrl_env import make_sim_env
from roboimi.utils.act_ex_utils import sample_transfer_pose

sys.path.append(os.getcwd())

log = logging.getLogger(__name__)

if not OmegaConf.has_resolver("len"):
    OmegaConf.register_new_resolver("len", lambda x: len(x))


def load_checkpoint(
    ckpt_path: str,
    agent_cfg: DictConfig,
    device: str = 'cuda'
) -> torch.nn.Module:
    """
    从检查点加载训练好的 VLA 模型，使用 Hydra agent 配置。

    Args:
        ckpt_path: 检查点文件路径 (.pt)
        agent_cfg: Hydra agent 配置，用于实例化
        device: 加载模型的设备

    Returns:
        加载后的 VLAAgent 模型
    """
    from pathlib import Path as PathLib

    ckpt_path = PathLib(ckpt_path).absolute()
    if not ckpt_path.exists():
        raise FileNotFoundError(f"检查点未找到: {ckpt_path}")

    log.info(f"从 {ckpt_path} 加载检查点")
    checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)
    log.info(f"检查点键值: {checkpoint.keys()}")

    # 加载数据集统计信息用于归一化
    stats = checkpoint.get('dataset_stats', None)

    # 使用数据集统计信息从 Hydra 配置实例化 agent
    log.info("从配置实例化 agent...")
    agent = instantiate(agent_cfg, dataset_stats=stats)

    # 加载模型状态
    agent.load_state_dict(checkpoint['model_state_dict'])
    log.info(f"✅ 模型状态已加载 (步数: {checkpoint.get('step', 'unknown')})")

    if stats is not None:
        log.info(f"✅ 数据集统计信息已加载 (归一化: {stats.get('normalization_type', 'gaussian')})")
    else:
        # 后备方案：尝试从外部 JSON 文件加载（兼容旧检查点）
        stats_path = ckpt_path.parent / 'dataset_stats.json'
        if stats_path.exists():
            with open(stats_path, 'r') as f:
                stats = json.load(f)
            log.info("✅ 数据集统计信息已从外部 JSON 加载（旧版本兼容）")
        else:
            log.warning("⚠️  未找到数据集统计信息。动作将无法反归一化！")

    agent.eval()
    agent.to(device)

    log.info(f"✅ 模型已成功加载到 {device}")
    return agent, stats


def prepare_observation(obs: Dict, camera_names: list) -> Dict:
    """
    将环境观测转换为 agent 格式。

    Args:
        obs: 环境观测字典，包含图像和 qpos
        camera_names: 摄像头名称列表

    Returns:
        agent 格式的观测字典
    """
    import cv2

    # 转换图像: numpy -> tensor, HWC -> CHW
    images = {}
    for cam_name in camera_names:
        img = obs['images'][cam_name]
        # Resize 到 224x224（与训练时一致）
        img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_LINEAR)
        img = rearrange(img, 'h w c -> c h w')
        img = torch.from_numpy(img / 255.0).float()
        images[cam_name] = img

    # 转换 qpos: numpy -> tensor
    qpos = torch.from_numpy(obs['qpos']).float()

    return {'qpos': qpos, 'images': images}


class ActionSmoother:
    """
    动作平滑器（指数移动平均）
    用于平滑执行动作以获得更稳定的控制
    """

    def __init__(self, alpha: float = 0.3):
        """
        Args:
            alpha: 平滑系数 (0-1)，值越大越重视当前动作
        """
        self.alpha = alpha
        self.prev_action = None

    def smooth(self, action: np.ndarray) -> np.ndarray:
        """
        平滑动作

        Args:
            action: 当前动作

        Returns:
            平滑后的动作
        """
        if self.prev_action is None:
            smoothed = action
        else:
            smoothed = self.alpha * action + (1 - self.alpha) * self.prev_action
        self.prev_action = smoothed
        return smoothed

    def reset(self):
        """重置平滑器状态"""
        self.prev_action = None


@hydra.main(version_base=None, config_path="../../vla/conf", config_name="config")
def main(cfg: DictConfig):
    """
    使用 agent 内置队列管理的简化版 VLA 评估

    所有评估参数来自 vla/conf/eval.yaml，合并到 cfg 中。
    命令行覆盖: python eval_vla_simple.py eval.ckpt_path=... eval.num_episodes=5
    """

    # 打印配置
    print("=" * 80)
    print("VLA 评估配置:")
    print("=" * 80)
    print(OmegaConf.to_yaml(cfg))
    print("=" * 80)

    eval_cfg = cfg.eval
    device = eval_cfg.device
    camera_names = list(eval_cfg.camera_names)

    # =========================================================================
    # 加载模型
    # =========================================================================
    log.info(f"🚀 从 {eval_cfg.ckpt_path} 加载模型...")
    agent, dataset_stats = load_checkpoint(
        ckpt_path=eval_cfg.ckpt_path,
        agent_cfg=cfg.agent,
        device=device
    )

    # 重置 agent 的队列
    agent.reset()

    # 可选：动作平滑器
    smoother = ActionSmoother(alpha=eval_cfg.smooth_alpha) if eval_cfg.use_smoothing else None

    # =========================================================================
    # 创建环境
    # =========================================================================
    env = make_sim_env(eval_cfg.task_name)

    # =========================================================================
    # 运行评估回合
    # =========================================================================
    all_stats = []

    for episode_idx in range(eval_cfg.num_episodes):
        print(f"\n{'='*60}")
        print(f"回合 {episode_idx + 1}/{eval_cfg.num_episodes}")
        print(f"{'='*60}\n")

        box_pos = sample_transfer_pose()
        env.reset(box_pos)

        # 为新回合重置 agent 队列
        agent.reset()
        if smoother:
            smoother.reset()

        # 计时统计
        inference_times = []
        total_times = []

        with torch.inference_mode():
            for t in tqdm(range(eval_cfg.max_timesteps), desc=f"回合 {episode_idx + 1}"):
                start_total = time.time()

                # 从环境获取观测
                obs = env._get_image_obs()
                qpos_obs = env._get_qpos_obs()
                obs['qpos'] = qpos_obs['qpos']

                # 准备给 agent 的观测
                observation = prepare_observation(obs, camera_names)

                # 选择动作（agent 内部处理队列管理）
                start_inference = time.time()
                action = agent.select_action(observation)

                if device == 'cuda':
                    torch.cuda.synchronize()
                end_inference = time.time()

                # 转换为 numpy
                action = action.cpu().numpy()

                # 调试：打印当前时间步的动作（由配置控制）
                if eval_cfg.get('verbose_action', False):
                    print(f"\n[Step {t:3d}] 预测动作: {action}")
                    print(f"  - 动作形状: {action.shape}")
                    print(f"  - 动作范围: [{action.min():.4f}, {action.max():.4f}]")
                    print(f"  - 动作均值: {action.mean():.4f}, 标准差: {action.std():.4f}")

                # 可选：平滑动作
                if smoother:
                    action = smoother.smooth(action)

                # 执行动作
                env.step_jnt(action)
                env.render()

                end_total = time.time()

                # 记录计时
                inference_times.append(end_inference - start_inference)
                total_times.append(end_total - start_total)

        # =========================================================================
        # 打印回合统计
        # =========================================================================
        avg_inference_time = np.mean(inference_times)
        avg_total_time = np.mean(total_times)

        stats = {
            'inference_fps': 1.0 / avg_inference_time if avg_inference_time > 0 else 0.0,
            'control_fps': 1.0 / avg_total_time if avg_total_time > 0 else 0.0,
            'avg_inference_time_ms': avg_inference_time * 1000,
            'avg_total_time_ms': avg_total_time * 1000,
            'num_inferences': len([t for t in inference_times if t > 0.001]),  # 统计实际推理次数
            'num_steps': len(total_times)
        }
        all_stats.append(stats)

        print(f"\n回合 {episode_idx + 1} 完成 ({eval_cfg.max_timesteps} 时间步)")
        print(f"  模型推理 FPS: {stats['inference_fps']:.2f} Hz")
        print(f"  控制循环 FPS: {stats['control_fps']:.2f} Hz")
        print(f"  平均推理时间: {stats['avg_inference_time_ms']:.2f} ms")
        print(f"  平均总时间: {stats['avg_total_time_ms']:.2f} ms")
        print(f"  总推理次数: {stats['num_inferences']}")

    # =========================================================================
    # 总体统计
    # =========================================================================
    print(f"\n{'='*60}")
    print("评估完成!")
    print(f"{'='*60}")

    if all_stats:
        avg_inference_fps = np.mean([s['inference_fps'] for s in all_stats])
        avg_control_fps = np.mean([s['control_fps'] for s in all_stats])
        avg_inference_time = np.mean([s['avg_inference_time_ms'] for s in all_stats])
        avg_total_time = np.mean([s['avg_total_time_ms'] for s in all_stats])

        print(f"\n总体统计 ({eval_cfg.num_episodes} 个回合):")
        print(f"  平均模型推理 FPS: {avg_inference_fps:.2f} Hz")
        print(f"  平均控制循环 FPS: {avg_control_fps:.2f} Hz")
        print(f"  平均推理时间: {avg_inference_time:.2f} ms")
        print(f"  平均总时间: {avg_total_time:.2f} ms")
    print()


if __name__ == '__main__':
    main()
