import torch
import torch.nn as nn


class SelfAttentionPooling(nn.Module):
    """
    Implementation of SelfAttentionPooling
    Original Paper: Self-Attention Encoding and Pooling for Speaker Recognition
    https://arxiv.org/pdf/2008.01077v1.pdf
    """

    def __init__(self, input_dim):
        super(SelfAttentionPooling, self).__init__()
        self.W = nn.Linear(input_dim, 1)


    def forward(self, batch_rep):
        """
        input:
            batch_rep : size (N, T, H), N: batch size, T: sequence length, H: Hidden dimension

        attention_weight:
            att_w : size (N, T, 1)

        return:
            utter_rep: size (N, H)
        """
        softmax = nn.functional.softmax
        att_w = softmax(self.W(batch_rep).squeeze(-1)).unsqueeze(-1)
        utter_rep = torch.sum(batch_rep * att_w, dim=1)

        return utter_rep


class Model(nn.Module):
    def __init__(self, input_dim, clipping=False, attention_pooling=False, num_judges=5000, **kwargs):
        super(Model, self).__init__()
        self.mean_net_linear = nn.Linear(input_dim, 1)
        self.mean_net_clipping = clipping
        self.mean_net_pooling = SelfAttentionPooling(input_dim) if attention_pooling else None
        self.bias_net_linear = nn.Linear(input_dim, 1)
        self.bias_net_pooling = SelfAttentionPooling(input_dim) if attention_pooling else None
        self.judge_embbeding = nn.Embedding(num_embeddings = num_judges, embedding_dim=input_dim)

    def forward(self, features, judge_ids=None):
        if self.mean_net_pooling is not None:
            x = self.mean_net_pooling(features)
            segment_score = self.mean_net_linear(x)
        else:
            x = self.mean_net_linear(features)
            segment_score = x.squeeze(-1).mean(dim=-1)
            
        if self.mean_net_clipping:
            segment_score = torch.tanh(segment_score) * 2 + 3
        
        if judge_ids is None:
            return segment_score.squeeze(-1)
        
        else:
            time = features.shape[1]
            judge_features = self.judge_embbeding(judge_ids)
            judge_features = torch.stack([judge_features for i in range(time)], dim = 1)
            bias_features = features + judge_features
            
            if self.bias_net_pooling is not None:
                y = self.bias_net_pooling(bias_features)
                bias_score = self.bias_net_linear(y)
            else:
                y = self.bias_net_linear(bias_features)
                bias_score = y.squeeze(-1).mean(dim=-1)
            bias_score = bias_score + segment_score

        return segment_score.squeeze(-1), bias_score.squeeze(-1)