import torch
import torch.nn as nn
import torch.nn.functional as F

from utils import init_random_seed
from type_def import *
from layers.MLP import MLP

class MainNet(nn.Module):
    def __init__(self, configs):
        super(MainNet, self).__init__()
        self.configs = configs
        
        # Embedding function
        self.encoder = MLP(configs['in_features'], configs['latent_dim'], [512], False,
                           configs['drop_ratio'], "relu")
        
        # Classifiers
        self.classifier = nn.Conv1d(configs['num_classes'], configs['num_classes'], 
                                    configs['latent_dim'],
                                    groups=configs['num_classes'])
        
        # Move model to the right device for consistent initialization
        self.to(configs['device'])
        
    def reset_parameters(self):
        self.encoder.reset_parameters()
        self.classifier.reset_parameters()
        
    def forward(self, input: Tensor, scale: Tensor=None, bias: Tensor=None) -> Tuple[Tensor, ...]:
        # Obtain latent representation of data
        z = self.encoder(input) # [B x D]
        
        # Perform label-specific feature correction
        if self.training and (scale is not None):
            z_corrected = bias + scale * z.unsqueeze(1).expand(-1, self.configs['num_classes'], -1) # [B x Q x D]
        else:
            z_corrected = z.unsqueeze(1).expand(-1, self.configs['num_classes'], -1) # [B x Q x D]
        
        # Classification
        preds = self.classifier(z_corrected).squeeze(2) # [B x Q]
 
        return preds
    
    def forward_enc(self, input: Tensor) -> Tensor:
        return self.encoder(input)


class MetaNet(nn.Module):
    def __init__(self, configs):
        super(MetaNet, self).__init__()
        self.configs = configs
        
        # Scaling and translation factors for feature correction
        self.scale = nn.Parameter(torch.zeros(configs['train_batch_size'], configs['num_classes'], 1))
        self.bias = nn.Parameter(torch.zeros(configs['train_batch_size'], configs['num_classes'], 1))
        
        # Move model to the right device for consistent initialization
        self.to(configs['device'])
        
    def reset_parameters(self):
        nn.init.zeros_(self.scale)
        nn.init.zeros_(self.bias)
        
    def reset_coef(self):
        nn.init.zeros_(self.scale)
        nn.init.zeros_(self.bias)
        
    def get_config_optim(self):
        return [{'params': self.scale, 'weight_decay': 0.0},
                {'params': self.bias, 'weight_decay': 0.0}]
        
    def forward(self, input: Tensor, reuse_eps=False) -> Tuple[Tensor, ...]:
        # Generate scaling and translation factors to correct features for each class label
        if self.training:
            scale = F.softplus(self.scale[:input.size(0)]) # [B x Q x 1]
            bias = self.bias[:input.size(0)] # [B x Q x 1]
        else:
            scale = None
            bias = None
 
        return scale, bias


class PASENet(nn.Module):
    def __init__(self, configs):
        super(PASENet, self).__init__()
        self.configs = configs
        init_random_seed(self.configs['rand_seed'])
        
        # Create model
        self.main_net = MainNet(self.configs) # prediction model
        
        self.meta_net = MetaNet(self.configs) # feature correction function
        
        # Move model to the right device for consistent initialization
        self.to(configs['device'])
        
        self.reset_parameters()
        
    def reset_parameters(self):
        init_random_seed(self.configs['rand_seed'])
        self.meta_net.reset_parameters()
        self.main_net.reset_parameters()
    
    def copy_main_net(self):
        """
        Return a copy of main_net for meta-training.
        """
        main_net_temp = MainNet(self.configs)
        main_net_temp.load_state_dict(self.main_net.state_dict())
        main_net_temp.train()
        
        return main_net_temp
    
    def forward(self, input: Tensor, main_net_temp=None) -> Tuple[Tensor, ...]:
        if main_net_temp is not None:
            scale, bias = self.meta_net(input)
            preds = main_net_temp(input, scale, bias)
        else:
            if self.training:
                scale, bias = self.meta_net(input, reuse_eps=True)
                preds = self.main_net(input, scale, bias)
            else:
                preds = self.main_net(input)
            
        return preds,
    
    def forward_meta(self, input: Tensor, main_net_temp) -> Tuple[Tensor, ...]:
        preds = main_net_temp(input)
        
        return preds,
    
    def training_start(self, train_dataloader, val_dataloader=None):
        '''
        Prepare for training.
        '''
        self.iters_per_epoch = len(train_dataloader)
    
    def loss_function_train(self, preds: Tuple[Tensor, ...], targets: Tensor) -> dict:
        Loss, Cls_loss = self._compute_loss(*preds, targets) 
        
        return {'Loss': Loss,
                'Cls_loss': Cls_loss}
    
    def loss_function_meta(self, preds: Tuple[Tensor, ...], targets: Tensor) -> dict:
        Loss = self._compute_loss_meta(*preds, targets)
        
        return {'Loss': Loss}
    
    def loss_function_eval(self, preds: Tuple[Tensor, ...], targets: Tensor) -> dict:
        Loss, Cls_loss = self._compute_loss(*preds, targets)
        
        return {'Loss': Loss.detach().item(),
                'Cls_loss': Cls_loss.detach().item()}
    
    def predict(self, input: Tensor) -> Tuple[Tensor, Tensor]:
        self.eval()
        with torch.no_grad():
            # Classification
            pred_probs = self.main_net(input).sigmoid_() # [B x Q]
            pred_labels = (pred_probs > self.configs['threshold']).type_as(pred_probs)
        
        return pred_labels, pred_probs
    
    def configure_optimizers(self) -> Tuple[Any, Any]:
        main_optimizer = torch.optim.Adam(self.main_net.parameters(), lr=self.configs['lr'],
                                          weight_decay=self.configs['weight_decay'])
        meta_optimizer = torch.optim.Adam(self.meta_net.get_config_optim(), lr=self.configs['lr_meta'],
                                          weight_decay=self.configs['weight_decay'])
            
        return main_optimizer, None, meta_optimizer, None
    
    def _compute_loss(self, preds: Tensor, targets: Tensor) -> Tuple[Tensor, ...]:
        # Cls loss
        Cls_loss = F.multilabel_soft_margin_loss(preds, targets) * targets.size(1)
        
        Loss = Cls_loss
        
        return Loss, Cls_loss
    
    def _compute_loss_meta(self, preds: Tensor, targets: Tensor) -> Tuple[Tensor, ...]:
        # Cls loss
        Loss = F.multilabel_soft_margin_loss(preds, targets) * targets.size(1)
        
        return Loss
