import torch
import torch.nn as nn
import torch.nn.functional as F
from model.emb import EMB
from open_vocabulary.util.consistency_util import attention_sup
from open_vocabulary.util.consistency_util import AttentionGate
from open_vocabulary.util.consistency_util import relevance_scores as rs
from open_vocabulary.util.consistency_util import TokenedSelfAttention
from open_vocabulary.util.TextRandomMask import TextDropout
from open_vocabulary.model.transformer import TransformerDecoder, TransformerEncoder, MLP, MultiheadAttention
# The final model used
class Parallel_TSGV(EMB):
    def __init__(self, configs, word_vectors):
        super().__init__(configs, word_vectors)
        self.parallel_num = configs.parallel_num
        self.layers_encoder = nn.ModuleList(TransformerEncoder(configs.dim, num_layers=2, FFN_dim=configs.dim // 2, dropout=(0.8-0.2) / (self.parallel_num - 2) * i if i!=0 else 0 + 0.2) for i in range(self.parallel_num-1))
        self.attngate = AttentionGate(configs.dim, 8, 0.1) if configs.use_attngate == 1 else lambda x,y,z,h: x
        self.text_mask = TextDropout(0.1)
        self.relation_scores = rs(input_dim=configs.dim, dropout_p=0.1)

        self.weight_attn = TransformerDecoder(configs.dim, num_layers=2, FFN_dim=configs.dim // 2, dropout=0.1)
        self.weight_predict = MLP(configs.dim, hidden_dim=configs.dim // 2, output_dim=1, num_layers=2, dropout=0.1)
        self.tokens = [torch.nn.Parameter(torch.randn(configs.dim)) for _ in range(self.parallel_num)]
        self.hidden_dim = configs.dim
        self.init_parameters()
    
    def position_embedding(self, video_features, query_features):
        video_features, query_features = self.position(video_features), self.position(query_features)
        return video_features, query_features

    def multimodel_embedding(self, word_ids, char_ids, video_features):
        video_features = self.video_affine(video_features)
        query_features = self.embedding_net(word_ids, char_ids)
        return video_features, query_features
    
    def emb_train(self, video_features, v_mask, query_features, q_mask, moments):

        video_features = self.attngate(video_features, query_features, v_mask, q_mask)

        ((video_features, vquery_features), 
         (proposal_features, pquery_features), 
         proposals, p_mask) = self.encoder(vfeats=video_features, 
                            qfeats=query_features, vmask=v_mask, qmask=q_mask,
                            moments=moments)
        # conduct frame-level regression
        features = self.fuser1(vfeats=video_features, qfeats=vquery_features,
                            vmask=v_mask, qmask=q_mask)
        proposal_free_features = features
        relation_scores = self.relation_scores(proposal_free_features)
        
        h_score = self.highlight_layer(features, v_mask)  
        features = features * h_score.unsqueeze(2)
        start_logits, end_logits = self.predictor(features, mask=v_mask) 
        # conduct proposal level ranking
        features = self.fuser2(vfeats=proposal_features, qfeats=pquery_features,
                            vmask=p_mask, qmask=q_mask)
        # convert proposal's feature and mask back to 2d for ranking
        p_score = self.ranker(features, p_mask)
        # flatten proposal's bounds, mask and scores as outputs
        return h_score, start_logits, end_logits, proposals, p_score, p_mask, proposal_free_features, relation_scores

    def one_train(self, video_features, v_mask, query_features, q_mask, moments):
        video_features, query_features = self.position_embedding(video_features, query_features)
        h_score, start_logits, end_logits, proposals, p_score, p_mask, fuser, fuser_rs = self.emb_train(video_features, v_mask, query_features, q_mask, moments)

        return h_score, start_logits, end_logits, proposals, p_score, p_mask, \
            fuser, fuser_rs
    
    def mask_train(self, video_features, v_mask, query_features, q_mask, moments):
        h_score, start_logits, end_logits, proposals, p_score, p_mask, fuser, fuser_rs = self.one_train(video_features, v_mask, query_features, q_mask, moments)
        mask_query = self.text_mask(query_features)
        _, _, _, _, _, _, mask_fuser, mask_fuser_rs = self.one_train(video_features, v_mask, mask_query, q_mask, moments)
        return h_score, start_logits, end_logits, proposals, p_score, p_mask, \
            fuser, fuser_rs, mask_fuser, mask_fuser_rs

    def mlp_fuse(self, inputs: list, idx: int):
        assert len(inputs) == self.parallel_num
        input = torch.stack(inputs, dim=-1)
        out = self.output_fusers[idx](input)
        out = out.squeeze(-1)
        # out = self.norm(out)
        return out
    
    def token_weights_MLP(self, video, querys, v_mask, q_mask):
        weights = []
        for i in range(self.parallel_num):
            q = querys[i]
            token = self.tokens[i].reshape([1, 1, self.hidden_dim]).repeat(q.shape[0], 1, 1).to(q.device)
            q = torch.cat([token, q], dim=1)
            q = self.weight_predict(q)
            weights.append(q[:, 0, :])
        weights = torch.cat(weights, dim=-1)
        weights = F.softmax(weights, dim=-1)
        return weights

    # Using MLP to predict weights
    def weights_compute(self, video, querys, v_mask, q_mask):
        weights = None
        weights = self.token_weights_MLP(video, querys, v_mask, q_mask)
        # weights = self.attn_weights(video, querys, v_mask, q_mask)
        assert weights is not None
        return weights

    def forward(self, word_ids, char_ids, video_features, v_mask, q_mask, moments=None):
        video_features, query_features = self.multimodel_embedding(word_ids, char_ids, video_features)
        
        querys = []
        querys.append(query_features)
        
        # Compute querys
        # ...


        # Using MLP to predict weights
        weights = self.weights_compute(video_features, querys, v_mask, q_mask)

        h_score, start_logits, end_logits, proposals, p_score, p_mask, \
            fuser1, enhance_fuser, fuser1_rs, enhance_fuser_rs = 0, 0, 0, 0, 0, 0, 0, 0, 0, 0


        # Layer querys learning
        h_score, start_logits, end_logits, proposals, p_score, p_mask, \
            fuser1, fuser1_rs, enhance_fuser, enhance_fuser_rs= self.maskrain(video_features, v_mask, query, q_mask, moments)

        return h_score, start_logits, end_logits, proposals, p_score, p_mask, \
            fuser1, fuser1_rs, enhance_fuser, enhance_fuser_rs
    