import torch.nn as nn
import torch.nn.init as init
import torch

class Swish(nn.Module):
    def __init__(self):
        super(Swish, self).__init__()

    def forward(self, x):
        return x * torch.sigmoid(x)

class Adapter(nn.Module):
    def __init__(self, input_size, output_size,  adapter_norm="layer_norm", init_type="glorot",  query_length=32, dropout_prob=0.1):
        super().__init__()
        self.query_length = query_length
        self.dropout_prob = dropout_prob
        self.adapter_norm = adapter_norm

        self.dropout = nn.Dropout(p=self.dropout_prob)
        
        self.c_fc = nn.Linear(input_size, input_size*2)
        self.act = Swish()
        self.c_proj = nn.Linear(input_size*2, output_size)
        
        if adapter_norm == "layer_norm":
            self.norm = nn.LayerNorm([self.query_length, output_size])
        elif adapter_norm == "batch_norm":
            self.norm = nn.BatchNorm1d(self.query_length)

        self.init_type = init_type.lower()
        self._initialize_weights()

    def forward(self, hidden_states):
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.c_fc(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.c_proj(hidden_states)
        hidden_states = self.norm(hidden_states)
        return hidden_states
    
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                if self.init_type == "glorot":
                    init.xavier_uniform_(m.weight)
                    if m.bias is not None:
                        init.constant_(m.bias, 0)
                elif self.init_type == "normal":
                    init.normal_(m.weight, mean=0, std=0.01)
                    if m.bias is not None:
                        init.constant_(m.bias, 0)
                else:
                    raise ValueError("Invalid initialization type specified.")