import numpy as np
import torch
import framework.ops
import t2vretrieval.encoders.mlsent
import t2vretrieval.encoders.mlvideo

import t2vretrieval.models.globalmatch
from t2vretrieval.models.criterion import cosine_sim

from t2vretrieval.models.globalmatch import VISENC, TXTENC

class RoleGraphMatchModelConfig(t2vretrieval.models.globalmatch.GlobalMatchModelConfig):
  def __init__(self):
    super().__init__()
    self.num_verbs = 4
    self.num_nouns = 6
    
    self.attn_fusion = 'embed'  # sim, embed
    self.simattn_sigma = 4

    self.hard_topk = 1
    self.max_violation = True

    self.loss_weights = None
    ## this config will be covered by model.json due to the functions of load and load_from_dict
    self.subcfgs[VISENC] = t2vretrieval.encoders.mlvideo.MultilevelEncoderConfig()
    self.subcfgs[TXTENC] = t2vretrieval.encoders.mlsent.RoleGraphEncoderConfig()


class RoleGraphMatchModel(t2vretrieval.models.globalmatch.GlobalMatchModel):
  def build_submods(self):
    return {
      VISENC: t2vretrieval.encoders.mlvideo.MultilevelEncoder(self.config.subcfgs[VISENC]),
      TXTENC: t2vretrieval.encoders.mlsent.RoleGraphEncoder(self.config.subcfgs[TXTENC])
    }

  def forward_video_embed(self, batch_data):
    vid_fts = torch.FloatTensor(batch_data['attn_fts']).to(self.device)
    vid_lens = torch.LongTensor(batch_data['attn_lens']).to(self.device)
    # (batch, max_vis_len, dim_embed)
    vid_sent_embeds, vid_verb_embeds, vid_noun_embeds, local_sent_embeds, logits, max_len = self.submods[VISENC](vid_fts, vid_lens)
    return {
      'vid_sent_embeds': vid_sent_embeds,
      'vid_verb_embeds': vid_verb_embeds, 
      'vid_noun_embeds': vid_noun_embeds,
      'local_vid_embeds': local_sent_embeds,
      'vid_lens': vid_lens,
      'max_len': max_len,
      'logits': logits,
    }

  def forward_text_embed(self, batch_data):
    sent_ids = torch.LongTensor(batch_data['sent_ids']).to(self.device)  ## sentence
    sent_lens = torch.LongTensor(batch_data['sent_lens']).to(self.device) ## length
    verb_masks = torch.BoolTensor(batch_data['verb_masks']).to(self.device) ## batch*nv*max_sen_len
    noun_masks = torch.BoolTensor(batch_data['noun_masks']).to(self.device)
    node_roles = torch.LongTensor(batch_data['node_roles']).to(self.device) ## batch*(n_v+n_n)
    rel_edges = torch.FloatTensor(batch_data['rel_edges']).to(self.device)  ## batch*(1+n_v+n_n)*(1+n_v+n_n)

    verb_lens = torch.sum(verb_masks, 2)
    noun_lens = torch.sum(noun_masks, 2)

    # sent_embeds: (batch, dim_embed)
    # verb_embeds, noun_embeds: (batch, num_xxx, dim_embed)
    sent_embeds, verb_embeds, noun_embeds, local_sent_embeds, sent_logits = self.submods[TXTENC](
      sent_ids, sent_lens, verb_masks, noun_masks, node_roles, rel_edges)
    return {
      'sent_embeds': sent_embeds, 'sent_lens': sent_lens, 
      'verb_embeds': verb_embeds, 'verb_lens': verb_lens, 
      'noun_embeds': noun_embeds, 'noun_lens': noun_lens,
      'sent_logits': sent_logits, 'local_sent_embeds': local_sent_embeds,
      }

  def generate_phrase_scores(self, vid_embeds, vid_masks, phrase_embeds, phrase_masks, mask_flag=False):
    '''Args:
      - vid_embeds: (batch, num_frames, embed_size)
      - vid_masks: (batch, num_frames)
      - phrase_embeds: (batch, num_phrases, embed_size)
      - phrase_masks: (batch, num_phrases)
    '''
    batch_vids, num_frames, _ = vid_embeds.size()
    vid_pad_masks = (vid_masks == 0).unsqueeze(1).unsqueeze(3)
    batch_phrases, num_phrases, dim_embed = phrase_embeds.size()

    # compute component-wise similarity
    vid_2d_embeds = vid_embeds.view(-1, dim_embed)
    phrase_2d_embeds = phrase_embeds.view(-1, dim_embed)
    # size = (batch_vids, batch_phrases, num_frames, num_phrases)
    ground_sims = cosine_sim(vid_2d_embeds, phrase_2d_embeds).view(
      batch_vids, num_frames, batch_phrases, num_phrases).transpose(1, 2)
    ###
    if mask_flag:
      vid_attn_per_word = ground_sims.masked_fill(vid_pad_masks, 0) ##############
    else:
      vid_attn_per_word = ground_sims
    vid_attn_per_word[vid_attn_per_word < 0] = 0
    vid_attn_per_word = framework.ops.l2norm(vid_attn_per_word, dim=2)
    if mask_flag:
      vid_attn_per_word = vid_attn_per_word.masked_fill(vid_pad_masks, -1e18) #################
    vid_attn_per_word = torch.softmax(self.config.simattn_sigma * vid_attn_per_word, dim=2)

    if self.config.attn_fusion == 'embed':
      vid_attned_embeds = torch.einsum('abcd,ace->abde', vid_attn_per_word, vid_embeds)
      word_attn_sims = torch.einsum('abde,bde->abd',
        framework.ops.l2norm(vid_attned_embeds),
        framework.ops.l2norm(phrase_embeds))
    elif self.config.attn_fusion == 'sim':
      # (batch_vids, batch_phrases, num_phrases)
      word_attn_sims = torch.sum(ground_sims * vid_attn_per_word, dim=2)

    # sum: (batch_vid, batch_phrases)
    phrase_scores = torch.sum(word_attn_sims * phrase_masks.float().unsqueeze(0), 2) \
                   / torch.sum(phrase_masks, 1).float().unsqueeze(0).clamp(min=1)
    return phrase_scores

  def generate_scores(self, **kwargs):
    ##### shared #####
    vid_lens = kwargs['vid_lens'] # (batch, )
    num_frames = int(kwargs['max_len'])###########################kwargs['vid_verb_embeds'].size(1)
    vid_masks = framework.ops.sequence_mask(vid_lens, num_frames, inverse=False)
    # batch*max_len
    ##### sentence-level scores #####
    sent_scores = cosine_sim(kwargs['vid_sent_embeds'], kwargs['sent_embeds'])
    #######################################################
    # concept scores use jaccard similarity
    concept_verb_scores = self.jaccard_sim(kwargs['logits'][0], kwargs['sent_logits'][0])
    concept_noun_scores = self.jaccard_sim(kwargs['logits'][1], kwargs['sent_logits'][1])
    #######################################################
    ##### verb-level scores #####
    vid_verb_embeds = kwargs['vid_verb_embeds'] # (batch, num_frames, dim_embed)
    verb_embeds = kwargs['verb_embeds'] # (batch, num_verbs, dim_embed)
    verb_lens = kwargs['verb_lens'] # (batch, num_verbs)
    local_vid_embeds =kwargs['local_vid_embeds']
    local_sent_embeds = kwargs['local_sent_embeds']
    verb_masks = framework.ops.sequence_mask(torch.sum(verb_lens > 0, 1).long(), 
      self.config.num_verbs, inverse=False)
    # sum: (batch_vids, batch_sents)
    verb_scores = self.generate_phrase_scores(vid_verb_embeds, vid_masks, verb_embeds, verb_masks)
    ind_verb_scores = self.generate_phrase_scores(local_vid_embeds[0], vid_masks, local_sent_embeds[0], verb_masks, True)
    ##### noun-level scores #####
    vid_noun_embeds = kwargs['vid_noun_embeds'] # (batch, num_frames, dim_embed)
    noun_embeds = kwargs['noun_embeds'] # (batch, num_nouns, dim_embed)
    noun_lens = kwargs['noun_lens'] # (batch, num_nouns)
    noun_masks = framework.ops.sequence_mask(torch.sum(noun_lens > 0, 1).long(), 
      self.config.num_nouns, inverse=False)
    # sum: (batch_vids, batch_sents)
    noun_scores = self.generate_phrase_scores(vid_noun_embeds, vid_masks, noun_embeds, noun_masks)
    ind_noun_scores = self.generate_phrase_scores(local_vid_embeds[1], vid_masks, local_sent_embeds[1], noun_masks, True)
    return sent_scores, verb_scores, noun_scores, concept_verb_scores, concept_noun_scores, ind_verb_scores, ind_noun_scores
  
  def jaccard_sim(self, im, s):
    im_bs = im.size(0)
    s_bs = s.size(0)
    im = im.unsqueeze(1).expand(-1, s_bs, -1)
    s = s.unsqueeze(0).expand(im_bs, -1, -1)
    intersection = torch.min(im, s).sum(-1)
    union = torch.max(im, s).sum(-1)
    score = intersection / union
    return score

  def forward_loss(self, batch_data, step=None):
    enc_outs = self.forward_video_embed(batch_data)
    cap_enc_outs = self.forward_text_embed(batch_data)
    enc_outs.update(cap_enc_outs)
    sent_scores, verb_scores, noun_scores, concept_verb_scores, concept_noun_scores, local_verb_scores, local_noun_scores = self.generate_scores(**enc_outs)
    scores = (sent_scores + verb_scores + noun_scores + local_verb_scores + local_noun_scores) / 5
    scores2 = (concept_verb_scores + concept_noun_scores) / 2

    sent_loss = self.criterion(sent_scores)
    verb_loss = self.criterion(verb_scores)
    noun_loss = self.criterion(noun_scores)

    eta = 0.1
    mu = 0.01
    concept_verb_loss = 0.5*self.criterion(concept_verb_scores)
    concept_noun_loss = 0.5*self.criterion(concept_noun_scores)
    concept_loss = eta*self.criterion(scores2)

    verb_concept_label = torch.FloatTensor(batch_data['verb_concept_label']).to(self.device)
    noun_concept_label = torch.FloatTensor(batch_data['noun_concept_label']).to(self.device)
    verb_concept_mask = torch.FloatTensor(batch_data['verb_concept_mask']).to(self.device)
    noun_concept_mask = torch.FloatTensor(batch_data['noun_concept_mask']).to(self.device)

    v_mask_sum = torch.sum(verb_concept_mask, dim=1)
    n_mask_sum = torch.sum(noun_concept_mask, dim=1)

    vbce_loss = torch.sum(verb_concept_mask*self.criterion_bce(enc_outs['logits'][0], verb_concept_label), dim=1)
    vbce_loss = mu*torch.mean(vbce_loss/v_mask_sum)
    nbce_loss = torch.sum(noun_concept_mask*self.criterion_bce(enc_outs['logits'][1], noun_concept_label), dim=1)
    nbce_loss = mu*torch.mean(nbce_loss/n_mask_sum)
    vbce_sent_loss = torch.sum(verb_concept_mask*self.criterion_bce(enc_outs['sent_logits'][0], verb_concept_label), dim=1)
    vbce_sent_loss = mu*torch.mean(vbce_sent_loss/v_mask_sum)
    nbce_sent_loss = torch.sum(noun_concept_mask*self.criterion_bce(enc_outs['sent_logits'][1], noun_concept_label), dim=1)
    nbce_sent_loss = mu*torch.mean(nbce_sent_loss/n_mask_sum)

    fusion_loss = self.criterion(scores) 

    if self.config.loss_weights is None:
      loss = fusion_loss + 1*(vbce_loss+nbce_loss) + 1*(vbce_sent_loss+nbce_sent_loss) + concept_loss
    else:
      loss = self.config.loss_weights[0] * fusion_loss + \
             self.config.loss_weights[1] * sent_loss + \
             self.config.loss_weights[2] * verb_loss + \
             self.config.loss_weights[3] * noun_loss + \
             vbce_loss + nbce_loss
    if step is not None and self.config.monitor_iter > 0 and step % self.config.monitor_iter == 0:
      neg_scores = scores.masked_fill(torch.eye(len(scores), dtype=torch.bool).to(self.device), -1e10)
      self.print_fn('\tstep %d: pos mean scores %.2f, hard neg mean scores i2t %.2f, t2i %.2f'%(
        step, torch.mean(torch.diag(scores)), torch.mean(torch.max(neg_scores, 1)[0]),
        torch.mean(torch.max(neg_scores, 0)[0])))
      self.print_fn('\tstep %d: sent_loss %.4f, verb_loss %.4f, noun_loss %.4f, fusion_loss %.4f'%(
        step, sent_loss.data.item(), verb_loss.data.item(), noun_loss.data.item(), fusion_loss.data.item()))
      self.print_fn('\tstep %d: vbce_loss %.4f, nbce_loss %.4f'%(step, vbce_loss.item(), nbce_loss.item()))
      self.print_fn('\tstep %d: vbce_sent_loss %.4f, nbce_sent_loss %.4f'%(step, vbce_sent_loss.item(), nbce_sent_loss.item()))
      self.print_fn('\tstep %d: sim_loss %.4f, vsim_loss %.4f, nsim_loss %.4f'%(step, concept_loss.item(),
                                                                                concept_verb_loss.item(), concept_noun_loss.item()))

    return loss
      
  def evaluate_scores(self, tst_reader):
    K = self.config.subcfgs[VISENC].num_levels
    K = K + 4
    assert K == 7, 'Note that this error indicates losing other scores!'
    vid_names, all_scores = [], [[] for _ in range(K)]
    cap_names = tst_reader.dataset.captions

    for vid_data in tst_reader:
      vid_names.extend(vid_data['names'])
      vid_enc_outs = self.forward_video_embed(vid_data)
      for k in range(K):
        all_scores[k].append([])
      ijj = 0
      for cap_data in tst_reader.dataset.iterate_over_captions(self.config.tst_batch_size):
        cap_enc_outs = self.forward_text_embed(cap_data)
        cap_enc_outs.update(vid_enc_outs)
        indv_scores = self.generate_scores(**cap_enc_outs)
        for k in range(K):
          all_scores[k][-1].append(indv_scores[k].data.cpu().numpy())
        ijj += 0
      for k in range(K):
        all_scores[k][-1] = np.concatenate(all_scores[k][-1], axis=1)
    for k in range(K):
      all_scores[k] = np.concatenate(all_scores[k], axis=0)  # (n_img, n_cap)
    all_scores = np.array(all_scores)  # (k, n_img, n_cap)

    return vid_names, cap_names, all_scores

  def evaluate(self, tst_reader, return_outs=False):
    vid_names, cap_names, scores = self.evaluate_scores(tst_reader)

    i2t_gts = []
    for vid_name in vid_names:
      i2t_gts.append([])
      for i, cap_name in enumerate(cap_names):
        if cap_name in tst_reader.dataset.ref_captions[vid_name]:
          i2t_gts[-1].append(i)

    t2i_gts = {}
    for i, t_gts in enumerate(i2t_gts):
      for t_gt in t_gts:
        t2i_gts.setdefault(t_gt, [])
        t2i_gts[t_gt].append(i)

    idx = [0, 1, 2, 5, 6]

    fused_scores = (np.mean(scores[idx], 0) + np.mean(scores[3:5], 0))/2
    metrics = self.calculate_metrics(fused_scores, i2t_gts, t2i_gts)

    if return_outs:
      outs = {
        'vid_names': vid_names,
        'cap_names': cap_names,
        'scores': scores,
      }
      return metrics, outs
    else:
      return metrics





[((60, 16, 60, 40), 'torch.sum', 'torch.sum', ({(60, 26, 60, 36): 'verb_masks', (60, 38, 60, 39): '2'}, {}), '(verb_masks, 2)', False, 'import torch\n'), ((61, 16, 61, 40), 'torch.sum', 'torch.sum', ({(61, 26, 61, 36): 'noun_masks', (61, 38, 61, 39): '2'}, {}), '(noun_masks, 2)', False, 'import torch\n'), ((100, 24, 100, 91), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((123, 18, 123, 78), 't2vretrieval.models.criterion.cosine_sim', 'cosine_sim', ({(123, 29, 123, 54): "kwargs['vid_sent_embeds']", (123, 56, 123, 77): "kwargs['sent_embeds']"}, {}), "(kwargs['vid_sent_embeds'], kwargs['sent_embeds'])", False, 'from t2vretrieval.models.criterion import cosine_sim\n'), ((184, 17, 184, 52), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((185, 17, 185, 52), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((244, 17, 244, 37), 'numpy.array', 'np.array', ({(244, 26, 244, 36): 'all_scores'}, {}), '(all_scores)', True, 'import numpy as np\n'), ((103, 26, 103, 87), 'torch.einsum', 'torch.einsum', ({(103, 39, 103, 55): '"""abcd,ace->abde"""', (103, 57, 103, 74): 'vid_attn_per_word', (103, 76, 103, 86): 'vid_embeds'}, {}), "('abcd,ace->abde', vid_attn_per_word, vid_embeds)", False, 'import torch\n'), ((188, 19, 188, 51), 'torch.mean', 'torch.mean', ({(188, 30, 188, 50): '(vbce_loss / v_mask_sum)'}, {}), '(vbce_loss / v_mask_sum)', False, 'import torch\n'), ((190, 19, 190, 51), 'torch.mean', 'torch.mean', ({(190, 30, 190, 50): '(nbce_loss / n_mask_sum)'}, {}), '(nbce_loss / n_mask_sum)', False, 'import torch\n'), ((192, 24, 192, 61), 'torch.mean', 'torch.mean', ({(192, 35, 192, 60): '(vbce_sent_loss / v_mask_sum)'}, {}), '(vbce_sent_loss / v_mask_sum)', False, 'import torch\n'), ((194, 24, 194, 61), 'torch.mean', 'torch.mean', ({(194, 35, 194, 60): '(nbce_sent_loss / n_mask_sum)'}, {}), '(nbce_sent_loss / n_mask_sum)', False, 'import torch\n'), ((243, 22, 243, 59), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((38, 14, 38, 55), 'torch.FloatTensor', 'torch.FloatTensor', ({(38, 32, 38, 54): "batch_data['attn_fts']"}, {}), "(batch_data['attn_fts'])", False, 'import torch\n'), ((39, 15, 39, 56), 'torch.LongTensor', 'torch.LongTensor', ({(39, 32, 39, 55): "batch_data['attn_lens']"}, {}), "(batch_data['attn_lens'])", False, 'import torch\n'), ((53, 15, 53, 55), 'torch.LongTensor', 'torch.LongTensor', ({(53, 32, 53, 54): "batch_data['sent_ids']"}, {}), "(batch_data['sent_ids'])", False, 'import torch\n'), ((54, 16, 54, 57), 'torch.LongTensor', 'torch.LongTensor', ({(54, 33, 54, 56): "batch_data['sent_lens']"}, {}), "(batch_data['sent_lens'])", False, 'import torch\n'), ((55, 17, 55, 59), 'torch.BoolTensor', 'torch.BoolTensor', ({(55, 34, 55, 58): "batch_data['verb_masks']"}, {}), "(batch_data['verb_masks'])", False, 'import torch\n'), ((56, 17, 56, 59), 'torch.BoolTensor', 'torch.BoolTensor', ({(56, 34, 56, 58): "batch_data['noun_masks']"}, {}), "(batch_data['noun_masks'])", False, 'import torch\n'), ((57, 17, 57, 59), 'torch.LongTensor', 'torch.LongTensor', ({(57, 34, 57, 58): "batch_data['node_roles']"}, {}), "(batch_data['node_roles'])", False, 'import torch\n'), ((58, 16, 58, 58), 'torch.FloatTensor', 'torch.FloatTensor', ({(58, 34, 58, 57): "batch_data['rel_edges']"}, {}), "(batch_data['rel_edges'])", False, 'import torch\n'), ((109, 23, 109, 72), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((156, 19, 156, 35), 'torch.min', 'torch.min', ({(156, 29, 156, 31): 'im', (156, 33, 156, 34): 's'}, {}), '(im, s)', False, 'import torch\n'), ((157, 12, 157, 28), 'torch.max', 'torch.max', ({(157, 22, 157, 24): 'im', (157, 26, 157, 27): 's'}, {}), '(im, s)', False, 'import torch\n'), ((179, 25, 179, 76), 'torch.FloatTensor', 'torch.FloatTensor', ({(179, 43, 179, 75): "batch_data['verb_concept_label']"}, {}), "(batch_data['verb_concept_label'])", False, 'import torch\n'), ((180, 25, 180, 76), 'torch.FloatTensor', 'torch.FloatTensor', ({(180, 43, 180, 75): "batch_data['noun_concept_label']"}, {}), "(batch_data['noun_concept_label'])", False, 'import torch\n'), ((181, 24, 181, 74), 'torch.FloatTensor', 'torch.FloatTensor', ({(181, 42, 181, 73): "batch_data['verb_concept_mask']"}, {}), "(batch_data['verb_concept_mask'])", False, 'import torch\n'), ((182, 24, 182, 74), 'torch.FloatTensor', 'torch.FloatTensor', ({(182, 42, 182, 73): "batch_data['noun_concept_mask']"}, {}), "(batch_data['noun_concept_mask'])", False, 'import torch\n'), ((241, 28, 241, 69), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((266, 20, 266, 43), 'numpy.mean', 'np.mean', ({(266, 28, 266, 39): 'scores[idx]', (266, 41, 266, 42): '(0)'}, {}), '(scores[idx], 0)', True, 'import numpy as np\n'), ((266, 46, 266, 69), 'numpy.mean', 'np.mean', ({(266, 54, 266, 65): 'scores[3:5]', (266, 67, 266, 68): '(0)'}, {}), '(scores[3:5], 0)', True, 'import numpy as np\n'), ((135, 45, 135, 72), 'torch.sum', 'torch.sum', ({(135, 55, 135, 68): 'verb_lens > 0', (135, 70, 135, 71): '1'}, {}), '(verb_lens > 0, 1)', False, 'import torch\n'), ((144, 45, 144, 72), 'torch.sum', 'torch.sum', ({(144, 55, 144, 68): 'noun_lens > 0', (144, 70, 144, 71): '1'}, {}), '(noun_lens > 0, 1)', False, 'import torch\n'), ((89, 18, 89, 61), 't2vretrieval.models.criterion.cosine_sim', 'cosine_sim', ({(89, 29, 89, 42): 'vid_2d_embeds', (89, 44, 89, 60): 'phrase_2d_embeds'}, {}), '(vid_2d_embeds, phrase_2d_embeds)', False, 'from t2vretrieval.models.criterion import cosine_sim\n'), ((209, 25, 209, 43), 'torch.diag', 'torch.diag', ({(209, 36, 209, 42): 'scores'}, {}), '(scores)', False, 'import torch\n'), ((209, 57, 209, 81), 'torch.max', 'torch.max', ({(209, 67, 209, 77): 'neg_scores', (209, 79, 209, 80): '(1)'}, {}), '(neg_scores, 1)', False, 'import torch\n'), ((210, 19, 210, 43), 'torch.max', 'torch.max', ({(210, 29, 210, 39): 'neg_scores', (210, 41, 210, 42): '(0)'}, {}), '(neg_scores, 0)', False, 'import torch\n'), ((113, 21, 113, 47), 'torch.sum', 'torch.sum', ({(113, 31, 113, 43): 'phrase_masks', (113, 45, 113, 46): '(1)'}, {}), '(phrase_masks, 1)', False, 'import torch\n')]