from roboimi.vla.core.interfaces import VLABackbone
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import numpy as np
from typing import Callable, Optional, Tuple, Union

def _replace_submodules(
    root_module: nn.Module, predicate: Callable[[nn.Module], bool], func: Callable[[nn.Module], nn.Module]
) -> nn.Module:
    """
    Args:
        root_module: 需要替换子模块的根模块
        predicate: 接受一个模块作为参数，如果该模块需要被替换则返回 True。
        func: 接受一个模块作为参数，并返回一个新的模块来替换它。
    Returns:
        子模块已被替换的根模块。
    """
    if predicate(root_module):
        return func(root_module)

    replace_list = [k.split(".") for k, m in root_module.named_modules(remove_duplicate=True) if predicate(m)]
    for *parents, k in replace_list:
        parent_module = root_module
        if len(parents) > 0:
            parent_module = root_module.get_submodule(".".join(parents))
        if isinstance(parent_module, nn.Sequential):
            src_module = parent_module[int(k)]
        else:
            src_module = getattr(parent_module, k)
        tgt_module = func(src_module)
        if isinstance(parent_module, nn.Sequential):
            parent_module[int(k)] = tgt_module
        else:
            setattr(parent_module, k, tgt_module)
    # 验证所有 BN 是否已被替换
    assert not any(predicate(m) for _, m in root_module.named_modules(remove_duplicate=True))
    return root_module

class SpatialSoftmax(nn.Module):
    """
    Finn 等人在 "Deep Spatial Autoencoders for Visuomotor Learning" 中描述的空间软 Argmax 操作
    (https://huggingface.co/papers/1509.06113)。这是 robomimic 实现的一个最小移植版本。
    """

    def __init__(self, input_shape, num_kp=None):
        """
        Args:
            input_shape (list): (C, H, W) 输入特征图形状。
            num_kp (int): 输出中的关键点数量。如果为 None，输出将具有与输入相同的通道数。
        """
        super().__init__()

        assert len(input_shape) == 3
        self._in_c, self._in_h, self._in_w = input_shape

        if num_kp is not None:
            self.nets = torch.nn.Conv2d(self._in_c, num_kp, kernel_size=1)
            self._out_c = num_kp
        else:
            self.nets = None
            self._out_c = self._in_c

        # 我们可以直接使用 torch.linspace，但这似乎与 numpy 的行为略有不同
        # 并且会导致预训练模型的 pc_success 略有下降。
        pos_x, pos_y = np.meshgrid(np.linspace(-1.0, 1.0, self._in_w), np.linspace(-1.0, 1.0, self._in_h))
        pos_x = torch.from_numpy(pos_x.reshape(self._in_h * self._in_w, 1)).float()
        pos_y = torch.from_numpy(pos_y.reshape(self._in_h * self._in_w, 1)).float()
        # 注册为 buffer，以便将其移动到正确的设备。
        self.register_buffer("pos_grid", torch.cat([pos_x, pos_y], dim=1))

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        """
        Args:
            features: (B, C, H, W) 输入特征图。
        Returns:
            (B, K, 2) 关键点的图像空间坐标。
        """
        if self.nets is not None:
            features = self.nets(features)

        # [B, K, H, W] -> [B * K, H * W]，其中 K 是关键点数量
        features = features.reshape(-1, self._in_h * self._in_w)
        # 2d softmax 归一化
        attention = F.softmax(features, dim=-1)
        # [B * K, H * W] x [H * W, 2] -> [B * K, 2] 用于 x 和 y 维度的空间坐标均值
        expected_xy = attention @ self.pos_grid
        # 重塑为 [B, K, 2]
        feature_keypoints = expected_xy.view(-1, self._out_c, 2)

        return feature_keypoints

class _SingleRgbEncoder(nn.Module):
    """单个摄像头的 RGB 编码器，支持独立或共享使用"""
    def __init__(
        self,
        vision_backbone: str,
        pretrained_backbone_weights: str | None,
        input_shape: Tuple[int, int, int],
        crop_shape: Optional[Tuple[int, int]],
        crop_is_random: bool,
        use_group_norm: bool,
        spatial_softmax_num_keypoints: int,
        freeze_backbone: bool = True,  # 新增：是否冻结backbone
    ):
        super().__init__()

        # 设置可选的预处理
        if crop_shape is not None:
            self.do_crop = True
            # 评估时始终使用中心裁剪
            self.center_crop = torchvision.transforms.CenterCrop(crop_shape)
            if crop_is_random:
                self.maybe_random_crop = torchvision.transforms.RandomCrop(crop_shape)
            else:
                self.maybe_random_crop = self.center_crop
        else:
            self.do_crop = False
            crop_shape = input_shape[1:]

        # 设置骨干网络
        backbone_model = getattr(torchvision.models, vision_backbone)(
            weights=pretrained_backbone_weights
        )

        # 移除 AvgPool 和 FC (假设 layer4 是 children()[-3])
        self.backbone = nn.Sequential(*(list(backbone_model.children())[:-2]))

        if use_group_norm:
            self.backbone = _replace_submodules(
                root_module=self.backbone,
                predicate=lambda x: isinstance(x, nn.BatchNorm2d),
                func=lambda x: nn.GroupNorm(num_groups=x.num_features // 16, num_channels=x.num_features),
            )

        # 冻结backbone参数（可选）
        if freeze_backbone:
            for param in self.backbone.parameters():
                param.requires_grad = False

        # 设置池化和最终层
        # 使用试运行来获取特征图形状
        dummy_shape = (1, input_shape[0], *crop_shape)
        with torch.no_grad():
            dummy_out = self.backbone(torch.zeros(dummy_shape))
        feature_map_shape = dummy_out.shape[1:]  # (C, H, W)

        self.pool = SpatialSoftmax(feature_map_shape, num_kp=spatial_softmax_num_keypoints)
        self.feature_dim = spatial_softmax_num_keypoints * 2
        self.out = nn.Linear(spatial_softmax_num_keypoints * 2, self.feature_dim)
        self.relu = nn.ReLU()

        # 注册ImageNet标准化参数为buffer（会自动移到GPU）
        self.register_buffer('mean', torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer('std', torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))

    def forward_single_image(self, x: torch.Tensor) -> torch.Tensor:
        if self.do_crop:
            x = self.maybe_random_crop(x) if self.training else self.center_crop(x)

        # ImageNet标准化（预训练权重期望的输入分布）
        x = (x - self.mean) / self.std

        x = self.relu(self.out(torch.flatten(self.pool(self.backbone(x)), start_dim=1)))
        return x


class ResNetDiffusionBackbone(VLABackbone):
    def __init__(
        self,
        vision_backbone: str = "resnet18",
        pretrained_backbone_weights: str | None = None,
        input_shape: Tuple[int, int, int] = (3, 84, 84), # (C, H, W)
        crop_shape: Optional[Tuple[int, int]] = None,
        crop_is_random: bool = True,
        use_group_norm: bool = True,
        spatial_softmax_num_keypoints: int = 32,
        use_separate_rgb_encoder_per_camera: bool = False,  # 新增：是否为每个摄像头使用独立编码器
        num_cameras: int = 1,  # 新增：摄像头数量（仅在独立编码器模式下使用）
        freeze_backbone: bool = True,  # 新增：是否冻结ResNet backbone（推荐True）
    ):
        super().__init__()

        self.use_separate_rgb_encoder_per_camera = use_separate_rgb_encoder_per_camera
        self.num_cameras = num_cameras

        if use_separate_rgb_encoder_per_camera:
            # 独立编码器模式：为每个摄像头创建独立的编码器
            encoders = [
                _SingleRgbEncoder(
                    vision_backbone=vision_backbone,
                    pretrained_backbone_weights=pretrained_backbone_weights,
                    input_shape=input_shape,
                    crop_shape=crop_shape,
                    crop_is_random=crop_is_random,
                    use_group_norm=use_group_norm,
                    spatial_softmax_num_keypoints=spatial_softmax_num_keypoints,
                    freeze_backbone=freeze_backbone,
                )
                for _ in range(num_cameras)
            ]
            self.rgb_encoder = nn.ModuleList(encoders)
            # 重要：output_dim 始终表示单个编码器的特征维度（与 lerobot 保持一致）
            self.feature_dim = encoders[0].feature_dim
        else:
            # 共享编码器模式：所有摄像头共享同一个编码器
            self.rgb_encoder = _SingleRgbEncoder(
                vision_backbone=vision_backbone,
                pretrained_backbone_weights=pretrained_backbone_weights,
                input_shape=input_shape,
                crop_shape=crop_shape,
                crop_is_random=crop_is_random,
                use_group_norm=use_group_norm,
                spatial_softmax_num_keypoints=spatial_softmax_num_keypoints,
                freeze_backbone=freeze_backbone,
            )
            self.feature_dim = self.rgb_encoder.feature_dim

    def forward(self, images):
        """
        Args:
            images: Dict[str, Tensor], 每个摄像头的图像
                    形状: {cam_name: (B, T, C, H, W)}

        Returns:
            Tensor: (B, T, total_feature_dim)
        """
        any_tensor = next(iter(images.values()))
        B, T = any_tensor.shape[:2]
        cam_names = sorted(images.keys())

        if self.use_separate_rgb_encoder_per_camera:
            # 独立编码器模式：每个摄像头使用对应的编码器
            features_all = []
            for cam_idx, cam_name in enumerate(cam_names):
                img = images[cam_name]
                encoder = self.rgb_encoder[cam_idx]
                features = encoder.forward_single_image(img.view(B * T, *img.shape[2:]))
                features_all.append(features)
            return torch.cat(features_all, dim=1).view(B, T, -1)
        else:
            # 共享编码器模式：所有摄像头共享同一个编码器
            features_all = []
            for cam_name in cam_names:
                img = images[cam_name]
                features = self.rgb_encoder.forward_single_image(img.view(B * T, *img.shape[2:]))
                features_all.append(features)
            return torch.cat(features_all, dim=1).view(B, T, -1)

    @property
    def output_dim(self):
        return self.feature_dim

if __name__ == "__main__":
    print("=" * 60)
    print("🚀 Testing ResNetDiffusionBackbone")
    print("=" * 60)

    # Configuration
    B, T = 2, 5
    C, H, W = 3, 96, 96
    crop_h, crop_w = 84, 84
    num_keypoints = 32
    feature_dim_per_cam = num_keypoints * 2

    # Create dummy input (2 cameras)
    images = {
        "cam_high": torch.randn(B, T, C, H, W),
        "cam_wrist": torch.randn(B, T, C, H, W)
    }
    num_cameras = len(images)

    # ============================================================================
    # Test 1: Shared Encoder (默认模式)
    # ============================================================================
    print("\n[Test 1] Shared Encoder Mode")
    print("-" * 60)
    backbone_shared = ResNetDiffusionBackbone(
        vision_backbone="resnet18",
        pretrained_backbone_weights=None,  # Speed up test
        input_shape=(C, H, W),
        crop_shape=(crop_h, crop_w),
        crop_is_random=True,
        use_group_norm=True,
        spatial_softmax_num_keypoints=num_keypoints,
        use_separate_rgb_encoder_per_camera=False,  # 共享编码器
    )

    print(f"✅ Shared encoder model instantiated")
    print(f"   Output dim per camera: {feature_dim_per_cam}")
    print(f"   Number of cameras: {num_cameras}")
    print(f"   Expected total dim: {num_cameras * feature_dim_per_cam}")

    output = backbone_shared(images)
    print(f"\n🔄 Forward pass completed")
    print(f"   Input shapes: {[v.shape for v in images.values()]}")
    print(f"   Output shape: {output.shape}")

    expected_dim = num_cameras * feature_dim_per_cam
    assert output.shape == (B, T, expected_dim), f"Expected shape {(B, T, expected_dim)}, got {output.shape}"
    print(f"✨ Test passed!")

    # ============================================================================
    # Test 2: Separate Encoders (独立编码器模式)
    # ============================================================================
    print("\n[Test 2] Separate Encoders Mode")
    print("-" * 60)
    backbone_separate = ResNetDiffusionBackbone(
        vision_backbone="resnet18",
        pretrained_backbone_weights=None,  # Speed up test
        input_shape=(C, H, W),
        crop_shape=(crop_h, crop_w),
        crop_is_random=True,
        use_group_norm=True,
        spatial_softmax_num_keypoints=num_keypoints,
        use_separate_rgb_encoder_per_camera=True,  # 独立编码器
        num_cameras=num_cameras,
    )

    print(f"✅ Separate encoders model instantiated")
    print(f"   Output dim per camera: {feature_dim_per_cam}")
    print(f"   Number of cameras: {num_cameras}")
    print(f"   Number of encoders: {len(backbone_separate.rgb_encoder)}")

    output = backbone_separate(images)
    print(f"\n🔄 Forward pass completed")
    print(f"   Input shapes: {[v.shape for v in images.values()]}")
    print(f"   Output shape: {output.shape}")

    expected_dim = num_cameras * feature_dim_per_cam
    assert output.shape == (B, T, expected_dim), f"Expected shape {(B, T, expected_dim)}, got {output.shape}"
    print(f"✨ Test passed!")

    # ============================================================================
    # Test 3: Verify parameters count
    # ============================================================================
    print("\n[Test 3] Parameter Count Comparison")
    print("-" * 60)
    shared_params = sum(p.numel() for p in backbone_shared.parameters())
    separate_params = sum(p.numel() for p in backbone_separate.parameters())

    print(f"   Shared encoder parameters: {shared_params:,}")
    print(f"   Separate encoders parameters: {separate_params:,}")
    print(f"   Ratio: {separate_params / shared_params:.2f}x")

    assert separate_params > shared_params, "Separate encoders should have more parameters"
    print(f"✨ Verification passed!")

    # ============================================================================
    # Test 4: Verify independent parameters
    # ============================================================================
    print("\n[Test 4] Verify Independent Parameters")
    print("-" * 60)
    # Check that encoders have independent parameters
    encoder_0_first_param = list(backbone_separate.rgb_encoder[0].parameters())[0]
    encoder_1_first_param = list(backbone_separate.rgb_encoder[1].parameters())[0]

    # Modify first encoder's parameter
    with torch.no_grad():
        encoder_0_first_param += 1.0

    # Verify they are not the same tensor
    assert not torch.allclose(encoder_0_first_param, encoder_1_first_param), \
        "Encoders should have independent parameters"

    print(f"✅ Encoders have independent parameters")
    print(f"✨ All tests passed!")

    print("\n" + "=" * 60)
    print("🎉 All tests completed successfully!")
    print("=" * 60)