import torch
import torch.nn as nn
import torch.nn.functional as F


class CondGraspNet(nn.Module):
    def __init__(self):
        super(CondGraspNet, self).__init__()

        # === 1. 定义输入维度 ===
        # 触觉特征: 12维 (3指 * 2单元 * 2分量)
        self.tactile_dim = 12
        # 构型特征: 3维 (One-Hot编码: [1,0,0], [0,1,0], [0,0,1])
        self.config_dim = 3

        input_total_dim = self.tactile_dim + self.config_dim  # 15维

        # === 2. 定义网络层 (MLP结构) ===

        # Layer 1: 特征融合层
        # 将触觉信息和构型信息混合
        self.fc1 = nn.Linear(input_total_dim, 64)
        self.bn1 = nn.BatchNorm1d(64)  # 批归一化: 防止梯度消失，加速训练

        # Layer 2: 非线性映射层
        # 增加网络宽度，拟合复杂的力学关系
        self.fc2 = nn.Linear(64, 128)
        self.bn2 = nn.BatchNorm1d(128)

        # Layer 3: 特征压缩层
        self.fc3 = nn.Linear(128, 64)

        # Layer 4: 输出层 (Regression Head)
        # 输出3个值: [Delta_X, Delta_Y, Delta_Theta]
        self.output = nn.Linear(64, 3)

        # === 3. 权重初始化 (Xavier) ===
        # 这一步对小数据集训练非常重要，能让模型收敛得更快
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(self, tactile_data, config_id_idx):
        """
        前向传播函数
        :param tactile_data: [Batch_Size, 12] 的触觉数据张量
        :param config_id_idx: [Batch_Size] 的构型索引 (例如 [0, 2, 1...])
        :return: [Batch_Size, 3] 的预测偏差
        """

        # Step 1: 处理构型 ID (One-Hot Encoding)
        # 必须把整数 ID (0,1,2) 变成向量 ([1,0,0]...) 才能喂给神经网络
        batch_size = tactile_data.size(0)

        # 创建一个全0的容器
        config_one_hot = torch.zeros(batch_size, self.config_dim).to(tactile_data.device)

        # 使用 scatter_ 方法进行填充
        # config_id_idx 需要升维: [Batch] -> [Batch, 1]
        config_one_hot.scatter_(1, config_id_idx.unsqueeze(1).long(), 1)

        # Step 2: 特征拼接 (Concatenate)
        # 将触觉数据和构型向量拼在一起 -> [Batch, 15]
        x = torch.cat((tactile_data, config_one_hot), dim=1)

        # Step 3: 通过隐藏层
        x = F.relu(self.bn1(self.fc1(x)))  # Linear -> BN -> ReLU
        x = F.relu(self.bn2(self.fc2(x)))  # Linear -> BN -> ReLU
        x = F.relu(self.fc3(x))  # Linear -> ReLU (最后一层通常不用BN)

        # Step 4: 输出结果
        prediction = self.output(x)

        return prediction


# === 单元测试 (Unit Test) ===
# 运行此文件，检查网络结构和输入输出形状是否正确
if __name__ == "__main__":
    print("Testing CondGraspNet Model...")

    # 1. 实例化模型
    model = CondGraspNet()
    print(f"Model Structure:\n{model}")

    # 2. 创建模拟输入数据 (Batch Size = 8)
    # 模拟8条触觉数据 (随机数)
    fake_tactile = torch.randn(8, 12)
    # 模拟8个构型ID (随机 0, 1, 2)
    fake_config = torch.tensor([0, 0, 1, 1, 2, 2, 0, 2], dtype=torch.long)

    # 3. 前向推理
    print("\nProcessing forward pass...")
    try:
        output = model(fake_tactile, fake_config)

        # 4. 验证结果
        print("Input Shape (Tactile):", fake_tactile.shape)
        print("Output Shape (Pred):  ", output.shape)  # 期望是 [8, 3]
        print("\nSample Prediction (Row 0):")
        print(f"Delta X: {output[0][0].item():.4f} mm")
        print(f"Delta Y: {output[0][1].item():.4f} mm")
        print(f"Delta θ: {output[0][2].item():.4f} deg")

        if output.shape == (8, 3):
            print("\n✅ 测试通过：网络维度正确！")
    except Exception as e:
        print(f"\n❌ 测试失败：{e}")