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
from gumbel_softmax_trick import gumbel_sigmoid, clamp_probs

class PARDNet(nn.Module):
    def __init__(self, configs):
        super(PARDNet, self).__init__()
        self.configs = configs
        init_random_seed(self.configs['rand_seed'])
        
        # Inference model (q(y|x,s))
        self.inf_model = MLP(configs['in_features']+configs['num_classes'], 256, [256, 512], False,
                           configs['drop_ratio'], "relu")
        self.inf_model_fc = nn.Linear(256, configs['num_classes'])
        
        # Generative model (p(s|x,y))
        self.gen_model = MLP(configs['in_features']+configs['num_classes'], 512, [256], False,
                             nonlinearity="relu")
        self.gen_model_fc = nn.Linear(512, configs['num_classes'])
        
        # Prediction model (p(y|x))
        self.classifier = nn.Linear(configs['in_features'], configs['num_classes'])
        
        # 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.inf_model.reset_parameters()
        self.inf_model_fc.reset_parameters()
        self.gen_model.reset_parameters()
        self.gen_model_fc.reset_parameters()
        self.classifier.reset_parameters()
    
    def get_config_optim(self):
        return [{'params': self.inf_model.parameters()},
                {'params': self.inf_model_fc.parameters()},
                {'params': self.gen_model.parameters()},
                {'params': self.gen_model_fc.parameters()},
                {'params': self.classifier.parameters()}]
    
    def forward(self, input: Tensor, candidate_labels: Tensor) -> Tuple[Tensor, ...]:
        # Ground-truth label identification (q(y|x,s))
        gt_convidence = self._encode(input, candidate_labels) # [B x Q]
        
        # Sample ground-truch labels with gumbel-softmax trick
        if self.training:
            gt_targets = gumbel_sigmoid(gt_convidence, tau=self.configs['tau'], gumbel_noise=True, hard=False) # [B x Q]
        else:
            # Straight through.
            gt_targets = (gt_convidence > 0.5).float()
        
        # Recover observed candidate labels (p(s|x,y))
        rec_convidence = self._decode(input, gt_targets) # [B x Q]
        
        # Classification (p(y|x))
        preds = self.classifier(input) # [B x Q]
 
        return gt_convidence, rec_convidence, 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 = self._compute_loss(*preds, targets) 
        
        return {'Loss': Loss}
    
    def loss_function_eval(self, preds: Tuple[Tensor, ...], targets: Tensor) -> dict:
        Loss = self._compute_loss(*preds, targets) 
        
        return {'Loss': Loss.detach().item()}
    
    def predict(self, input: Tensor) -> Tuple[Tensor, Tensor]:
        self.eval()
        with torch.no_grad():
            # Classification (p(y|x))
            pred_probs = self.classifier(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]:
        optimizer = torch.optim.Adam(self.get_config_optim(), lr=self.configs['lr'],
                                     weight_decay=self.configs['weight_decay'])
        
        return optimizer
        
    def _encode(self, input: Tensor, candidate_labels: Tensor) -> Tensor:
        result = self.inf_model(torch.cat([input, candidate_labels], dim=1))
        result = torch.sigmoid(self.inf_model_fc(result)) * candidate_labels # [B x Q]
        
        return result
    
    def _decode(self, input: Tensor, gt_labels: Tensor) -> Tensor:
        result = self.gen_model(torch.cat([input, gt_labels], dim=1))
        result = torch.sigmoid(self.gen_model_fc(result)) # [B x Q]
        
        return result
    
    def _compute_loss(self, gt_convidence: Tensor, rec_convidence: Tensor,
                      preds: Tensor, targets: Tensor) -> Tuple[Tensor, ...]:
        # KL-divergence
        Cls_loss = F.multilabel_soft_margin_loss(preds, gt_convidence.detach())
        gt_convidence = clamp_probs(gt_convidence)
        Ent_loss = -torch.mean(gt_convidence * torch.log(gt_convidence) + (1-gt_convidence) * torch.log1p(-gt_convidence))
        Kl_loss = Cls_loss - Ent_loss
        
        # Reconstruction loss
        Rec_loss = F.binary_cross_entropy(rec_convidence, targets) * targets.size(1)
        
        Loss = Rec_loss + self.configs['alpha'] * Kl_loss
        
        return Loss