import sys
import os
import logging
import json
import pickle
import hydra
import torch
from tqdm import tqdm
from omegaconf import DictConfig, OmegaConf
from torch.utils.data import DataLoader, random_split
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from pathlib import Path

# 确保正确的导入路径
sys.path.append(os.getcwd())

from hydra.utils import instantiate

log = logging.getLogger(__name__)

# 注册列表长度解析器（用于配置中如 ${len:${data.camera_names}}）
if not OmegaConf.has_resolver("len"):
    OmegaConf.register_new_resolver("len", lambda x: len(x))


def recursive_to_device(data, device):
    """
    递归地将嵌套字典/列表中的张量移动到指定设备。

    Args:
        data: 字典、列表或张量
        device: 目标设备 (例如 'cuda', 'cpu')

    Returns:
        所有张量已移动到指定设备的数据结构
    """
    if isinstance(data, torch.Tensor):
        return data.to(device)
    elif isinstance(data, dict):
        return {k: recursive_to_device(v, device) for k, v in data.items()}
    elif isinstance(data, list):
        return [recursive_to_device(v, device) for v in data]
    return data


def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_type='cosine', min_lr=0):
    """
    创建带预热的学习率调度器。

    Args:
        optimizer: PyTorch 优化器
        warmup_steps: 预热步数
        max_steps: 总训练步数
        scheduler_type: 预热后的调度器类型 ('cosine' 或 'constant')
        min_lr: 最小学习率（用于余弦衰减）

    Returns:
        LambdaLR 调度器
    """
    import math
    # 在 LambdaLR 修改前捕获初始学习率
    base_lr = optimizer.param_groups[0]['lr']
    min_lr_ratio = min_lr / base_lr if base_lr > 0 else 0.0

    def lr_lambda(step):
        # 预热阶段：从 0 线性增加到 1
        if step < warmup_steps:
            return float(step) / float(max(1, warmup_steps))

        # 预热后阶段
        if scheduler_type == 'cosine':
            # 从 1 到 min_lr_ratio 的余弦退火
            progress = float(step - warmup_steps) / float(max(1, max_steps - warmup_steps))
            cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
            return max(min_lr_ratio, cosine_decay)
        else:
            # 恒定学习率
            return 1.0

    return LambdaLR(optimizer, lr_lambda)


@hydra.main(version_base=None, config_path="../../vla/conf", config_name="config")
def main(cfg: DictConfig):
    """
    VLA 训练脚本（ResNet 骨干网络 + Diffusion 策略）

    该脚本功能：
    1. 从 HDF5 文件加载数据集
    2. 实例化带 ResNet 视觉编码器的 VLAAgent
    3. 训练基于扩散的动作预测模型
    4. 定期保存检查点
    """

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

    log.info(f"🚀 开始 VLA 训练 (设备: {cfg.train.device})")

    # 创建检查点目录
    checkpoint_dir = Path("checkpoints")
    checkpoint_dir.mkdir(exist_ok=True)

    # =========================================================================
    # 1. 实例化数据集与 DataLoader
    # =========================================================================
    log.info("📦 加载数据集...")
    try:
        dataset = instantiate(cfg.data)
        log.info(f"✅ 数据集加载成功。总样本数: {len(dataset)}")
    except Exception as e:
        log.error(f"❌ 数据集加载失败: {e}")
        raise

    # 训练/验证集划分
    val_split = float(cfg.train.get('val_split', 0.1))
    seed = int(cfg.train.get('seed', 42))
    val_size = int(len(dataset) * val_split)
    train_size = len(dataset) - val_size
    if val_size > 0:
        train_dataset, val_dataset = random_split(
            dataset,
            [train_size, val_size],
            generator=torch.Generator().manual_seed(seed)
        )
        log.info(f"✅ 数据集划分: 训练集={train_size}, 验证集={val_size} (验证比例={val_split})")
    else:
        train_dataset, val_dataset = dataset, None
        log.info("✅ 数据集划分: 全部用于训练, 验证集=0 (验证比例=0)")

    train_loader = DataLoader(
        train_dataset,
        batch_size=cfg.train.batch_size,
        shuffle=True,
        num_workers=cfg.train.num_workers,
        pin_memory=(cfg.train.device != "cpu"),
        drop_last=True  # 丢弃不完整批次以稳定训练
    )

    val_loader = None
    if val_dataset is not None:
        val_loader = DataLoader(
            val_dataset,
            batch_size=cfg.train.batch_size,
            shuffle=False,
            num_workers=cfg.train.num_workers,
            pin_memory=(cfg.train.device != "cpu"),
            drop_last=False
        )

    log.info(f"✅ 训练加载器每轮批次数: {len(train_loader)}")
    if val_loader is not None:
        log.info(f"✅ 验证加载器每轮批次数: {len(val_loader)}")

    # =========================================================================
    # 2. 加载数据集统计信息（将传递给 agent）
    # =========================================================================
    log.info("💾 加载数据集统计信息...")
    dataset_stats = None
    try:
        dataset_dir = cfg.data.get('dataset_dir', 'roboimi/demos/dataset/sim_transfer')
        stats_path = Path(dataset_dir) / 'dataset_stats.pkl'

        if stats_path.exists():
            with open(stats_path, 'rb') as f:
                stats = pickle.load(f)

            # 扁平化stats字典（嵌套结构→扁平结构）以匹配NormalizationModule的期望格式
            dataset_stats = {
                'action_mean': stats['action_mean'].tolist(),
                'action_std': stats['action_std'].tolist(),
                'action_min': stats['action_min'].tolist(),
                'action_max': stats['action_max'].tolist(),
                'qpos_mean': stats['qpos_mean'].tolist(),
                'qpos_std': stats['qpos_std'].tolist(),
                'qpos_min': stats['qpos_min'].tolist(),
                'qpos_max': stats['qpos_max'].tolist(),
            }
            log.info(f"✅ 数据集统计信息加载完成 (归一化: {cfg.agent.normalization_type})")
        else:
            log.warning(f"⚠️  统计文件未找到: {stats_path}")
            log.warning("⚠️  推理时动作将无法反归一化！")

    except Exception as e:
        log.warning(f"⚠️  统计信息加载失败: {e}")
        log.warning("⚠️  训练将继续，但推理可能无法正常工作")

    # =========================================================================
    # 3. 实例化 VLA Agent
    # =========================================================================
    log.info("🤖 初始化 VLA Agent...")
    try:
        # 将 dataset_stats 和 normalization_type 传递给 agent
        agent = instantiate(cfg.agent, dataset_stats=dataset_stats)
        agent.to(cfg.train.device)
        agent.train()
        log.info(f"✅ Agent 初始化完成并已移至 {cfg.train.device}")

        # 统计参数量
        total_params = sum(p.numel() for p in agent.parameters())
        trainable_params = sum(p.numel() for p in agent.parameters() if p.requires_grad)
        log.info(f"📊 总参数量: {total_params:,}")
        log.info(f"📊 可训练参数量: {trainable_params:,}")

    except Exception as e:
        log.error(f"❌ Agent 初始化失败: {e}")
        raise

    # =========================================================================
    # 3.1 从预训练 checkpoint 加载权重（微调）
    # =========================================================================
    pretrained_ckpt = cfg.train.get('pretrained_ckpt', None)
    if pretrained_ckpt is not None:
        ckpt_path = Path(pretrained_ckpt)
        if ckpt_path.exists():
            log.info(f"🔄 [Finetune] 从预训练 checkpoint 加载权重: {ckpt_path}")
            try:
                checkpoint = torch.load(ckpt_path, map_location=cfg.train.device)

                # 只加载模型权重（不加载 optimizer、scheduler）
                missing_keys, unexpected_keys = agent.load_state_dict(
                    checkpoint['model_state_dict'],
                    strict=False  # 允许部分加载（结构不完全匹配时）
                )

                log.info(f"✅ [Finetune] 模型权重加载成功")

                if missing_keys:
                    log.warning(f"⚠️  [Finetune] 缺少的键 ({len(missing_keys)} 个): {missing_keys[:5]}...")
                if unexpected_keys:
                    log.warning(f"⚠️  [Finetune] 多余的键 ({len(unexpected_keys)} 个): {unexpected_keys[:5]}...")

                log.info(f"📊 [Finetune] 预训练信息: 步骤={checkpoint.get('step', 'N/A')}, 损失={checkpoint.get('loss', 'N/A')}")
                log.info(f"📈 [Finetune] 使用新的训练配置（lr={cfg.train.lr}, max_steps={cfg.train.max_steps}）")

            except Exception as e:
                log.error(f"❌ [Finetune] 加载 checkpoint 失败: {e}")
                log.warning("⚠️  将从头开始训练")
        else:
            log.error(f"❌ [Finetune] Checkpoint 文件不存在: {ckpt_path}")
            log.warning("⚠️  将从头开始训练")

    # =========================================================================
    # 4. 设置优化器与学习率调度器
    # =========================================================================
    optimizer = AdamW(agent.parameters(), lr=cfg.train.lr, weight_decay=1e-5)
    log.info(f"🔧 优化器: AdamW (学习率={cfg.train.lr})")

    # 设置带预热的学習率调度器
    warmup_steps = int(cfg.train.get('warmup_steps', 500))
    scheduler_type = cfg.train.get('scheduler_type', 'cosine')
    min_lr = float(cfg.train.get('min_lr', 1e-6))

    scheduler = get_lr_schedule_with_warmup(
        optimizer,
        warmup_steps=warmup_steps,
        max_steps=cfg.train.max_steps,
        scheduler_type=scheduler_type,
        min_lr=min_lr
    )
    log.info(f"📈 学习率调度器: {scheduler_type}，{warmup_steps} 步预热 (最小学习率={min_lr})")

    # =========================================================================
    # 5. 训练循环
    # =========================================================================
    log.info("🏋️ 开始训练循环...")

    def build_agent_input(batch_data):
        """构建 agent 输入格式"""
        images = {}
        # SimpleRobotDataset 返回 observation.{cam_name} 格式
        for cam_name in cfg.data.camera_names:
            key = f"observation.{cam_name}"
            if key in batch_data:
                images[cam_name] = batch_data[key]

        return {
            'images': images,
            'qpos': batch_data['observation.state'],  # SimpleRobotDataset 使用 observation.state
            'action': batch_data['action'],
            'action_is_pad': batch_data.get('action_is_pad', None)  # 传递padding mask
        }

    def run_validation():
        """运行验证"""
        if val_loader is None:
            return None
        agent.eval()

        # 设置确定性种子以获得可重现的损失
        # 这确保验证损失在不同步骤之间可比较
        torch.manual_seed(42)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(42)

        total_loss = 0.0
        num_batches = 0
        with torch.no_grad():
            for val_batch in val_loader:
                val_batch = recursive_to_device(val_batch, cfg.train.device)
                val_input = build_agent_input(val_batch)
                val_loss = agent.compute_loss(val_input)
                total_loss += val_loss.item()
                num_batches += 1
        agent.train()
        return total_loss / max(num_batches, 1)

    data_iter = iter(train_loader)
    pbar = tqdm(range(cfg.train.max_steps), desc="训练中", ncols=100)

    best_loss = float('inf')

    for step in pbar:
        try:
            batch = next(data_iter)
        except StopIteration:
            # 轮次结束时重启迭代器
            data_iter = iter(train_loader)
            batch = next(data_iter)

        # =====================================================================
        # 将批次移至设备
        # =====================================================================
        batch = recursive_to_device(batch, cfg.train.device)

        # =====================================================================
        # 准备 agent 输入
        # =====================================================================
        # 数据集返回: {action, qpos, image_<cam_name>, ...}
        # Agent 期望: {images: dict, qpos: tensor, action: tensor}

        # 准备 agent 输入
        agent_input = build_agent_input(batch)

        # =====================================================================
        # 前向传播与损失计算
        # =====================================================================
        try:
            loss = agent.compute_loss(agent_input)
        except Exception as e:
            log.error(f"❌ 步骤 {step} 前向传播失败: {e}")
            raise

        # =====================================================================
        # 反向传播与优化
        # =====================================================================
        optimizer.zero_grad()
        loss.backward()

        # 梯度裁剪以稳定训练
        torch.nn.utils.clip_grad_norm_(agent.parameters(), max_norm=1.0)

        optimizer.step()
        scheduler.step()

        # =====================================================================
        # 日志记录
        # =====================================================================
        if step % cfg.train.log_freq == 0:
            current_lr = optimizer.param_groups[0]['lr']
            pbar.set_postfix({
                "loss": f"{loss.item():.4f}",
                "lr": f"{current_lr:.2e}",
                "best_loss": f"{best_loss:.4f}"
            })
            log.info(f"步骤 {step}/{cfg.train.max_steps} | 损失: {loss.item():.4f} | 学习率: {current_lr:.2e}")

        # =====================================================================
        # 检查点保存与验证
        # =====================================================================
        if step > 0 and step % cfg.train.save_freq == 0:
            # 运行验证
            val_loss = run_validation()
            if val_loss is not None:
                log.info(f"步骤 {step}/{cfg.train.max_steps} | 验证损失: {val_loss:.4f}")

            checkpoint_path = checkpoint_dir / f"vla_model_step_{step}.pt"
            # 使用agent的归一化统计信息（包含normalization_type）
            agent_stats = agent.get_normalization_stats()
            torch.save({
                'step': step,
                'model_state_dict': agent.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
                'loss': loss.item(),
                'val_loss': val_loss,
                'dataset_stats': agent_stats,  # 保存agent的统计信息
                'current_lr': optimizer.param_groups[0]['lr'],
            }, checkpoint_path)
            log.info(f"💾 检查点已保存: {checkpoint_path}")

            # 根据验证损失保存最佳模型
            eval_loss = val_loss if val_loss is not None else loss.item()
            if eval_loss < best_loss:
                best_loss = eval_loss
                best_model_path = checkpoint_dir / "vla_model_best.pt"
                agent_stats = agent.get_normalization_stats()
                torch.save({
                    'step': step,
                    'model_state_dict': agent.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'scheduler_state_dict': scheduler.state_dict(),
                    'loss': loss.item(),
                    'val_loss': val_loss,
                    'dataset_stats': agent_stats,  # 保存agent的统计信息
                    'current_lr': optimizer.param_groups[0]['lr'],
                }, best_model_path)
                log.info(f"🌟 最佳模型已更新: {best_model_path} (验证损失: {best_loss:.4f})")

    # =========================================================================
    # 6. 保存最终模型
    # =========================================================================
    final_model_path = checkpoint_dir / "vla_model_final.pt"
    agent_stats = agent.get_normalization_stats()
    torch.save({
        'step': cfg.train.max_steps,
        'model_state_dict': agent.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        'loss': loss.item(),
        'dataset_stats': agent_stats,  # 保存agent的统计信息
        'current_lr': optimizer.param_groups[0]['lr'],
    }, final_model_path)
    log.info(f"💾 最终模型已保存: {final_model_path}")

    log.info("✅ 训练成功完成!")
    log.info(f"📊 最终损失: {loss.item():.4f}")
    log.info(f"📊 最佳损失: {best_loss:.4f}")


if __name__ == "__main__":
    main()
