import torch
import torch.nn as nn
import numpy as np

from Qtorch.Models.Qnn import Qnn


class QCNN(Qnn):
    def __init__(
        self,
        data,
        labels=None,
        conv_channels=(16, 32),
        kernel_size=3,
        hidden_size=128,
        dropout_rate=0.3,
        test_size=0.2,
        random_state=None,
        batch_size=64,
        learning_rate=0.00001,
        weight_decay=1e-5,
        lr_scheduler_patience=10,
        early_stop_patience=100,
        early_stop_threshold=0.99,
    ):
        super(QCNN, self).__init__(
            data=data,
            labels=labels,
            test_size=test_size,
            random_state=random_state,
            batch_size=batch_size,
            learning_rate=learning_rate,
            weight_decay=weight_decay,
            lr_scheduler_patience=lr_scheduler_patience,
            early_stop_patience=early_stop_patience,
            early_stop_threshold=early_stop_threshold,
        )

        self.conv_channels = tuple(conv_channels)
        self.kernel_size = kernel_size
        self.hidden_size = hidden_size
        self.dropout_rate = dropout_rate

        self.feature_extractor = nn.Sequential()
        self.classifier = nn.Sequential()

        # 构造 1D CNN 网络结构
        self.build_model(input_shape=self.X_train.shape[1:], num_classes=self.num_classes)
        self._model_built = True

    def _transform_features(self, features):
        # 1D CNN 输入格式: [batch, channel=1, length]
        return torch.tensor(features, dtype=torch.float32).unsqueeze(1)

    def build_model(self, input_shape, num_classes):
        if len(self.conv_channels) == 0:
            raise ValueError("'conv_channels' must contain at least one channel size.")

        input_length = int(np.prod(input_shape))

        conv_layers = []
        in_channels = 1
        for out_channels in self.conv_channels:
            conv_layers.append(nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=self.kernel_size))
            conv_layers.append(nn.ReLU())
            conv_layers.append(nn.MaxPool1d(kernel_size=2))
            in_channels = out_channels
        self.feature_extractor = nn.Sequential(*conv_layers)

        conv_output_size = self._get_conv_output_size(input_length)
        self.classifier = nn.Sequential(
            nn.Linear(conv_output_size, self.hidden_size),
            nn.ReLU(),
            nn.Dropout(self.dropout_rate),
            nn.Linear(self.hidden_size, num_classes),
        )

        self.__init_weights()

    def _get_conv_output_size(self, input_length):
        x = torch.randn(1, 1, input_length)
        x = self.feature_extractor(x)
        return int(x.numel())

    def forward(self, x):
        x = self.feature_extractor(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

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