# This file contains the neural network model for CTMeth
# Tokenizer class is used to tokenize the input data - it is a Conv1d->ReLu->MaxPool1d
# TransformerClassifier class is used to classify the input data

import torch
import torch.nn as nn
from torch.nn import (
    Module,
    ModuleList,
    Linear,
    Dropout,
    LayerNorm,
    Parameter,
    init,
)
import torch.nn.functional


# Method used for weights initialization in tokenizer.
# It is used to prevent exploding/vanishing gradient
# It uses kaiming_normal because ReLu is used as a activation layer


def init_weight(m):
    if isinstance(m, nn.Conv1d):
        nn.init.kaiming_normal_(m.weight)


# Tokenizeer


class Tokenizer(nn.Module):
    def __init__(
        self,
        # MaxPool1d
        MaxPool1d_kernel_size,
        MaxPool1d_stride,
        MaxPool1d_padding,
        # Conv1d
        conv1d_bias,
        conv1d_kernel_size,
        conv1d_stride,
        conv1d_padding,
        d_model,
        # Sequence length
        beta_values_array_length,
        # Else
        *args,
        **kwargs
    ):
        super(Tokenizer, self).__init__()
        self.conv1d_output_channels = d_model
        self.conv1d_lay = nn.Sequential(
            nn.Conv1d(
                1,
                out_channels=self.conv1d_output_channels,
                kernel_size=conv1d_kernel_size,
                stride=conv1d_stride,
                padding=conv1d_padding,
                bias=conv1d_bias,
            ),
            nn.ReLU(),
            nn.MaxPool1d(
                kernel_size=MaxPool1d_kernel_size,
                stride=MaxPool1d_stride,
                padding=MaxPool1d_padding,
            ),
        )

        self.apply(init_weight)

    # Method used to calculate the length of the sequence after passing through the tokenizer
    def sequence_length(self, beta_values_array_length, *args, **kwargs):
        forward_sequence_length = self.forward(
            torch.zeros((1, 1, beta_values_array_length))
        ).shape[1]
        return forward_sequence_length

    def forward(self, x):

        to_transformer = self.conv1d_lay(x)  # Sequence of -> Conv1d->ReLu->MaxPool1d
        to_transformer = to_transformer.transpose(-2, -1)

        # data ir ready to be passed to transformer
        return to_transformer


# Transformer


class TransformerClassifier(Module):
    def __init__(
        self,
        d_model,  # Dimensionality of the input
        num_classes,  # Number of classes
        sequence_length,  # Length of the sequence
        num_layers,  # Number of transformer block
        nhead,  # Number of heads in the multi-head attention models
        mlp_ratio,  # Ratio of the hidden layer size to the input size
        dropout_stage1,  # Dropout rate for embedding layer
        dropout_stage2,  # Dropout rate for transformer block
        *args,
        **kwargs
    ):
        super().__init__()

        self.dropout_stage1 = dropout_stage1  # Dropout rate for embedding layer
        self.drouput_stage2 = dropout_stage2  # Dropout rate for transformer block
        self.d_model = d_model  # Dimensionality of the input
        self.dim_feedforward = int(self.d_model * mlp_ratio)

        self.sequence_length = sequence_length
        self.nhead = nhead
        self.attention_pool = Linear(self.d_model, 1)
        self.attention_pool2 = Linear(self.d_model * self.sequence_length, 3)
        self.attention_pool3 = Linear(self.sequence_length, 3)

        self.attention_pool4 = Linear(self.d_model, 3)
        self.decoder2 = Linear(3, 3)

        self.positional_emb = Parameter(
            torch.zeros(1, self.sequence_length, self.d_model), requires_grad=True
        )
        self.a_one = torch.ones(1, self.d_model)
        init.trunc_normal_(self.positional_emb, std=0.2)

        self.dropout = Dropout(p=dropout_stage1)

        self.blocks = ModuleList(
            [
                torch.nn.TransformerEncoderLayer(
                    self.d_model,
                    self.nhead,
                    dim_feedforward=self.dim_feedforward,
                    dropout=dropout_stage2,
                    activation="relu",
                    batch_first=False,
                    norm_first=False,
                )
                for i in range(num_layers)
            ]
        )

        self.norm = LayerNorm(self.d_model)

        self.decoder = Linear(self.d_model, num_classes)
        self.apply(self.init_weight)

    def forward(self, x):

        x += self.positional_emb
        x = self.dropout(x)

        n = 0
        for a_block in self.blocks:
            x = a_block(x)
            n += 1

        x = self.norm(x)
        z = self.attention_pool(x)
        z = z.transpose(-1, -2)
        zz = self.attention_pool3(z)
        zz = self.decoder2(zz)
        zz = zz.reshape([1, 3])

        # retuns labels
        return zz

    @staticmethod
    def init_weight(m):
        if isinstance(m, Linear):
            init.trunc_normal_(m.weight, std=0.02)
            if isinstance(m, Linear) and m.bias is not None:
                init.constant_(m.bias, 0)
        elif isinstance(m, LayerNorm):
            init.constant_(m.bias, 0)
            init.constant_(m.weight, 1.0)


class CTMeth_neural(nn.Module):
    def __init__(self, params, *args, **kwargs):
        super(CTMeth_neural, self).__init__()
        self.tokenizer = Tokenizer(**params)
        self.classifier = TransformerClassifier(
            sequence_length=self.tokenizer.sequence_length(**params), **params
        )

    def forward(self, x):
        x = self.tokenizer(x)
        return self.classifier(x)
