import timm
import torch
import torch.nn as nn
from src.gumbel_softmax import GumbelSoftmax, GumbelSoftmaxTopK

class SaliencyNet(nn.Module):
    def __init__(self, cfg):

        super(SaliencyNet, self).__init__()
        self.k = 8
        self.tau = cfg.init_tau
        self.min_tau = cfg.min_tau
        self.decay_factor = cfg.decay_factor
        self.raw_backbone_path = cfg.raw_backbone_path
        
        self.configure_pretrained(cfg.backbone, cfg.n_trainable_backbone_layers) 
        embed_dim = self.pretrained.num_features

        self.fc1 = nn.Linear(embed_dim, cfg.fc1_size)
        self.fc2 = nn.Linear(cfg.fc1_size, cfg.fc2_size)
        self.fc3 = nn.Linear(cfg.fc2_size, cfg.fc3_size)
        self.fc4 = nn.Linear(cfg.fc3_size, 1)
        # self.bn1 = nn.BatchNorm1d(fc1_size)  # Batch normalization layer
        self.dropout = nn.Dropout(cfg.dropout_ratio)
        self.sigmoid = nn.Sigmoid()
        
        if cfg.activation == 'relu':
        self.activation = nn.ReLU()
        elif cfg.activation == 'leakyrelu':
            self.activation = nn.LeakyReLU()
        elif cfg.activation == 'gelu':
            self.activation = nn.GELU()
        else:
            raise ValueError("Invalid activation type. Choose 'relu', 'leakyrelu', or 'gelu'.")

        self.normalization_layer = self.init_normalization_layer(cfg.normalization_layer, self.tau, self.k)

        self.initialize_weights()


    def init_normalization_layer(self, normalization_type, tau, k):
        if normalization_type == 'raw':
            normalization_layer = lambda x: x
        elif normalization_type == 'sigmoid':
            normalization_layer = nn.Sigmoid()
        elif normalization_type == 'softmax':
            normalization_layer = nn.Softmax(dim=1)
        elif normalization_type == 'gumbel':
            normalization_layer = GumbelSoftmax(tau)
        elif normalization_type == 'k_top':
            normalization_layer = GumbelSoftmaxTopK(k, tau)

        return normalization_layer

    def initialize_weights(self):
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)

    def forward(self, preprocessed_images):
        features = self.pretrained(preprocessed_images)

        x = self.activation(self.fc1(features))
        x = self.dropout(x)
        x = self.activation(self.fc2(x))
        x = self.dropout(x)
        x = self.activation(self.fc3(x))
        x = self.fc4(x)

        frame_scores = torch.flatten(x)
        frame_weights = torch.flatten(self.normalization_layer(x))
        
        return frame_scores, frame_weights, features
    
                

    def anneal_tau(self):
        """ Anneals the temperature parameter τ using exponential decay """
        self.tau = self.tau * torch.exp(-self.decay_factor * torch.tensor(1.0))
        self.tau = max(self.tau, self.min_tau)  # To prevent τ from becoming too small

    def configure_pretrained(self, backbone, n_trainable_backbone_layers=1):
        # backbones = ['resnet18', 'efficientnet_b0', 'mobilenetv2_100', 'mobilenetv3_large_100']
        self.pretrained = timm.create_model(backbone, pretrained=False, checkpoint_path=self.raw_backbone_path)
        
        # Freeze the early layers of the model
        for param in self.pretrained.parameters():
            param.requires_grad = False
            
        # disable last layer of resnet18
        if backbone =='resnet18':
            self.pretrained.fc = nn.Identity()
            sequential_modules = [child for child in self.pretrained.children() if isinstance(child, nn.Sequential)]
        
        elif backbone in ['efficientnet_b0', 'mobilenetv2_100','mobilenetv3_large_100'] :
            self.pretrained.classifier = nn.Identity()
            sequential_modules = [child for child in self.pretrained.blocks.children()]
        
        else:
            assert('Not supported backbone model!')

        num_seq_blocks = len(sequential_modules)
        
        if n_trainable_backbone_layers ==0:
            pass

        elif n_trainable_backbone_layers >= num_seq_blocks:
            for param in self.pretrained.parameters():
                param.requires_grad = True
        else:
            # Select the last `n_trainable_backbone_layers` of sequential modules to be trainable
            layers_to_train = sequential_modules[-n_trainable_backbone_layers:]
            for layer in layers_to_train:
                for param in layer.parameters():
                    param.requires_grad = True
            

            if n_trainable_backbone_layers > 0:
                last_layers= []

                if hasattr(self.pretrained, 'conv_head'):
                    last_layers.append(self.pretrained.conv_head.weight)
                
                if hasattr(self.pretrained, 'bn2'):
                    last_layers.append(self.pretrained.bn2.weight)
                    last_layers.append(self.pretrained.bn2.bias)
                
                for param in last_layers: 
                    param.requires_grad = True