# Code Taken from https://github.com/LYH-YF/MWPToolkit
# -*- encoding: utf-8 -*-
# @Author: Yihuai Lan
# @Time: 2021/08/21 04:59:55
# @File: sausolver.py

import random
import torch
from torch import nn
import copy

from module.Encoder.rnn_encoder import BasicRNNEncoder
from module.Embedder.basic_embedder import BasicEmbedder
from module.Decoder.tree_decoder import SARTreeDecoder
from module.Layer.tree_layers import NodeGenerater, SubTreeMerger, TreeNode, TreeEmbedding
from module.Layer.tree_layers import Prediction, GenerateNode, Merge, SemanticAlignmentModule
from module.Strategy.beam_search import TreeBeam
from loss.masked_cross_entropy_loss import MaskedCrossEntropyLoss, masked_cross_entropy
from loss.mse_loss import MSELoss
from utils.utils import copy_list
from utils.enum_type import NumMask, SpecialTokens


class SAUSolver(nn.Module):
    """
    Reference:
        Qin et al. "Semantically-Aligned Universal Tree-Structured Solver for Math Word Problems" in EMNLP 2020.
    """

    def __init__(self, config, dataset):
        super(SAUSolver, self).__init__()
        # parameter
        self.hidden_size = config["hidden_size"]
        self.device = config["device"]
        self.USE_CUDA = True if self.device == torch.device('cuda') else False
        self.beam_size = config['beam_size']
        self.max_out_len = config['max_output_len']
        self.embedding_size = config["embedding_size"]
        self.dropout_ratio = config["dropout_ratio"]
        self.num_layers = config["num_layers"]
        self.rnn_cell_type = config["rnn_cell_type"]
        self.loss_weight = config['loss_weight']

        self.vocab_size = len(dataset.in_idx2word)
        self.out_symbol2idx = dataset.out_symbol2idx
        self.out_idx2symbol = dataset.out_idx2symbol
        generate_list = dataset.generate_list
        self.generate_nums = [self.out_symbol2idx[symbol] for symbol in generate_list]
        self.mask_list = NumMask.number
        self.num_start = dataset.num_start
        self.operator_nums = dataset.operator_nums
        self.generate_size = len(generate_list)

        self.unk_token = self.out_symbol2idx[SpecialTokens.UNK_TOKEN]
        try:
            self.out_sos_token = self.out_symbol2idx[SpecialTokens.SOS_TOKEN]
        except:
            self.out_sos_token = None
        try:
            self.out_eos_token = self.out_symbol2idx[SpecialTokens.EOS_TOKEN]
        except:
            self.out_eos_token = None
        try:
            self.out_pad_token = self.out_symbol2idx[SpecialTokens.PAD_TOKEN]
        except:
            self.out_pad_token = None
        # module
        self.embedder = BasicEmbedder(self.vocab_size, self.embedding_size, self.dropout_ratio)
        # self.t_encoder = BasicRNNEncoder(self.embedding_size, self.hidden_size, self.num_layers, self.rnn_cell_type, self.dropout_ratio)
        self.encoder = BasicRNNEncoder(self.embedding_size, self.hidden_size, self.num_layers, self.rnn_cell_type,
                                       self.dropout_ratio, batch_first=False)
        #self.decoder = SARTreeDecoder(self.hidden_size, self.operator_nums, self.generate_size, self.dropout_ratio)
        self.decoder = Prediction(self.hidden_size,self.operator_nums,self.generate_size,self.dropout_ratio)
        self.node_generater = GenerateNode(self.hidden_size, self.operator_nums, self.embedding_size,
                                           self.dropout_ratio)
        self.merge = Merge(self.hidden_size, self.embedding_size, self.dropout_ratio)
        self.sa = SemanticAlignmentModule(self.hidden_size,self.hidden_size,self.hidden_size)
        self.loss1 = MaskedCrossEntropyLoss()
        #

    def calculate_loss(self, batch_data:dict) -> float:
        """Finish forward-propagating, calculating loss and back-propagation.

        :param batch_data: one batch data.
        :return: loss value.

        batch_data should include keywords 'question', 'ques len', 'equation', 'equ len',
        'num stack', 'num size', 'num pos'
        """
        seq = torch.tensor(batch_data["question"]).to(self.device)
        seq_length = torch.tensor(batch_data["ques len"]).long()
        target = torch.tensor(batch_data["equation"]).to(self.device)
        target_length = torch.LongTensor(batch_data["equ len"]).to(self.device)
        nums_stack = copy.deepcopy(batch_data["num stack"])
        num_size = batch_data["num size"]
        num_pos = batch_data["num pos"]
        generate_nums = self.generate_nums
        num_start = self.num_start
        # sequence mask for attention
        unk = self.unk_token

        loss = self.train_tree(seq, seq_length, target, target_length, nums_stack, num_size, generate_nums, num_pos, unk, num_start)
        return loss

    def model_test(self, batch_data:dict) -> tuple:
        """Model test.

        :param batch_data: one batch data.
        :return: predicted equation, target equation.
        batch_data should include keywords 'question', 'ques len', 'equation',
        'num stack', 'num pos', 'num list'
        """
        seq = torch.tensor(batch_data["question"]).to(self.device)
        seq_length = torch.tensor(batch_data["ques len"]).long()
        target = torch.tensor(batch_data["equation"]).to(self.device)
        nums_stack = copy.deepcopy(batch_data["num stack"])
        num_pos = batch_data["num pos"]
        num_list = batch_data['num list']
        generate_nums = self.generate_nums
        num_start = self.num_start
        # sequence mask for attention
        all_node_output = self.evaluate_tree(seq, seq_length, generate_nums, num_pos, num_start, self.beam_size,
                                             self.max_out_len)

        all_output = self.convert_idx2symbol(all_node_output, num_list[0], copy_list(nums_stack[0]))
        targets = self.convert_idx2symbol(target[0], num_list[0], copy_list(nums_stack[0]))
        return all_output, targets

    def train_tree(self,input_batch, input_length, target_batch, target_length, nums_stack_batch, num_size_batch, generate_nums, num_pos, unk, num_start, 
                                                   english=False,var_nums=[], batch_first=False):
        # sequence mask for attention
        seq_mask = []
        max_len = max(input_length)
        for i in input_length:
            seq_mask.append([0 for _ in range(i)] + [1 for _ in range(i, max_len)])
        seq_mask = torch.ByteTensor(seq_mask)

        num_mask = []
        max_num_size = max(num_size_batch) + len(generate_nums) + len(var_nums)  # 最大的位置列表数目+常识数字数目+未知数列表
        for i in num_size_batch:
            d = i + len(generate_nums) + len(var_nums)
            num_mask.append([0] * d + [1] * (max_num_size - d))
        num_mask = torch.ByteTensor(num_mask)  # 用于屏蔽无关数字，防止生成错误的Nx

        #unk = output_lang.word2index["UNK"]

        # Turn padded arrays into (batch_size x max_len) tensors, transpose into (max_len x batch_size)
        input_var = input_batch.transpose(0, 1)
        target = target_batch.transpose(0, 1)

        padding_hidden = torch.FloatTensor([0.0 for _ in range(self.decoder.hidden_size)]).unsqueeze(0)
        batch_size = len(input_length)

        if self.USE_CUDA:
            input_var = input_var.cuda()
            seq_mask = seq_mask.cuda()
            padding_hidden = padding_hidden.cuda()
            num_mask = num_mask.cuda()

        # Zero gradients of both optimizers
        # Run words through encoder

        #encoder_outputs, problem_output = self.encoder(input_var, input_length)
        seq_emb = self.embedder(input_var)
        pade_outputs, _ = self.encoder(seq_emb, input_length)
        problem_output = pade_outputs[-1, :, :self.hidden_size] + pade_outputs[0, :, self.hidden_size:]
        encoder_outputs = pade_outputs[:, :, :self.hidden_size] + pade_outputs[:, :, self.hidden_size:]
        # Prepare input and output variables
        node_stacks = [[TreeNode(_)] for _ in problem_output.split(1, dim=0)]  # root embedding B x 1

        max_target_length = max(target_length)

        all_node_outputs = []
        all_sa_outputs = []
        # all_leafs = []

        copy_num_len = [len(_) for _ in num_pos]
        num_size = max(copy_num_len)
        # 提取与问题相关的数字embedding
        all_nums_encoder_outputs = self.get_all_number_encoder_outputs(encoder_outputs, num_pos, batch_size, num_size,
                                                                  self.encoder.hidden_size)

        embeddings_stacks = [[] for _ in range(batch_size)]  # B x 1  当前的tree state/ subtree embedding / output
        left_childs = [None for _ in range(batch_size)]  # B x 1

        for t in range(max_target_length):
            num_score, op, current_embeddings, current_context, current_nums_embeddings = self.decoder(
                node_stacks, left_childs, encoder_outputs, all_nums_encoder_outputs, padding_hidden, seq_mask, num_mask)

            # all_leafs.append(p_leaf)
            outputs = torch.cat((op, num_score), 1)
            all_node_outputs.append(outputs)

            target_t, generate_input = self.generate_tree_input(target[t].tolist(), outputs, nums_stack_batch, num_start,
                                                           unk)
            target[t] = target_t
            if self.USE_CUDA:
                generate_input = generate_input.cuda()
            left_child, right_child, node_label = self.node_generater(current_embeddings, generate_input, current_context)
            left_childs = []
            for idx, l, r, node_stack, i, o in zip(range(batch_size), left_child.split(1), right_child.split(1),
                                                   node_stacks, target[t].tolist(), embeddings_stacks):
                if len(node_stack) != 0:
                    node = node_stack.pop()
                else:
                    left_childs.append(None)
                    continue

                # 未知数当数字处理，SEP当操作符处理
                if i < num_start:  # 非数字
                    node_stack.append(TreeNode(r))
                    node_stack.append(TreeNode(l, left_flag=True))
                    o.append(TreeEmbedding(node_label[idx].unsqueeze(0), terminal=False))
                    # print(o[-1].embedding.size())
                    # print(encoder_outputs[idx].size())
                else:  # 数字
                    current_num = current_nums_embeddings[idx, i - num_start].unsqueeze(0)
                    while len(o) > 0 and o[-1].terminal:
                        sub_stree = o.pop()
                        op = o.pop()
                        current_num = self.merge(op.embedding, sub_stree.embedding, current_num)  # Subtree embedding
                        if batch_first:
                            encoder_mapping, decoder_mapping = self.sa(current_num, encoder_outputs[idx])
                        else:
                            temp_encoder_outputs = encoder_outputs.transpose(0, 1)
                            encoder_mapping, decoder_mapping = self.sa(current_num,temp_encoder_outputs[idx])
                        all_sa_outputs.append((encoder_mapping, decoder_mapping))
                    o.append(TreeEmbedding(current_num, terminal=True))

                if len(o) > 0 and o[-1].terminal:
                    left_childs.append(o[-1].embedding)

                else:
                    left_childs.append(None)

        # all_leafs = torch.stack(all_leafs, dim=1)  # B x S x 2
        all_node_outputs = torch.stack(all_node_outputs, dim=1)  # B x S x N

        target = target.transpose(0, 1).contiguous()  # B x S

        if self.USE_CUDA:
            # all_leafs = all_leafs.cuda()
            all_node_outputs = all_node_outputs.cuda()
            target = target.cuda()
            new_all_sa_outputs = []
            for sa_pair in all_sa_outputs:
                new_all_sa_outputs.append((sa_pair[0].cuda(), sa_pair[1].cuda()))
            all_sa_outputs = new_all_sa_outputs
            # target_length = torch.LongTensor(target_length).cuda()
        else:
            pass
            # target_length = torch.LongTensor(target_length)

        semantic_alignment_loss = nn.MSELoss()
        total_semanti_alognment_loss = 0
        sa_len = len(all_sa_outputs)
        for sa_pair in all_sa_outputs:
            total_semanti_alognment_loss += semantic_alignment_loss(sa_pair[0], sa_pair[1])
        # print(total_semanti_alognment_loss)
        total_semanti_alognment_loss = total_semanti_alognment_loss / sa_len
        # print(total_semanti_alognment_loss)

        # op_target = target < num_start
        # loss_0 = masked_cross_entropy_without_logit(all_leafs, op_target.long(), target_length)
        loss = masked_cross_entropy(all_node_outputs, target,target_length) + 0.01 * total_semanti_alognment_loss
        # loss = loss_0 + loss_1
        loss.backward()
        # clip the grad
        # torch.nn.utils.clip_grad_norm_(encoder.parameters(), 5)
        # torch.nn.utils.clip_grad_norm_(predict.parameters(), 5)
        # torch.nn.utils.clip_grad_norm_(generate.parameters(), 5)

        # Update parameters with optimizers
        return loss.item()  # , loss_0.item(), loss_1.item()

    def evaluate_tree(self, input_batch, input_length, generate_nums, num_pos, num_start, beam_size=5, max_length=30):
        seq_mask = torch.BoolTensor(1, input_length).fill_(0)
        # Turn padded arrays into (batch_size x max_len) tensors, transpose into (max_len x batch_size)
        input_var = input_batch.transpose(0, 1)

        num_mask = torch.BoolTensor(1, len(num_pos[0]) + len(generate_nums)).fill_(0)

        padding_hidden = torch.FloatTensor([0.0 for _ in range(self.hidden_size)]).unsqueeze(0)

        batch_size = 1

        if self.USE_CUDA:
            input_var = input_var.cuda()
            seq_mask = seq_mask.cuda()
            padding_hidden = padding_hidden.cuda()
            num_mask = num_mask.cuda()
        # Run words through encoder

        seq_emb = self.embedder(input_var)
        pade_outputs, _ = self.encoder(seq_emb, input_length)
        problem_output = pade_outputs[-1, :, :self.hidden_size] + pade_outputs[0, :, self.hidden_size:]
        encoder_outputs = pade_outputs[:, :, :self.hidden_size] + pade_outputs[:, :, self.hidden_size:]

        # Prepare input and output variables
        node_stacks = [[TreeNode(_)] for _ in problem_output.split(1, dim=0)]

        num_size = len(num_pos[0])
        all_nums_encoder_outputs = self.get_all_number_encoder_outputs(encoder_outputs, num_pos, batch_size, num_size,
                                                                       self.hidden_size)
        # B x P x N
        embeddings_stacks = [[] for _ in range(batch_size)]
        left_childs = [None for _ in range(batch_size)]

        beams = [TreeBeam(0.0, node_stacks, embeddings_stacks, left_childs, [])]

        for t in range(max_length):
            current_beams = []
            while len(beams) > 0:
                b = beams.pop()
                if len(b.node_stack[0]) == 0:
                    current_beams.append(b)
                    continue
                # left_childs = torch.stack(b.left_childs)
                left_childs = b.left_childs

                num_score, op, current_embeddings, current_context, current_nums_embeddings = self.decoder(b.node_stack,
                                                                                                           left_childs,
                                                                                                           encoder_outputs,
                                                                                                           all_nums_encoder_outputs,
                                                                                                           padding_hidden,
                                                                                                           seq_mask,
                                                                                                           num_mask)

                out_score = nn.functional.log_softmax(torch.cat((op, num_score), dim=1), dim=1)

                # out_score = p_leaf * out_score

                topv, topi = out_score.topk(beam_size)

                for tv, ti in zip(topv.split(1, dim=1), topi.split(1, dim=1)):
                    current_node_stack = copy_list(b.node_stack)
                    current_left_childs = []
                    current_embeddings_stacks = copy_list(b.embedding_stack)
                    current_out = copy.deepcopy(b.out)
                    
                    out_token = int(ti)
                    current_out.append(out_token)

                    node = current_node_stack[0].pop()

                    if out_token < num_start:
                        generate_input = torch.LongTensor([out_token])
                        if self.USE_CUDA:
                            generate_input = generate_input.cuda()
                        left_child, right_child, node_label = self.node_generater(current_embeddings, generate_input,
                                                                                  current_context)

                        current_node_stack[0].append(TreeNode(right_child))
                        current_node_stack[0].append(TreeNode(left_child, left_flag=True))

                        current_embeddings_stacks[0].append(TreeEmbedding(node_label[0].unsqueeze(0), False))
                    else:
                        current_num = current_nums_embeddings[0, out_token - num_start].unsqueeze(0)

                        while len(current_embeddings_stacks[0]) > 0 and current_embeddings_stacks[0][-1].terminal:
                            sub_stree = current_embeddings_stacks[0].pop()
                            op = current_embeddings_stacks[0].pop()
                            current_num = self.merge(op.embedding, sub_stree.embedding, current_num)
                        current_embeddings_stacks[0].append(TreeEmbedding(current_num, True))
                    if len(current_embeddings_stacks[0]) > 0 and current_embeddings_stacks[0][-1].terminal:
                        current_left_childs.append(current_embeddings_stacks[0][-1].embedding)
                    else:
                        current_left_childs.append(None)
                    current_beams.append(TreeBeam(b.score + float(tv), current_node_stack, current_embeddings_stacks,
                                                  current_left_childs, current_out))
            beams = sorted(current_beams, key=lambda x: x.score, reverse=True)
            beams = beams[:beam_size]
            flag = True
            for b in beams:
                if len(b.node_stack[0]) != 0:
                    flag = False
            if flag:
                break

        return beams[0].out

    def get_all_number_encoder_outputs(self, encoder_outputs, num_pos, batch_size, num_size, hidden_size):
        indices = list()
        sen_len = encoder_outputs.size(0)
        masked_index = []
        temp_1 = [1 for _ in range(hidden_size)]
        temp_0 = [0 for _ in range(hidden_size)]
        for b in range(batch_size):
            for i in num_pos[b]:
                indices.append(i + b * sen_len)
                masked_index.append(temp_0)
            indices += [0 for _ in range(len(num_pos[b]), num_size)]
            masked_index += [temp_1 for _ in range(len(num_pos[b]), num_size)]
        indices = torch.LongTensor(indices)
        masked_index = torch.BoolTensor(masked_index)
        masked_index = masked_index.view(batch_size, num_size, hidden_size)
        if self.USE_CUDA:
            indices = indices.cuda()
            masked_index = masked_index.cuda()
        all_outputs = encoder_outputs.transpose(0, 1).contiguous()
        all_embedding = all_outputs.view(-1, encoder_outputs.size(2))  # S x B x H -> (B x S) x H
        all_num = all_embedding.index_select(0, indices)
        all_num = all_num.view(batch_size, num_size, hidden_size)
        return all_num.masked_fill_(masked_index, 0.0)

    def generate_tree_input(self, target, decoder_output, nums_stack_batch, num_start, unk):
        # when the decoder input is copied num but the num has two pos, chose the max
        target_input = copy.deepcopy(target)
        for i in range(len(target)):
            if target[i] == unk:
                num_stack = nums_stack_batch[i].pop()
                max_score = -float("1e12")
                for num in num_stack:
                    if decoder_output[i, num_start + num] > max_score:
                        target[i] = num + num_start
                        max_score = decoder_output[i, num_start + num]
            if target_input[i] >= num_start:
                target_input[i] = 0
        return torch.LongTensor(target), torch.LongTensor(target_input)

    def mse_loss(self, outputs, targets, mask=None):
        # outputs   : [batch_size,output_len,hidden_size]
        # targets   : [batch_size,output_len,hidden_size]
        # mask      : [batch_size,output_len]
        mask = mask.to(self.device)
        x = torch.sqrt(torch.sum(torch.square((outputs - targets)), dim=-1))  # [batch_size,output_len]
        y = torch.sum(x * mask, dim=-1) / torch.sum(mask, dim=-1)  # [batch_size]
        return torch.sum(y)

    def convert_idx2symbol(self, output, num_list, num_stack):
        # batch_size=output.size(0)
        '''batch_size=1'''
        seq_len = len(output)
        num_len = len(num_list)
        output_list = []
        res = []
        for s_i in range(seq_len):
            idx = output[s_i]
            if idx in [self.out_sos_token, self.out_eos_token, self.out_pad_token]:
                break
            symbol = self.out_idx2symbol[idx]
            if "NUM" in symbol:
                num_idx = self.mask_list.index(symbol)
                if num_idx >= num_len:
                    res = []
                    break
                res.append(num_list[num_idx])
            elif symbol == SpecialTokens.UNK_TOKEN:
                try:
                    pos_list = num_stack.pop()
                    c = num_list[pos_list[0]]
                    res.append(c)
                except:
                    return None
            else:
                res.append(symbol)
        output_list.append(res)
        return output_list

[((68, 24, 68, 95), 'module.Embedder.basic_embedder.BasicEmbedder', 'BasicEmbedder', ({(68, 38, 68, 53): 'self.vocab_size', (68, 55, 68, 74): 'self.embedding_size', (68, 76, 68, 94): 'self.dropout_ratio'}, {}), '(self.vocab_size, self.embedding_size, self.dropout_ratio)', False, 'from module.Embedder.basic_embedder import BasicEmbedder\n'), ((70, 23, 71, 77), 'module.Encoder.rnn_encoder.BasicRNNEncoder', 'BasicRNNEncoder', (), '', False, 'from module.Encoder.rnn_encoder import BasicRNNEncoder\n'), ((73, 23, 73, 108), 'module.Layer.tree_layers.Prediction', 'Prediction', ({(73, 34, 73, 50): 'self.hidden_size', (73, 51, 73, 69): 'self.operator_nums', (73, 70, 73, 88): 'self.generate_size', (73, 89, 73, 107): 'self.dropout_ratio'}, {}), '(self.hidden_size, self.operator_nums, self.generate_size, self.\n    dropout_ratio)', False, 'from module.Layer.tree_layers import Prediction, GenerateNode, Merge, SemanticAlignmentModule\n'), ((74, 30, 75, 62), 'module.Layer.tree_layers.GenerateNode', 'GenerateNode', ({(74, 43, 74, 59): 'self.hidden_size', (74, 61, 74, 79): 'self.operator_nums', (74, 81, 74, 100): 'self.embedding_size', (75, 43, 75, 61): 'self.dropout_ratio'}, {}), '(self.hidden_size, self.operator_nums, self.embedding_size,\n    self.dropout_ratio)', False, 'from module.Layer.tree_layers import Prediction, GenerateNode, Merge, SemanticAlignmentModule\n'), ((76, 21, 76, 85), 'module.Layer.tree_layers.Merge', 'Merge', ({(76, 27, 76, 43): 'self.hidden_size', (76, 45, 76, 64): 'self.embedding_size', (76, 66, 76, 84): 'self.dropout_ratio'}, {}), '(self.hidden_size, self.embedding_size, self.dropout_ratio)', False, 'from module.Layer.tree_layers import Prediction, GenerateNode, Merge, SemanticAlignmentModule\n'), ((77, 18, 77, 93), 'module.Layer.tree_layers.SemanticAlignmentModule', 'SemanticAlignmentModule', ({(77, 42, 77, 58): 'self.hidden_size', (77, 59, 77, 75): 'self.hidden_size', (77, 76, 77, 92): 'self.hidden_size'}, {}), '(self.hidden_size, self.hidden_size, self.hidden_size)', False, 'from module.Layer.tree_layers import Prediction, GenerateNode, Merge, SemanticAlignmentModule\n'), ((78, 21, 78, 45), 'loss.masked_cross_entropy_loss.MaskedCrossEntropyLoss', 'MaskedCrossEntropyLoss', ({}, {}), '()', False, 'from loss.masked_cross_entropy_loss import MaskedCrossEntropyLoss, masked_cross_entropy\n'), ((94, 21, 94, 59), 'copy.deepcopy', 'copy.deepcopy', ({(94, 35, 94, 58): "batch_data['num stack']"}, {}), "(batch_data['num stack'])", False, 'import copy\n'), ((116, 21, 116, 59), 'copy.deepcopy', 'copy.deepcopy', ({(116, 35, 116, 58): "batch_data['num stack']"}, {}), "(batch_data['num stack'])", False, 'import copy\n'), ((136, 19, 136, 45), 'torch.ByteTensor', 'torch.ByteTensor', ({(136, 36, 136, 44): 'seq_mask'}, {}), '(seq_mask)', False, 'import torch\n'), ((143, 19, 143, 45), 'torch.ByteTensor', 'torch.ByteTensor', ({(143, 36, 143, 44): 'num_mask'}, {}), '(num_mask)', False, 'import torch\n'), ((237, 27, 237, 63), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((254, 34, 254, 46), 'torch.nn.MSELoss', 'nn.MSELoss', ({}, {}), '()', False, 'from torch import nn\n'), ((394, 18, 394, 43), 'torch.LongTensor', 'torch.LongTensor', ({(394, 35, 394, 42): 'indices'}, {}), '(indices)', False, 'import torch\n'), ((395, 23, 395, 53), 'torch.BoolTensor', 'torch.BoolTensor', ({(395, 40, 395, 52): 'masked_index'}, {}), '(masked_index)', False, 'import torch\n'), ((408, 23, 408, 44), 'copy.deepcopy', 'copy.deepcopy', ({(408, 37, 408, 43): 'target'}, {}), '(target)', False, 'import copy\n'), ((428, 15, 428, 27), 'torch.sum', 'torch.sum', ({(428, 25, 428, 26): 'y'}, {}), '(y)', False, 'import torch\n'), ((125, 75, 125, 99), 'utils.utils.copy_list', 'copy_list', ({(125, 85, 125, 98): 'nums_stack[0]'}, {}), '(nums_stack[0])', False, 'from utils.utils import copy_list\n'), ((126, 66, 126, 90), 'utils.utils.copy_list', 'copy_list', ({(126, 76, 126, 89): 'nums_stack[0]'}, {}), '(nums_stack[0])', False, 'from utils.utils import copy_list\n'), ((191, 22, 191, 51), 'torch.cat', 'torch.cat', ({(191, 32, 191, 47): '(op, num_score)', (191, 49, 191, 50): '1'}, {}), '((op, num_score), 1)', False, 'import torch\n'), ((265, 15, 265, 75), 'loss.masked_cross_entropy_loss.masked_cross_entropy', 'masked_cross_entropy', ({(265, 36, 265, 52): 'all_node_outputs', (265, 54, 265, 60): 'target', (265, 61, 265, 74): 'target_length'}, {}), '(all_node_outputs, target, target_length)', False, 'from loss.masked_cross_entropy_loss import MaskedCrossEntropyLoss, masked_cross_entropy\n'), ((309, 17, 309, 79), 'module.Strategy.beam_search.TreeBeam', 'TreeBeam', ({(309, 26, 309, 29): '(0.0)', (309, 31, 309, 42): 'node_stacks', (309, 44, 309, 61): 'embeddings_stacks', (309, 63, 309, 74): 'left_childs', (309, 76, 309, 78): '[]'}, {}), '(0.0, node_stacks, embeddings_stacks, left_childs, [])', False, 'from module.Strategy.beam_search import TreeBeam\n'), ((419, 15, 419, 39), 'torch.LongTensor', 'torch.LongTensor', ({(419, 32, 419, 38): 'target'}, {}), '(target)', False, 'import torch\n'), ((419, 41, 419, 71), 'torch.LongTensor', 'torch.LongTensor', ({(419, 58, 419, 70): 'target_input'}, {}), '(target_input)', False, 'import torch\n'), ((427, 12, 427, 39), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((427, 42, 427, 65), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((35, 47, 35, 67), 'torch.device', 'torch.device', ({(35, 60, 35, 66): '"""cuda"""'}, {}), "('cuda')", False, 'import torch\n'), ((90, 14, 90, 50), 'torch.tensor', 'torch.tensor', ({(90, 27, 90, 49): "batch_data['question']"}, {}), "(batch_data['question'])", False, 'import torch\n'), ((91, 21, 91, 57), 'torch.tensor', 'torch.tensor', ({(91, 34, 91, 56): "batch_data['ques len']"}, {}), "(batch_data['ques len'])", False, 'import torch\n'), ((92, 17, 92, 53), 'torch.tensor', 'torch.tensor', ({(92, 30, 92, 52): "batch_data['equation']"}, {}), "(batch_data['equation'])", False, 'import torch\n'), ((93, 24, 93, 63), 'torch.LongTensor', 'torch.LongTensor', ({(93, 41, 93, 62): "batch_data['equ len']"}, {}), "(batch_data['equ len'])", False, 'import torch\n'), ((113, 14, 113, 50), 'torch.tensor', 'torch.tensor', ({(113, 27, 113, 49): "batch_data['question']"}, {}), "(batch_data['question'])", False, 'import torch\n'), ((114, 21, 114, 57), 'torch.tensor', 'torch.tensor', ({(114, 34, 114, 56): "batch_data['ques len']"}, {}), "(batch_data['ques len'])", False, 'import torch\n'), ((115, 17, 115, 53), 'torch.tensor', 'torch.tensor', ({(115, 30, 115, 52): "batch_data['equation']"}, {}), "(batch_data['equation'])", False, 'import torch\n'), ((169, 24, 169, 35), 'module.Layer.tree_layers.TreeNode', 'TreeNode', ({(169, 33, 169, 34): '_'}, {}), '(_)', False, 'from module.Layer.tree_layers import NodeGenerater, SubTreeMerger, TreeNode, TreeEmbedding\n'), ((277, 19, 277, 52), 'torch.BoolTensor', 'torch.BoolTensor', ({(277, 36, 277, 37): '1', (277, 39, 277, 51): 'input_length'}, {}), '(1, input_length)', False, 'import torch\n'), ((300, 24, 300, 35), 'module.Layer.tree_layers.TreeNode', 'TreeNode', ({(300, 33, 300, 34): '_'}, {}), '(_)', False, 'from module.Layer.tree_layers import NodeGenerater, SubTreeMerger, TreeNode, TreeEmbedding\n'), ((426, 33, 426, 66), 'torch.square', 'torch.square', ({(426, 47, 426, 64): 'outputs - targets'}, {}), '(outputs - targets)', False, 'import torch\n'), ((329, 54, 329, 87), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((336, 41, 336, 64), 'utils.utils.copy_list', 'copy_list', ({(336, 51, 336, 63): 'b.node_stack'}, {}), '(b.node_stack)', False, 'from utils.utils import copy_list\n'), ((338, 48, 338, 76), 'utils.utils.copy_list', 'copy_list', ({(338, 58, 338, 75): 'b.embedding_stack'}, {}), '(b.embedding_stack)', False, 'from utils.utils import copy_list\n'), ((339, 34, 339, 54), 'copy.deepcopy', 'copy.deepcopy', ({(339, 48, 339, 53): 'b.out'}, {}), '(b.out)', False, 'import copy\n'), ((211, 38, 211, 49), 'module.Layer.tree_layers.TreeNode', 'TreeNode', ({(211, 47, 211, 48): 'r'}, {}), '(r)', False, 'from module.Layer.tree_layers import NodeGenerater, SubTreeMerger, TreeNode, TreeEmbedding\n'), ((212, 38, 212, 65), 'module.Layer.tree_layers.TreeNode', 'TreeNode', (), '', False, 'from module.Layer.tree_layers import NodeGenerater, SubTreeMerger, TreeNode, TreeEmbedding\n'), ((228, 29, 228, 70), 'module.Layer.tree_layers.TreeEmbedding', 'TreeEmbedding', (), '', False, 'from module.Layer.tree_layers import NodeGenerater, SubTreeMerger, TreeNode, TreeEmbedding\n'), ((347, 41, 347, 70), 'torch.LongTensor', 'torch.LongTensor', ({(347, 58, 347, 69): '[out_token]'}, {}), '([out_token])', False, 'import torch\n'), ((353, 53, 353, 74), 'module.Layer.tree_layers.TreeNode', 'TreeNode', ({(353, 62, 353, 73): 'right_child'}, {}), '(right_child)', False, 'from module.Layer.tree_layers import NodeGenerater, SubTreeMerger, TreeNode, TreeEmbedding\n'), ((354, 53, 354, 89), 'module.Layer.tree_layers.TreeNode', 'TreeNode', (), '', False, 'from module.Layer.tree_layers import NodeGenerater, SubTreeMerger, TreeNode, TreeEmbedding\n'), ((364, 60, 364, 92), 'module.Layer.tree_layers.TreeEmbedding', 'TreeEmbedding', ({(364, 74, 364, 85): 'current_num', (364, 87, 364, 91): '(True)'}, {}), '(current_num, True)', False, 'from module.Layer.tree_layers import NodeGenerater, SubTreeMerger, TreeNode, TreeEmbedding\n')]