import torch
from torch import nn
from GlobalAttention import GlobalAttention
from torch.autograd import Variable
from Beam import TreeBeam
from UtilClass import bottle, unbottle
from preprocess import rhs, CDDataset
from decoders import DecoderState, Prediction
import torch.nn.functional as F

class ConcodeDecoder(nn.Module):

  def __init__(self, vocabs, opt):
    super(ConcodeDecoder, self).__init__()

    self.opt = opt
    self.vocabs = vocabs

    self.nt_embedding = nn.Embedding(
      len(vocabs['nt']),
      opt.tgt_word_vec_size,
      padding_idx=vocabs['nt'].stoi['<blank>'])

    self.rule_embedding = nn.Embedding(
      len(vocabs['prev_rules']),
      opt.tgt_word_vec_size,
      padding_idx=vocabs['prev_rules'].stoi['<blank>'])

    self.attn = GlobalAttention(
        opt.decoder_rnn_size,
        attn_type='general',
        include_rnn=False)

    self.attn_linear = nn.Linear(self.opt.decoder_rnn_size * 3, self.opt.decoder_rnn_size)

    self.var_attn = GlobalAttention(
        opt.decoder_rnn_size,
        attn_type='general',
        include_rnn=False)

    if opt.copy_attn:
      self.copy_attn = GlobalAttention(
          opt.decoder_rnn_size,
          attn_type='general')

    self.decoder_rnn = nn.LSTM(
      input_size=opt.tgt_word_vec_size * 3 + opt.decoder_rnn_size, # nt and prev_rule
      hidden_size=opt.decoder_rnn_size,
      num_layers=opt.dec_layers,
      dropout=opt.dropout,
      batch_first=True)

    self.decoder_dropout = nn.Dropout(opt.dropout)

  def forward(self, batch, all_context, context_masks, decState):

    src_context = all_context[0]
    src_context_mask = context_masks[0]
    rest_context = torch.cat(all_context[1:], 1)
    rest_context_mask = torch.cat(context_masks[1:], 1)

    context = torch.cat(all_context, 1)
    context_lengths = torch.cat(context_masks, 1)

    # embed everything
    nt_embeddings = self.nt_embedding(Variable(batch['nt'].cuda(), requires_grad=False))
    rule_embeddings = self.rule_embedding(Variable(batch['prev_rules'].cuda(), requires_grad=False))
    parent_rule_embeddings = self.rule_embedding(Variable(batch['parent_rules'].cuda(), requires_grad=False))

    attn_outputs, attn_scores, copy_attn_scores = [], [], []
    # For each batch we have to maintain states

    batch_size = batch['nt'].size(0) # 1 for predict
    num_decodes = 0

    attn_outputs, attn_scores, copy_attn_scores = [], [], []
    for i, (nt, rule, parent_rule) in enumerate(zip(nt_embeddings.split(1, 1), rule_embeddings.split(1, 1), parent_rule_embeddings.split(1, 1))):
      # accumulate parent decoder states
      parent_states = []
      for j in range(0, batch_size):
        try: # this is needed coz the batch is of different sizes
          parent_states.append(batch['parent_states'][j][i]) # one state for every batch
        except:
          parent_states.append(batch['parent_states'][j][0]) # one state for every batch
      parent_states = torch.cat(parent_states, 0)

      rnn_output, prev_hidden = self.decoder_rnn(torch.cat((nt, rule, parent_rule, parent_states), 2), decState.hidden)
      num_decodes += 1
      rnn_output.contiguous()

      if self.opt.twostep:
        src_attn_output, src_attn_score = self.attn(rnn_output, src_context, src_context_mask)
        varmet_attn_output, varmet_attn_score = self.var_attn(src_attn_output, rest_context, rest_context_mask)

        attn_output = F.tanh(self.attn_linear(torch.cat((rnn_output, src_attn_output, varmet_attn_output), 2)))
        attn_scores.append(varmet_attn_score)
        copy_attn_scores.append(varmet_attn_score)
      else:
        attn_output, attn_score = self.attn(rnn_output, context, context_lengths)
        # attn_score is b x tgt_len x src_len, src should be removed from this
        attn_scores.append(attn_score[:, :, src_context.size(1):])
        copy_attn_scores.append(attn_score[:, :, src_context.size(1):])


      attn_output = self.decoder_dropout(attn_output)
      attn_outputs.append(attn_output)

      decState.update_state(prev_hidden, attn_output)

      # update all children
      for j, elem in enumerate(rnn_output.split(1, 0)):
        # children wont be there during prediction
        if 'children' in batch and i in batch['children'][j]: # rule i has children
          for child in batch['children'][j][i]:
            batch['parent_states'][j][child] = elem

    output = torch.cat(attn_outputs, 1)
    attn_scores = torch.cat(attn_scores, 1)
    copy_attn_scores = torch.cat(copy_attn_scores, 1) if self.opt.copy_attn else None

    return output, attn_scores, copy_attn_scores

  def predict(self, enc_hidden, context, context_lengths, batch, beam_size, max_code_length, generator, replace_unk, vis_params):

    # This decoder does not have input feeding. Parent state replces that
    decState = DecoderState(
      enc_hidden, #encoder hidden
      Variable(torch.zeros(1, 1, self.opt.decoder_rnn_size).cuda(), requires_grad=False) # parent state
    )
    # Repeat everything beam_size times.
    def rvar(a, beam_size):
      return Variable(a.repeat(beam_size, 1, 1), volatile=True)

    context = tuple(rvar(context[i].data, beam_size) for i in range(0, len(context)))
    context_lengths = tuple(context_lengths[i].repeat(beam_size, 1) for i in range(0, len(context_lengths)))

    decState.repeat_beam_size_times(beam_size)

    # Use only one beam
    beam = TreeBeam(beam_size, True, self.vocabs, self.opt.decoder_rnn_size)

    for count in range(0, max_code_length): # We will break when we have the required number of terminals
        # to be consistent with seq2seq

        if beam.done():
          break

        # Construct batch x beam_size nxt words.
        # Get all the pending current beam words and arrange for forward.
        # Uses the start symbol in the beginning
        inp = beam.getCurrentState() # Should return a batch of the frontier

        # Run one step., decState gets automatically updated
        output, attn, copy_attn = self.forward(inp, context, context_lengths, decState)
        src_map = torch.zeros(0, 0)
        if self.opt.var_names:
          src_map = torch.cat((src_map, batch['concode_src_map_vars']), 1)
        if self.opt.method_names:
          src_map = torch.cat((src_map, batch['concode_src_map_methods']), 1)

        scores = generator(bottle(output), bottle(copy_attn), src_map, inp) #generator needs the non-terminals

        out = generator.collapseCopyScores(unbottle(scores.data.clone(), beam_size), batch) # needs seq2seq from batch
        out = out.log()

        # beam x tgt_vocab

        beam.advance(out[:, 0],  attn.data[:, 0], output)
        decState.beam_update(beam.getCurrentOrigin(), beam_size)

    pred_score_total = 0
    pred_words_total = 0

    score, times, k = beam.getFinal() # times is the length of the prediction
    hyp, att = beam.getHyp(times, k)
    goldNl = []
    if self.opt.var_names:
      goldNl += batch['concode_var'][0] # because batch = 1
    if self.opt.method_names:
      goldNl += batch['concode_method'][0] # because batch = 1

    goldCode = self.vocabs['code'].addStartOrEnd(batch['raw_code'][0])
    predSent, copied_tokens, replaced_tokens = self.buildTargetTokens(
      hyp,
      self.vocabs,
      goldNl,
      att,
      batch['concode_vocab'][0],
      replace_unk
    )
    predSent = ConcodeDecoder.rulesToCode(predSent)
    pred_score_total += score
    pred_words_total += len(predSent)

    return Prediction(goldNl, goldCode, predSent, att)

  @staticmethod
  def rulesToCode(rules):
    stack = []
    code = []
    for i in range(0, len(rules)):
      if not CDDataset._is_terminal_rule(rules[i]):
        stack.extend(rhs(rules[i]).split('___')[::-1])
      else:
        code.append(rhs(rules[i]))

      try:
        top = stack.pop()

        while not top[0].isupper():
          code.append(top)
          if len(stack) == 0:
            break
          top = stack.pop()
      except:
        pass

    return code

  def buildTargetTokens(self, pred, vocabs, src, attn, copy_vocab, replace_unk):
      vocab = vocabs['next_rules']
      tokens = []
      copied_tokens, replaced_tokens = [], []
      for tok in pred:
          if tok < len(vocab):
              tokens.append(vocab.itos[tok])
          else:
              tokens.append("IdentifierNT-->" + copy_vocab.itos[tok - len(vocab)])
              copied_tokens.append(copy_vocab.itos[tok - len(vocab)])

      if replace_unk and attn is not None:
          for i in range(len(tokens)):
              if tokens[i] == '<unk>':
                  _, maxIndex = attn[i].max(0)
                  tokens[i] = "IdentifierNT-->" + src[maxIndex[0]]
                  replaced_tokens.append(src[maxIndex[0]])

      return tokens, copied_tokens, replaced_tokens