import torch
import torch.nn as nn
import numpy as np
from typing import Dict, Optional, Any
from roboimi.vla.core.interfaces import VLABackbone, VLAProjector, VLAHead
from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
from diffusers.schedulers.scheduling_ddim import DDIMScheduler
from roboimi.vla.models.heads.diffusion import ConditionalUnet1D

class VLAAgent(nn.Module):

    def __init__(
        self,
        vision_backbone,      # 你之前定义的 ResNet 类
        action_dim,           # 机器人动作维度 (例如 7: xyz + rpy + gripper)
        obs_dim,              # 本体感知维度 (例如 关节角度)
        pred_horizon=16,      # 预测未来多少步动作
        obs_horizon=4,        # 使用多少步历史观测
        diffusion_steps=100,
        num_cams=3,          # 视觉输入的摄像头数量
    ):
        super().__init__()
        # Store parameters
        self.action_dim = action_dim
        self.obs_dim = obs_dim
        self.pred_horizon = pred_horizon
        self.obs_horizon = obs_horizon
        self.num_cams = num_cams

        self.vision_encoder = vision_backbone
        single_img_feat_dim = self.vision_encoder.output_dim
        total_vision_dim = single_img_feat_dim * num_cams * obs_horizon
        total_prop_dim = obs_dim * obs_horizon
        self.global_cond_dim = total_vision_dim + total_prop_dim

        self.noise_scheduler = DDPMScheduler(
                    num_train_timesteps=diffusion_steps,
                    beta_schedule='squaredcos_cap_v2', # 机器人任务常用的 schedule
                    clip_sample=True,
                    prediction_type='epsilon' # 预测噪声
                )

        # DDIM scheduler for faster inference
        self.infer_scheduler = DDIMScheduler(
            num_train_timesteps=diffusion_steps,
            beta_schedule='squaredcos_cap_v2',
            clip_sample=True,
            prediction_type='epsilon'
        )

        self.noise_pred_net = ConditionalUnet1D(
            input_dim=action_dim,
            global_cond_dim=self.global_cond_dim
        )

    # ==========================
    # 训练阶段 (Training)
    # ==========================
    def compute_loss(self, batch):
        """
        batch: 包含 images, qpos (proprioception), action
        """
        gt_actions = batch['action'] # Shape: (B, Horizon, Action_Dim)
        B = gt_actions.shape[0]
        images = batch['images']
        proprioception = batch['qpos'].view(B, -1) # (B, obs_horizon * obs_dim)


        # 1. 提取视觉特征
        visual_features = self.vision_encoder(images).view(B, -1) # (B, vision_dim)
        
        # 2. 融合特征 -> 全局条件 (Global Conditioning)
        global_cond = torch.cat([visual_features, proprioception], dim=-1)

        # 3. 采样噪声
        noise = torch.randn_like(gt_actions)
        
        # 4. 随机采样时间步 (Timesteps)
        timesteps = torch.randint(
            0, self.noise_scheduler.config.num_train_timesteps, 
            (B,), device=gt_actions.device
        ).long()

        # 5. 给动作加噪 (Forward Diffusion)
        noisy_actions = self.noise_scheduler.add_noise(
            gt_actions, noise, timesteps
        )

        # 6. 网络预测噪声
        pred_noise = self.noise_pred_net(
            sample=noisy_actions, 
            timestep=timesteps, 
            global_cond=global_cond
        )

        # 7. 计算 Loss (MSE)
        loss = nn.functional.mse_loss(pred_noise, noise)
        return loss

    # ==========================
    # 推理阶段 (Inference)
    # ==========================
    @torch.no_grad()
    def predict_action(self, images, proprioception):
        B = 1 # 假设单次推理

        # 1. 提取当前观测特征 (只做一次)
        visual_features = self.vision_encoder(images).view(B, -1)
        proprioception = proprioception.view(B, -1)
        global_cond = torch.cat([visual_features, proprioception], dim=-1)

        # 2. 初始化纯高斯噪声动作
        # Shape: (B, pred_horizon, action_dim)
        current_actions = torch.randn(
            (B, self.pred_horizon, self.action_dim), device=global_cond.device
        )

        # 3. 逐步去噪循环 (Reverse Diffusion)
        self.infer_scheduler.set_timesteps(10) # DDIM 推理步数
        
        for t in self.infer_scheduler.timesteps:
            model_input = current_actions
            
            # 预测噪声
            noise_pred = self.noise_pred_net(
                sample=model_input, 
                timestep=t, 
                global_cond=global_cond
            )

            # 移除噪声，更新 current_actions
            current_actions = self.infer_scheduler.step(
                noise_pred, t, current_actions
            ).prev_sample

        # 4. 输出最终动作序列（归一化空间，由调用方负责反归一化）
        return current_actions