import torch
import torch.nn as nn
import math
#from transformers import AutoModelForCausalLM, AutoTokenizer

class LSTM(nn.Module):
    def _init_(self, input_size, lstm_layer_sizes,linear_layer_size, output_size):
        super(LSTM, self)._init_()

        self.input_size = input_size
        self.linear_layer_size = linear_layer_size

        self.lstm_layer_1 = nn.LSTM(input_size, lstm_layer_sizes[0], batch_first=True)
        self.lstm_layer_2 = nn.LSTM(lstm_layer_sizes[0], lstm_layer_sizes[1], batch_first=True)
        self.lstm_layer_3 = nn.LSTM(lstm_layer_sizes[1], lstm_layer_sizes[2], batch_first=True)

        self.fc = Linear(lstm_layer_sizes[2], self.linear_layer_size,output_size)

        self.apply(self.initialize_weights)

    def forward(self, x):

        out, (hn_1, cn_1) = self.lstm_layer_1(x)
        out, (hn_2, cn_2) = self.lstm_layer_2(out)
        out, (hn_3, cn_3) = self.lstm_layer_3(out)

        out = hn_3[-1]
        out = self.fc(out)
        return out
    
    def initialize_weights(self, layer):
        if isinstance(layer, nn.Linear):
            nn.init.xavier_uniform_(layer.weight)
            nn.init.zeros_(layer.bias)
        elif isinstance(layer, nn.LSTM):
            for name, param in layer.named_parameters():
                if 'weight' in name:
                    nn.init.xavier_uniform_(param.data)
                elif 'bias' in name:
                    nn.init.zeros_(param.data)

class Linear(nn.Module):
    def _init_(self,input_size,hidden_sizes,output_size):
        super(Linear,self)._init_()
        
        self.relu =nn.ReLU()
        self.sigmoid =nn.Sigmoid()
        self.tanh = nn.Tanh()
        self.input = nn.Linear(input_size,hidden_sizes[0])
        self.fc = nn.Linear(hidden_sizes[0],hidden_sizes[1])
        self.output = nn.Linear(hidden_sizes[1],output_size)

        self.apply(self.initialize_weights)

    def forward(self,x):
        out = self.relu(self.input(x))
        out = self.relu(self.fc(out))
        out = self.relu(self.output(out))
        return out
    
    def initialize_weights(self, layer):
        if isinstance(layer, nn.Linear):
            nn.init.xavier_uniform_(layer.weight)
            nn.init.zeros_(layer.bias)
            
class LUCLSTM(nn.Module):
    def _init_(self, input_size, lstm_layer_sizes, output_size):
        super(LUCLSTM, self)._init_()

        self.input_size = input_size

        self.lstm_layer_1 = nn.LSTM(input_size, lstm_layer_sizes[0], batch_first=True)
        self.lstm_layer_2 = nn.LSTM(lstm_layer_sizes[0], lstm_layer_sizes[1], batch_first=True)
        self.lstm_layer_3 = nn.LSTM(lstm_layer_sizes[1], lstm_layer_sizes[2], batch_first=True)

        self.fc = nn.Linear(lstm_layer_sizes[2],64)
        self.fc2 = nn.Linear(64,output_size)
        self.tanh = nn.Tanh()
        self.relu =nn.ReLU()

        self.apply(self.initialize_weights)

    def forward(self, x):

        out, (hn_1, cn_1) = self.lstm_layer_1(x)
        out, (hn_2, cn_2) = self.lstm_layer_2(out)
        out, (hn_3, cn_3) = self.lstm_layer_3(out)

        out = hn_3[-1]
        out = self.tanh(self.fc(out))
        out = self.fc2(out)
        return out
    
    def initialize_weights(self, layer):
        if isinstance(layer, nn.Linear):
            nn.init.xavier_uniform_(layer.weight)
            nn.init.zeros_(layer.bias)
        elif isinstance(layer, nn.LSTM):
            for name, param in layer.named_parameters():
                if 'weight' in name:
                    nn.init.xavier_uniform_(param.data)
                elif 'bias' in name:
                    nn.init.zeros_(param.data)
    
    
class PositionalEncoding(nn.Module):
    def _init_(self, dim, max_len=300):
        super(PositionalEncoding, self)._init_()
        pe = torch.zeros(max_len, dim)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, dim, 2).float() * (-math.log(10000.0) / dim))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:x.size(0), :]
    
class Transformer(nn.Module):
    def _init_(self):
        super(Transformer,self)._init_()