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

from backbone.build_backbone import build_backbone

class MainNet(nn.Module):
    def __init__(self, args):
        super(MainNet, self).__init__()
        self.num_class = args.num_class
        # Embedding function
        self.backbone = build_backbone(args)
        self.pooling = nn.AdaptiveMaxPool2d((1, 1))

        # Classifiers
        self.classifier = nn.Conv1d(args.num_class, args.num_class, 
                                    self.backbone.num_features,
                                    groups=args.num_class)
        
        self.criterion = nn.MultiLabelSoftMarginLoss()
        
    def forward(self, input, scale=None, bias=None):
        # Obtain latent representation of data
        z = self.backbone(input)
        z = self.pooling(z).flatten(1) # [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.num_class, -1) # [B x Q x D]
        else:
            z_corrected = z.unsqueeze(1).expand(-1, self.num_class, -1) # [B x Q x D]
        
        # Classification
        preds = self.classifier(z_corrected).squeeze(2) # [B x Q]
 
        return preds
    
    def compute_loss(self, preds, targets):
        return self.criterion(preds, targets) * targets.size(1)
    
    def get_config_optim(self):
        param_dicts = [{"params": [p for n, p in self.named_parameters() if p.requires_grad]},]
        return param_dicts
    
    def get_params(self):
        return [p for n, p in self.named_parameters() if p.requires_grad]
    
    def get_param_names(self):
        return [n for n, p in self.named_parameters() if p.requires_grad]
    
    def grad_clip(self):
        pass
    
class MetaNet(nn.Module):
    def __init__(self, args):
        super(MetaNet, self).__init__()
        # Scaling and translation factors for feature correction
        self.scale = nn.Parameter(torch.zeros(args.batch_size, args.num_class, 1))
        self.bias = nn.Parameter(torch.zeros(args.batch_size, args.num_class, 1))
        
    def forward(self, input):
        # 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
    
    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}]

        