import torch
import torch.nn as nn
from sklearn.preprocessing import LabelEncoder
from torch.utils.data import DataLoader, TensorDataset

from Qtorch.Models.Qnn import Qnn

class QCNN(Qnn):
    def __init__(self, X_train, y_train, X_test, y_test, labels=None, dropout_rate=0.3):
        super(QCNN, self).__init__()

        self.LABEL_ENCODER = LabelEncoder()

        self.X_train, self.y_train, self.X_test, self.y_test = X_train, y_train, X_test, y_test
        self.labels = labels

        input_size = X_train.shape[1]  # 输入的长度
        num_classes = len(set(y_train))  # 分类数

        # 网络层：卷积层 + 池化层 + 全连接层
        self.layers = nn.ModuleList()
        self.layers.append(nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3))  # 卷积层
        self.layers.append(nn.MaxPool1d(kernel_size=2))  # 池化层
        self.layers.append(nn.Conv1d(in_channels=16, out_channels=32, kernel_size=3))  # 卷积层
        self.layers.append(nn.MaxPool1d(kernel_size=2))  # 池化层

        # 计算展平后的大小
        conv_output_size = self._get_conv_output_size(input_size)  # 卷积后的输出大小
        print(f"Conv output size: {conv_output_size}")  # 打印卷积后的输出大小
        self.layers.append(nn.Linear(conv_output_size, 128))  # 全连接层
        self.layers.append(nn.Linear(128, num_classes))  # 输出层

        self.__init_weights()

    def _get_conv_output_size(self, input_size):
        # 计算卷积后的输出尺寸
        x = torch.randn(1, 1, input_size)  # 创建一个假的输入张量
        for layer in self.layers:
            x = layer(x)  # 通过每一层
        return int(x.numel())  # 返回展平后的输出大小


    def forward(self, x):
        # 通过卷积和池化层
        for layer in self.layers[:-2]:  # 除去最后两个 Linear 层
            x = layer(x)
        
        # 展平卷积后的输出
        x = x.view(x.size(0), -1)  # 这样 x 会变成 (batch_size, conv_output_size)
        
        # 通过全连接层
        x = self.layers[-2](x)
        x = self.layers[-1](x)
        
        return x

    def __init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv1d) or isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    m.bias.data.fill_(0.01)
    
    def _prepare_data(self):

        # 将data转换为tensor形式, unsqueeze可以创建多一维度的，给
        X_train_tensor = torch.tensor(self.X_train, dtype=torch.float32).unsqueeze(1)
        self.y_train = self.LABEL_ENCODER.fit_transform(self.y_train)
        y_train_tensor = torch.tensor(self.y_train, dtype=torch.long)
        
        X_test_tensor = torch.tensor(self.X_test, dtype=torch.float32).unsqueeze(1)
        self.y_test = self.LABEL_ENCODER.transform(self.y_test)
        y_test_tensor = torch.tensor(self.y_test, dtype=torch.long)
        
        train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
        test_dataset = TensorDataset(X_test_tensor, y_test_tensor)
        
        train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
        test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
        
        return train_loader, test_loader