import numpy as np
import torch.nn as nn
from Qtorch.Models.Qnn import Qnn


class Qmlp(Qnn):

  def __init__(self, data,
               hidden_layers, 
               labels=None,
               dropout_rate=0.3,
               test_size = 0.2, 
               random_state=None
               ):
    super(Qmlp, self).__init__(data=data, labels=labels, test_size=test_size, random_state=random_state)

    input_size = self.X_train.shape[1]
    num_classes = len(np.unique(self.y_train))
    self.layers = nn.ModuleList()
    
    # 连接输入层和第一个隐藏层
    self.layers.append(nn.Linear(input_size, hidden_layers[0]))
    self.layers.append(nn.BatchNorm1d(hidden_layers[0]))
    self.layers.append(nn.ReLU())
    self.layers.append(nn.Dropout(dropout_rate))
    
    # 创建隐藏层
    for i in range(1, len(hidden_layers)):
        self.layers.append(nn.Linear(hidden_layers[i-1], hidden_layers[i]))
        self.layers.append(nn.BatchNorm1d(hidden_layers[i]))
        self.layers.append(nn.ReLU())
        self.layers.append(nn.Dropout(dropout_rate))
    
    # 创建输出层
    self.layers.append(nn.Linear(hidden_layers[-1], num_classes))
    self.__init_weights()
  
  def forward(self, x):
    for layer in self.layers:
        x = layer(x)
    return x

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