import mindspore
import mindspore.nn as nn
import mindspore.ops.operations as P


class AlexNet(nn.Cell):
    def __init__(self, input_channels, output_num, input_size):
        super().__init__()

        self.features = nn.SequentialCell([
            nn.Conv2d(in_channels=input_channels, out_channels=64, kernel_size=3, stride=2, pad_mode='pad', padding=1,
                      has_bias=True),
            nn.BatchNorm2d(num_features=64, momentum=0.9),  # 批量归一化层
            nn.MaxPool2d(kernel_size=2, stride=2, pad_mode='valid'),
            nn.ReLU(),

            nn.Conv2d(in_channels=64, out_channels=192, kernel_size=3, pad_mode='pad', padding=1, has_bias=True),
            nn.BatchNorm2d(num_features=192, momentum=0.9),  # 批量归一化层
            nn.MaxPool2d(kernel_size=2, stride=2, pad_mode='valid'),
            nn.ReLU(),

            nn.Conv2d(in_channels=192, out_channels=384, kernel_size=3, pad_mode='pad', padding=1, has_bias=True),
            nn.BatchNorm2d(num_features=384, momentum=0.9),  # 批量归一化层
            nn.ReLU(),

            nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, pad_mode='pad', padding=1, has_bias=True),
            nn.BatchNorm2d(num_features=256, momentum=0.9),  # 批量归一化层
            nn.ReLU(),

            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, pad_mode='pad', padding=1, has_bias=True),
            nn.BatchNorm2d(num_features=256, momentum=0.9),  # 批量归一化层
            nn.MaxPool2d(kernel_size=2, stride=2, pad_mode='valid'),
            nn.ReLU(),
        ])

        self.input_size = input_size
        self._init_classifier(output_num)

    def _init_classifier(self, output_num):
        # Forward a dummy input through the feature extractor part of the network
        dummy_input = mindspore.ops.zeros((1, 3, self.input_size, self.input_size))
        features_size = self.features(dummy_input).numel()

        self.classifier = nn.SequentialCell([
            nn.Dropout(p=0.5),
            nn.Dense(in_channels=features_size, out_channels=1000),
            nn.ReLU(),

            nn.Dropout(p=0.5),
            nn.Dense(in_channels=1000, out_channels=256),
            nn.ReLU(),

            nn.Dense(in_channels=256, out_channels=output_num)
        ])

    def construct(self, x):
        x = self.features(x)
        x = P.Reshape()(x, (P.Shape()(x)[0], -1,))
        x = self.classifier(x)
        return x


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='AlexNet Implementation')
    parser.add_argument('--input_channels', default=3, type=int)
    parser.add_argument('--output_num', default=10, type=int)
    parser.add_argument('--input_size', default=32, type=int)
    args = parser.parse_args()

    model = AlexNet(args.input_channels, args.output_num, args.input_size)
    tensor = mindspore.ops.rand((1, args.input_channels, args.input_size, args.input_size))
    pred = model(tensor)

    print(model)
    print("Predictions shape:", pred.shape)