import json
import os
import subprocess
from collections import OrderedDict
from logging import getLogger
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from sklearn.metrics import roc_auc_score, average_precision_score
import torch.multiprocessing as mp
from tqdm import tqdm

# try:
#     mp.set_start_method('spawn')
# except RuntimeError:
#     pass
from typing import Dict, List

from .comp_acc_computation import load_evosuite_transcoder_tests, eval_function_output
from .subtoken_score import run_subtoken_score
from ..data.loader import DATASET_SPLITS, TRAIN_SPLITS
from ..trainer import get_programming_language_name
from ..utils import (
    to_cuda,
    to_device,
    restore_segmentation,
    concat_batches,
    vizualize_translated_files,
    vizualize_do_files,
    show_batch,
    add_noise,
    convert_to_text,
    REPO_ROOT, remove_prefix,
)

from codegen_sources.model.src.model import build_model

import sys

sys.path.append(REPO_ROOT)
from codegen_sources.test_generation.test_runners.cpp_test_runner import CppTestRunner
from codegen_sources.test_generation.test_runners.python_test_runner import (
    PythonTestRunner,
)
import pdb

SRC_ST_LANGS = "java"

TARGET_ST_LANG =  {"python"} #{"cpp", "python"}

EVAL_OBF_PROBAS = []

BLEU_SCRIPT_PATH = os.path.join(
    os.path.abspath(os.path.dirname(__file__)), "multi-bleu.perl"
)
EVAL_DATASET_SPLITS = [ds for ds in DATASET_SPLITS if ds not in TRAIN_SPLITS]
assert os.path.isfile(BLEU_SCRIPT_PATH)
ROOT_FOLDER = Path(__file__).parents[4]
EVAL_SCRIPT_FOLDER = {
    "test": ROOT_FOLDER.joinpath("data/transcoder_evaluation_gfg"),
    "valid": ROOT_FOLDER.joinpath("data/transcoder_evaluation_gfg"),
    "test_st": None
}
logger = getLogger()

class Evaluator(object):
    def __init__(self, trainer, data, params):
        """
        Initialize evaluator.
        """
        self.trainer = trainer
        self.data = data
        self.dico = data["dico"]
        self.params = params
        EVAL_DATASET_SPLITS.extend(params.eval_splits)

        # create directory to store hypotheses, and reference files for BLEU evaluation
        if self.params.is_master:
            params.hyp_path = os.path.join(params.dump_path, "hypotheses")
            subprocess.Popen("mkdir -p %s" % params.hyp_path, shell=True).wait()
            params.eval_scripts_root = os.path.join(params.dump_path, "eval_scripts")
            subprocess.Popen(
                "mkdir -p %s" % params.eval_scripts_root, shell=True
            ).wait()
            self.params.ref_paths = {}
            self.params.id_paths = {}
            self.params.eval_scripts_folders = {}
            if params.eval_bleu or params.eval_subtoken_score:
                self.create_reference_files()
        if self.params.eval_st:
            logger.info("Loading evosuite tests")
            self.evosuite_tests_dico = load_evosuite_transcoder_tests(params)
        else:
            self.evosuite_tests_dico = None

        assert self.params.n_jobs > 0 and self.params.n_workers > 0
        self.parallel_eval = (self.params.n_jobs > 1 and self.params.n_workers > 1)
        if self.parallel_eval:
            self.eval_model_name = "eval_model"

    def get_iterator(
        self,
        data_set,
        lang1,
        lang2=None,
        stream=False,
        span=None,
        st_scores_cutoff=None,
        label=None,
    ):
        """
        Create a new iterator for a dataset.
        """
        assert data_set in EVAL_DATASET_SPLITS
        assert lang1 in self.params.langs
        assert (
            lang2 is None
            or lang2 in self.params.langs
            or (lang1, lang2) in self.params.classif_steps
            or (lang1, lang2, label) in self.params.seq_classif_steps
        )
        assert stream is False or lang2 is None

        n_sentences = self.params.n_sentences_eval
        subsample = 10

        if lang2 is None or lang2 == lang1:
            key = lang1 if span is None else (lang1, span)
            if stream and lang2 is None:
                iterator = self.data["mono_stream"][key][data_set].get_iterator(
                    shuffle=False, subsample=subsample
                )
            else:
                iterator = self.data["mono"][key][data_set].get_iterator(
                    tokens_per_batch=self.params.tokens_per_batch,
                    shuffle=False,
                    group_by_size=True,
                    n_sentences=n_sentences,
                )
        else:
            assert stream is False
            _lang1, _lang2 = (lang1, lang2) if (lang1 < lang2 or label is not None) else (lang2, lang1)
            key = (_lang1, _lang2, span) if span is not None else \
                  (_lang1, _lang2, label) if label is not None else (_lang1, _lang2)
            iterator = self.data["para"][key][data_set].get_iterator(
                shuffle=False,
                group_by_size=True,
                n_sentences=n_sentences,
                tokens_per_batch=self.params.tokens_per_batch,
            )
        for batch in iterator:
            batch = batch if lang2 is None or lang1 == lang2 or lang1 <= lang2 or label is not None else batch[::-1]
            # if there is dataset prefix, remove the prefix for target sequences
            # TODO temp maybe sketchy code
            # pdb.set_trace()
            if self.params.add_dataset_prefix:
                x2 = torch.cat([batch[1][0][[0]], batch[1][0][2:]])  # remove prefix from target sentence
                len2 = batch[1][1] - 1  # decrease target sentence len
                batch = (batch[0], (x2, len2, batch[1][2], batch[1][3]))
            yield batch

    def create_reference_files(self):
        """
        Create reference files for BLEU evaluation.
        """
        logger.info("Creating reference files for BLEU evaluation...")
        params = self.params
        for key, v in self.data["para"].items():
            span = None
            if len(key) == 3:
                lang1, lang2, span = key
            else:
                assert len(key) == 2
                lang1, lang2 = key
            assert lang1 < lang2, (lang1, lang2)

            for data_set in EVAL_DATASET_SPLITS:
                has_sent_ids = (data_set, (lang1, lang2)) in params.has_sentence_ids

                params.eval_scripts_folders[(lang1, lang2, data_set)] = os.path.join(
                    params.eval_scripts_root,
                    "{0}-{1}.{2}".format(lang1, lang2, data_set),
                )
                subprocess.Popen(
                    "mkdir -p %s"
                    % params.eval_scripts_folders[(lang1, lang2, data_set)],
                    shell=True,
                ).wait()
                params.eval_scripts_folders[(lang2, lang1, data_set)] = os.path.join(
                    params.eval_scripts_root,
                    "{0}-{1}.{2}".format(lang2, lang1, data_set),
                )
                subprocess.Popen(
                    "mkdir -p %s"
                    % params.eval_scripts_folders[(lang2, lang1, data_set)],
                    shell=True,
                ).wait()

                # define data paths
                lang1_path = os.path.join(
                    params.hyp_path,
                    "ref.{0}-{1}.{2}.txt".format(lang2, lang1, data_set),
                )
                lang2_path = os.path.join(
                    params.hyp_path,
                    "ref.{0}-{1}.{2}.txt".format(lang1, lang2, data_set),
                )
                spans_path = os.path.join(
                    params.hyp_path,
                    "ref.{0}-{1}-{3}.{2}.txt".format(lang1, lang2, span, data_set),
                )
                id_path = os.path.join(
                    params.hyp_path,
                    "ids.{0}-{1}.{2}.txt".format(lang1, lang2, data_set),
                )
                # store data paths
                params.ref_paths[(lang2, lang1, data_set)] = lang1_path
                params.ref_paths[(lang1, lang2, data_set)] = lang2_path
                params.id_paths[(lang1, lang2, data_set)] = id_path
                params.id_paths[(lang2, lang1, data_set)] = id_path

                # text sentences
                lang1_txt = []
                lang2_txt = []

                id_txt = []
                spans = []
                # convert to text
                for i, batch in enumerate(
                    self.get_iterator(data_set, lang1, lang2, span=span)
                ):
                    logger.info("converting to text for {0}".format(data_set))
                    if len(batch) == 2:
                        (sent1, len1, id1, lenid1), (sent2, len2, id2, lenid2) = batch
                    else:
                        (
                            (sent1, len1, id1, lenid1),
                            (sent2, len2, id2, lenid2),
                            (span_batch, len_span, _, _),
                        ) = batch
                        spans.extend(list(span_batch.T))
                    lang1_txt.extend(convert_to_text(sent1, len1, self.dico, params))
                    lang2_txt.extend(convert_to_text(sent2, len2, self.dico, params))
                    if has_sent_ids:
                        assert id1.equal(id2) and lenid1.equal(lenid2)
                        id_txt.extend(convert_to_text(id1, lenid1, self.dico, params))

                # replace <unk> by <<unk>> as these tokens cannot be counted in BLEU
                lang1_txt = [x.replace("<unk>", "<<unk>>") for x in lang1_txt]
                lang2_txt = [x.replace("<unk>", "<<unk>>") for x in lang2_txt]

                # export hypothesis
                with open(lang1_path, "w", encoding="utf-8") as f:
                    f.write("\n".join(lang1_txt) + "\n")
                with open(lang2_path, "w", encoding="utf-8") as f:
                    f.write("\n".join(lang2_txt) + "\n")
                if len(spans) > 0:
                    with open(spans_path, "w", encoding="utf-8") as f:
                        f.write("\n".join([str(s) for s in spans]) + "\n")

                # restore original segmentation
                restore_segmentation(
                    lang1_path, roberta_mode=params.roberta_mode, single_line=True
                )
                restore_segmentation(
                    lang2_path, roberta_mode=params.roberta_mode, single_line=True
                )

                if has_sent_ids:
                    with open(id_path, "w", encoding="utf-8") as f:
                        f.write("\n".join(id_txt) + "\n")
                    restore_segmentation(
                        id_path, roberta_mode=params.roberta_mode, single_line=True
                    )

    def save_qe_prediction_file(self, predictions: Dict[str, List]):
        # pdb.set_trace()
        predict_class_path = os.path.join(self.params.dump_path, "predictions_classes.txt")
        prediction_path = os.path.join(self.params.dump_path, "predictions.csv")
        with open(predict_class_path, "w", encoding="utf-8") as f:
            f.write("\n".join(predictions["TARGET_CLASS"]) + "\n")
        restore_segmentation(
            predict_class_path, roberta_mode=self.params.roberta_mode, single_line=True
        )
        restored_class_ids = [l.strip() for l in open(predict_class_path).readlines()]
        predictions["TARGET_CLASS"] = restored_class_ids
        df = pd.DataFrame(predictions)
        df.to_csv(prediction_path, index=False)

    def mask_out(self, x, lengths, rng):
        """
        Decide of random words to mask out.
        We specify the random generator to ensure that the test is the same at each epoch.
        """
        params = self.params
        slen, bs = x.size()

        # words to predict - be sure there is at least one word per sentence
        to_predict = rng.rand(slen, bs) <= params.word_pred
        to_predict[0] = 0
        for i in range(bs):
            to_predict[lengths[i] - 1 :, i] = 0
            if not np.any(to_predict[: lengths[i] - 1, i]):
                v = rng.randint(1, lengths[i] - 1)
                to_predict[v, i] = 1
        pred_mask = torch.from_numpy(to_predict.astype(np.uint8))
        pred_mask = pred_mask == 1

        # generate possible targets / update x input
        _x_real = x[pred_mask]
        _x_mask = _x_real.clone().fill_(params.mask_index)
        x = x.masked_scatter(pred_mask, _x_mask)

        assert 0 <= x.min() <= x.max() < params.n_words
        assert x.size() == (slen, bs)
        assert pred_mask.size() == (slen, bs)

        return x, _x_real, pred_mask

    def run_all_evals(self, trainer):
        """
        Run all evaluations.
        """
        if self.parallel_eval:
            trainer.save_checkpoint(name=self.eval_model_name,include_optimizers=False)

        params = self.params
        scores = OrderedDict({"epoch": trainer.epoch})
        deobf_probas_to_eval = EVAL_OBF_PROBAS
        deobfuscation_proba = 1 - params.obf_proba
        if deobfuscation_proba not in deobf_probas_to_eval:
            deobf_probas_to_eval.append(deobfuscation_proba)

        with torch.no_grad():

            for data_set in EVAL_DATASET_SPLITS if not params.eval_splits else params.eval_splits: #["test"]:

                # causal prediction task (evaluate perplexity and accuracy)
                for lang1, lang2 in params.clm_steps:
                    self.evaluate_clm(scores, data_set, lang1, lang2)

                # prediction task (evaluate perplexity and accuracy)
                for lang1, lang2 in params.mlm_steps:
                    self.evaluate_mlm(scores, data_set, lang1, lang2)

                # machine translation task (evaluate perplexity and accuracy)
                for keys in set(
                    params.mt_steps
                    + [(l2, l3) for _, l2, l3 in params.bt_steps]
                    + [(l1, l2) for l1, langs2 in params.st_steps for l2 in langs2]
                    + [(l2, l1) for l1, langs2 in params.st_steps for l2 in langs2]
                    + [
                        (l2_1, l2_2)
                        for l1, langs2 in params.st_steps
                        for l2_1 in langs2
                        for l2_2 in langs2
                        if l2_1 != l2_2
                    ]
                    + params.mt_spans_steps
                ):
                    spans = None
                    assert len(keys) == 2 or len(keys) == 3
                    lang1, lang2 = keys[0], keys[1]
                    if len(keys) == 3:
                        spans = keys[2]
                    # if (lang1 == "java" or lang1 == "java_sa" or lang1 == "java_obfuscated") and \
                    #    (lang2 == "python" or lang2 == "python_sa" or lang2 == "python_obfuscated"):
                    # only evaluate on directions where we keep validation metrics checkpoints
                    if f"{lang1}-{lang2}" in params.validation_metrics:
                        self.evaluate_mt(
                            scores,
                            data_set,
                            lang1,
                            lang2,
                            params.eval_bleu,
                            params.eval_computation,
                            params.eval_subtoken_score,
                            spans,
                        )
                if self.params.eval_denoising:
                    for lang in set(params.ae_steps):
                        assert lang in params.langs, lang
                        self.evaluate_mt(
                            scores,
                            data_set,
                            lang,
                            lang,
                            eval_bleu=False,
                            eval_computation=False,
                            eval_subtoken_score=False,
                            span=None,
                        )

                # machine translation task (evaluate perplexity and accuracy)
                if self.params.eval_do:
                    for lang1, lang2 in set(params.do_steps):
                        assert len(deobf_probas_to_eval) == len(
                            set(deobf_probas_to_eval)
                        ), f"deobf_probas_to_eval should have no duplicates, was {deobf_probas_to_eval}"
                        self.evaluate_mt(
                            scores,
                            data_set,
                            lang1,
                            lang2,
                            params.eval_bleu,
                            eval_computation=False,
                            eval_subtoken_score=params.eval_subtoken_score,
                            span=None,
                            deobfuscate=True,
                            deobfuscate_probas=deobf_probas_to_eval,
                        )

                # token-wise prediction task (evaluate perplexity and accuracy)
                for lang1, lang2 in params.classif_steps:
                    self.evaluate_classif(scores, data_set, lang1, lang2)

                # sequence prediction task (evaluate perplexity and accuracy)
                for lang1, lang2, label in params.seq_classif_steps:
                    self.evaluate_classif(scores, data_set, lang1, lang2, label)

                # report average metrics per language
                if self.params.eval_do and len(params.do_steps) > 0 and params.is_master:
                    for obfuscation_proba in deobf_probas_to_eval:
                        for score_type in ["precision", "recall", "F1"]:
                            scores[
                                "%s_obf_proba_%s_mt_subtoken_%s"
                                % (data_set, 1 - obfuscation_proba, score_type)
                            ] = np.mean(
                                [
                                    scores[
                                        "%s_%s_mt_subtoken_%s"
                                        % (
                                            data_set,
                                            get_l1l2_string(
                                                lang1, lang2, obfuscation_proba
                                            ),
                                            score_type,
                                        )
                                    ]
                                    for lang1, lang2 in params.do_steps
                                ]
                            )
                _clm_mono = [l1 for (l1, l2) in params.clm_steps if l2 is None]
                if len(_clm_mono) > 0:
                    scores["%s_clm_ppl" % data_set] = np.mean(
                        [
                            scores["%s_%s_clm_ppl" % (data_set, lang)]
                            for lang in _clm_mono
                        ]
                    )
                    scores["%s_clm_acc" % data_set] = np.mean(
                        [
                            scores["%s_%s_clm_acc" % (data_set, lang)]
                            for lang in _clm_mono
                        ]
                    )
                _mlm_mono = [l1 for (l1, l2) in params.mlm_steps if l2 is None]
                if len(_mlm_mono) > 0:
                    scores["%s_mlm_ppl" % data_set] = np.mean(
                        [
                            scores["%s_%s_mlm_ppl" % (data_set, lang)]
                            for lang in _mlm_mono
                        ]
                    )
                    scores["%s_mlm_acc" % data_set] = np.mean(
                        [
                            scores["%s_%s_mlm_acc" % (data_set, lang)]
                            for lang in _mlm_mono
                        ]
                    )
        return scores

    def eval_mode(self):
        [enc.eval() for enc in self.encoder]
        if self.decoder is not None:
            [dec.eval() for dec in self.decoder]

    def evaluate_clm(self, scores, data_set, lang1, lang2):
        """
        Evaluate perplexity and next word prediction accuracy.
        """
        params = self.params
        assert data_set in EVAL_DATASET_SPLITS
        assert lang1 in params.langs
        assert lang2 in params.langs or lang2 is None

        model = self.model if params.encoder_only else self.decoder
        model.eval()
        model = model.module if params.multi_gpu else model

        lang1_id = params.lang2id[lang1]
        lang2_id = params.lang2id[lang2] if lang2 is not None else None
        l1l2 = lang1 if lang2 is None else f"{lang1}-{lang2}"

        n_words = 0
        xe_loss = 0
        n_valid = 0

        for batch in self.get_iterator(data_set, lang1, lang2, stream=(lang2 is None)):

            # batch
            if lang2 is None:
                x, lengths = batch
                positions = None
                langs = x.clone().fill_(lang1_id) if params.n_langs > 1 else None
            else:
                (sent1, len1), (sent2, len2) = batch
                x, lengths, positions, langs = concat_batches(
                    sent1,
                    len1,
                    lang1_id,
                    sent2,
                    len2,
                    lang2_id,
                    params.pad_index,
                    params.eos_index,
                    reset_positions=True,
                )

            # words to predict
            alen = torch.arange(lengths.max(), dtype=torch.long, device=lengths.device)
            pred_mask = alen[:, None] < lengths[None] - 1
            y = x[1:].masked_select(pred_mask[:-1])
            assert pred_mask.sum().item() == y.size(0)

            # cuda
            x, lengths, positions, langs, pred_mask, y = to_cuda(
                x, lengths, positions, langs, pred_mask, y
            )

            # forward / loss
            tensor = model(
                "fwd",
                x=x,
                lengths=lengths,
                positions=positions,
                langs=langs,
                causal=True,
            )
            word_scores, loss = model(
                "predict", tensor=tensor, pred_mask=pred_mask, y=y, get_scores=True
            )

            # update stats
            n_words += y.size(0)
            xe_loss += loss.item() * len(y)
            n_valid += (word_scores.max(1)[1] == y).sum().item()

        # log
        logger.info(
            "Found %i words in %s. %i were predicted correctly."
            % (n_words, data_set, n_valid)
        )

        # compute perplexity and prediction accuracy
        ppl_name = "%s_%s_clm_ppl" % (data_set, l1l2)
        acc_name = "%s_%s_clm_acc" % (data_set, l1l2)
        scores[ppl_name] = np.exp(xe_loss / n_words)
        scores[acc_name] = 100.0 * n_valid / n_words

    def evaluate_mlm(self, scores, data_set, lang1, lang2):
        """
        Evaluate perplexity and next word prediction accuracy.
        """
        params = self.params
        assert data_set in EVAL_DATASET_SPLITS
        assert lang1 in params.langs
        assert lang2 in params.langs or lang2 is None

        model = self.model[0] if params.encoder_only else self.encoder[0]
        model.eval()
        model = model.module if params.multi_gpu else model

        rng = np.random.RandomState(0)

        lang1_id = params.lang2id[lang1]
        lang2_id = params.lang2id[lang2] if lang2 is not None else None
        l1l2 = lang1 if lang2 is None else f"{lang1}_{lang2}"

        n_words = 0
        xe_loss = 0
        n_valid = 0

        for i, batch in enumerate(
            self.get_iterator(data_set, lang1, lang2, stream=(lang2 is None))
        ):
            if i > 50:
                break
            # batch
            if lang2 is None:
                x, lengths = batch
                positions = None
                langs = x.clone().fill_(lang1_id) if params.n_langs > 1 else None
            else:
                (sent1, len1, _, _), (sent2, len2, _, _) = batch
                x, lengths, positions, langs = concat_batches(
                    sent1,
                    len1,
                    lang1_id,
                    sent2,
                    len2,
                    lang2_id,
                    params.pad_index,
                    params.eos_index,
                    reset_positions=True,
                )

            # words to predict
            x, y, pred_mask = self.mask_out(x, lengths, rng)

            # log first batch of training
            if i < 1:
                show_batch(
                    logger,
                    [("masked source", x.transpose(0, 1))],
                    self.data["dico"],
                    self.params.roberta_mode,
                    "Evaluation",
                )

            # cuda
            x, y, pred_mask, lengths, positions, langs = to_cuda(
                x, y, pred_mask, lengths, positions, langs
            )

            # forward / loss
            tensor = model(
                "fwd",
                x=x,
                lengths=lengths,
                positions=positions,
                langs=langs,
                causal=False,
            )
            word_scores, loss = model(
                "predict", tensor=tensor, pred_mask=pred_mask, y=y, get_scores=True
            )

            # update stats
            n_words += len(y)
            xe_loss += loss.item() * len(y)
            n_valid += (word_scores.max(1)[1] == y).sum().item()

        # compute perplexity and prediction accuracy
        ppl_name = "%s_%s_mlm_ppl" % (data_set, l1l2)
        acc_name = "%s_%s_mlm_acc" % (data_set, l1l2)
        scores[ppl_name] = np.exp(xe_loss / n_words) if n_words > 0 else 1e9
        scores[acc_name] = 100.0 * n_valid / n_words if n_words > 0 else 0.0

    def evaluate_classif(self, scores, data_set, lang1, lang2, label=None):
        params = self.params
        assert data_set in EVAL_DATASET_SPLITS
        assert lang1 in params.langs
        lang1_id = params.lang2id[lang1]

        model = self.model[0] if params.encoder_only else self.encoder[0]
        model.eval()
        model = model.module if params.multi_gpu else model
        assert self.classifier is not None
        classifier = self.classifier[0].eval()

        n_words = 0
        n_valid = 0
        labels = []
        word_probas = []
        n_words_by_cl = [0 for c in range(self.params.n_classes_classif)]
        n_valid_by_cl = [0 for c in range(self.params.n_classes_classif)]
        n_attribution_by_cl = [0 for c in range(self.params.n_classes_classif)]

        predictions = {f"logit{cl}":[] for cl in range(self.params.n_classes_classif)}
        predictions["TARGET_CLASS"]= []
        for batch in self.get_iterator(data_set, lang1, lang2, stream=False, label=label):
            if label is not None:
                (x1, len1, id1, lenid1), (y, len2, id2, lenid2) = batch
                pred_mask = torch.zeros(x1.shape[0], x1.shape[1]).bool()
                pred_mask[0, :] = True  # use first token (eos) as sentence token
                assert all(len2 == 3)  # label shape is 3 x n, where y[1,:] holds class index label
                if all(y[1, :] == -1):  # no label, has to be wild set, only need to predict and save logits
                    y = None
                    predictions["TARGET_CLASS"].extend(convert_to_text(id1, lenid1, self.dico, params))
            else:
                (x1, len1, _, _), (y, len2, _, _) = batch
                pred_mask = (x1 != self.params.eos_index) * (x1 != self.params.pad_index)
                assert len1.equal(len2)
            langs1 = x1.clone().fill_(lang1_id)
            # pdb.set_trace()
            # cuda
            x1, len1, langs1, y = to_cuda(x1, len1, langs1, y)

            # encode source sentence
            enc1 = model("fwd", x=x1, lengths=len1, langs=langs1, causal=False)
            if self.params.fp16:
                enc1 = enc1.half()

            # classification + loss
            word_scores, loss = classifier(enc1, y, pred_mask)

            # if no label provided, save word_scores and return
            if y is None:
                for cl in range(self.params.n_classes_classif):
                    predictions[f"logit{cl}"].extend(word_scores[:, cl].cpu().numpy().tolist())
                continue

            # update stats
            y_ = y[pred_mask].view(-1,) if label is None else y[1,:].view(-1,)
            n_words += len(y_)
            n_valid += (word_scores.max(1)[1] == y_).sum().item()
            labels.extend(y_.cpu().numpy())
            word_probas.extend(word_scores.cpu().numpy())

            for cl in range(self.params.n_classes_classif):
                n_words_by_cl[cl] += (y_ == cl).sum().item()
                n_valid_by_cl[cl] += (
                    ((word_scores.max(1)[1] == y_) * (y_ == cl)).sum().item()
                )
                n_attribution_by_cl[cl] += (word_scores.max(1)[1] == cl).sum().item()

        # if no labels, save predictions and exit
        if len(predictions["TARGET_CLASS"]) > 0:
            # pdb.set_trace()
            self.save_qe_prediction_file(predictions)
            return

        if len(set(labels)) > 1:
            for target_label in range(self.params.n_classes_classif):
                roc_auc_name = "%s_%s-%s_roc_auc_label_cl%i" % (
                    data_set,
                    lang1,
                    lang2,
                    target_label,
                )
                new_labels = [1 if l == target_label else 0 for l in labels]
                word_level_scores = [wp[target_label] for wp in word_probas]
                scores[roc_auc_name] = roc_auc_score(new_labels, word_level_scores)

                pr_auc_name = "%s_%s-%s_pr_auc_cl%i" % (
                    data_set,
                    lang1,
                    lang2,
                    target_label,
                )
                scores[pr_auc_name] = average_precision_score(
                    new_labels, word_level_scores
                )

            roc_auc_name = "%s_%s-%s_roc_auc_label_all_changes" % (
                data_set,
                lang1,
                lang2,
            )
            new_labels = [1 if l > 0 else 0 for l in labels]
            word_level_scores = [1 - s[0] for s in word_probas]
            scores[roc_auc_name] = roc_auc_score(new_labels, word_level_scores)

            pr_auc_name = "%s_%s-%s_pr_auc_label_all_changes" % (data_set, lang1, lang2)
            scores[pr_auc_name] = average_precision_score(new_labels, word_level_scores)

        # compute perplexity and prediction accuracy
        class_proportion_name = "%s_%s-%s_class_proportion" % (data_set, lang1, lang2)
        acc_name = "%s_%s-%s_classif_acc" % (data_set, lang1, lang2)
        recall_name = "%s_%s-%s_classif_recall" % (data_set, lang1, lang2)
        precision_name = "%s_%s-%s_classif_precision" % (data_set, lang1, lang2)

        scores[class_proportion_name] = [
            (100.0 * x / n_words) if n_words > 0 else 0.0 for x in n_words_by_cl
        ]
        scores[acc_name] = (100.0 * n_valid / n_words) if n_words > 0 else 0.0
        # scores[recall_name] = [(100. * n_valid_by_cl[cl] / n_words_by_cl[cl]) if n_words_by_cl[cl] > 0 else 0 for cl in range(self.params.n_classes_classif)]
        # scores[precision_name] = [(100. * n_valid_by_cl[cl] / n_attribution_by_cl[cl]) if n_attribution_by_cl[cl] > 0 else 0 for cl in range(self.params.n_classes_classif)]

        for cl in range(params.n_classes_classif):
            scores[f"{recall_name}_{cl}"] = (
                100.0 * n_valid_by_cl[cl] / n_words_by_cl[cl]
                if n_words_by_cl[cl] > 0
                else 0
            )
        for cl in range(params.n_classes_classif):
            scores[f"{precision_name}_{cl}"] = (
                100.0 * n_valid_by_cl[cl] / n_attribution_by_cl[cl]
                if n_attribution_by_cl[cl] > 0
                else 0
            )
        # pdb.set_trace()

class SingleEvaluator(Evaluator):
    def __init__(self, trainer, data, params):
        """
        Build language model evaluator.
        """
        super().__init__(trainer, data, params)
        self.model = trainer.model
        if params.use_classifier:
            self.classifier = trainer.classifier

class EncDecEvaluator(Evaluator):
    def __init__(self, trainer, data, params):
        """
        Build encoder / decoder evaluator.
        """
        super().__init__(trainer, data, params)
        self.encoder = trainer.encoder
        self.decoder = trainer.decoder

    def evaluate_mt(
        self,
        scores,
        data_set,
        lang1,
        lang2,
        eval_bleu,
        eval_computation,
        eval_subtoken_score,
        span,
        deobfuscate=False,
        deobfuscate_probas=None,
    ):
        """
        Evaluate perplexity and next word prediction accuracy.
        """
        params = self.params
        assert data_set in EVAL_DATASET_SPLITS
        assert lang1 in params.langs
        assert lang2 in params.langs
        rng = np.random.RandomState(0)
        torch_rng = torch.Generator().manual_seed(0)
        eval_st = params.eval_st
        if not params.is_master or "cl" in lang1:
            # Computing the accuracy on every node is useful for debugging but
            # no need to evaluate spend too much time on the evaluation when not on master
            eval_bleu = False
            eval_computation = False
            eval_subtoken_score = False
            eval_st = False

        # store hypothesis to compute BLEU score
        if params.eval_bleu_test_only:
            datasets_for_bleu = ["test"]
        else:
            datasets_for_bleu = [s for s in EVAL_DATASET_SPLITS if s not in TRAIN_SPLITS] #  "test"

        lang1_id = params.lang2id[lang1]
        lang2_id = params.lang2id[lang2]

        self.eval_mode()
        encoder = self.encoder[0].module if params.multi_gpu else self.encoder[0]
        decoder = (
            self.decoder[lang2_id] if params.separate_decoders else
            self.decoder[params.train_splits.index("wild")] if params.separate_dataset_decoders else
            self.decoder[0]
        )
        decoder = decoder.module if params.multi_gpu else decoder

        for deobfuscation_proba in (
            deobfuscate_probas if deobfuscate_probas is not None else [None]
        ):
            if deobfuscate:
                rng = np.random.RandomState(0)
            hypothesis, n_valid, n_words, references, sources, xe_loss = self.decode_and_test_data(data_set,
                                                                                datasets_for_bleu,
                                                                                decoder, deobfuscate,
                                                                                deobfuscation_proba,
                                                                                encoder, eval_bleu,
                                                                                eval_computation,
                                                                                eval_subtoken_score,
                                                                                lang1, lang1_id,
                                                                                lang2, lang2_id,
                                                                                params, rng, span,
                                                                                torch_rng,
                                                                                self.params.n_jobs,
                                                                                self.params.n_workers)

            # compute perplexity and prediction accuracy
            l1l2 = get_l1l2_string(lang1, lang2, deobfuscation_proba)
            scores["%s_%s_mt_ppl" % (data_set, l1l2)] = np.exp(xe_loss / n_words)
            scores["%s_%s_mt_acc" % (data_set, l1l2)] = 100.0 * n_valid / n_words

            # write hypotheses
            if (
                eval_bleu or eval_computation or eval_subtoken_score
            ) and data_set in datasets_for_bleu:
                hyp_paths, ref_path, src_path = self.write_hypo_ref_src(
                    data_set,
                    hypothesis,
                    lang1,
                    lang2,
                    params,
                    references,
                    scores,
                    sources,
                    deobfuscation_proba,
                )

            # check how many functions compiles + return same output as GT
            if eval_computation and data_set in datasets_for_bleu and data_set != "test_st":
                print("compute_comp_acc")
                self.compute_comp_acc(
                    data_set,
                    hyp_paths,
                    hypothesis,
                    lang1,
                    lang2,
                    params,
                    ref_path,
                    scores,
                    roberta_mode=params.roberta_mode,
                )
            # pdb.set_trace()
            if (
                eval_st
                and data_set in ("test_st") #datasets_for_bleu
                and get_programming_language_name(lang1) == SRC_ST_LANGS
                and get_programming_language_name(lang2) in TARGET_ST_LANG
            ):
                logger.info("Computing ST comp acc")
                self.compute_comp_acc(
                    data_set,
                    hyp_paths,
                    hypothesis,
                    lang1,
                    lang2,
                    params,
                    ref_path,
                    scores,
                    roberta_mode=params.roberta_mode,
                    evosuite_functions=True,
                )
            if eval_subtoken_score and data_set in datasets_for_bleu:
                subtoken_level_scores = run_subtoken_score(ref_path, hyp_paths)
                for score_type, value in subtoken_level_scores.items():
                    logger.info(
                        "Subtoken %s score %s %s : %f"
                        % (score_type, hyp_paths, ref_path, value)
                    )
                    scores[
                        "%s_%s_mt_subtoken_%s"
                        % (
                            data_set,
                            get_l1l2_string(lang1, lang2, deobfuscation_proba),
                            score_type,
                        )
                    ] = value

            # compute BLEU score
            if eval_bleu and data_set in datasets_for_bleu:
                # evaluate BLEU score
                bleu = eval_moses_bleu(ref_path, hyp_paths[0])
                logger.info("BLEU %s %s : %f" % (hyp_paths[0], ref_path, bleu))
                scores[
                    "%s_%s_mt_bleu"
                    % (data_set, get_l1l2_string(lang1, lang2, deobfuscation_proba))
                ] = bleu
                if eval_computation:
                    for hyp_path in hyp_paths:
                        Path(hyp_path).unlink()

            if (
                deobfuscate
                and eval_bleu
                or eval_subtoken_score
                and data_set in datasets_for_bleu
            ):
                # TODO clean lang1
                vizualize_do_files(lang1.split("_")[0], src_path, ref_path, hyp_paths)

    def decode_and_test_data(self, data_set, datasets_for_bleu, decoder, deobfuscate, deobfuscation_proba, encoder,
                             eval_bleu, eval_computation, eval_subtoken_score, lang1, lang1_id, lang2, lang2_id, params,
                             rng, span, torch_rng, n_jobs=1, n_workers=1):

        if self.parallel_eval:
            ## swap decoder and encoder with None (because instead we will reload it)
            return self.decode_and_test_data_parallel_driver(data_set, datasets_for_bleu, None, deobfuscate,
                                                             deobfuscation_proba, None, eval_bleu, eval_computation,
                                                             eval_subtoken_score, lang1, lang1_id, lang2, lang2_id,
                                                             params, rng, span, torch_rng, n_jobs, n_workers)
        else:
            return self.decode_and_test_data_single(data_set, datasets_for_bleu, decoder, deobfuscate,
                                                    deobfuscation_proba, encoder,eval_bleu, eval_computation,
                                                    eval_subtoken_score, lang1, lang1_id, lang2,  lang2_id,
                                                    params, rng, span, torch_rng)

    def decode_and_test_data_single(self, data_set, datasets_for_bleu, decoder, deobfuscate, deobfuscation_proba, encoder,
                             eval_bleu, eval_computation, eval_subtoken_score, lang1, lang1_id, lang2, lang2_id, params,
                             rng, span, torch_rng):
        n_words = 0
        xe_loss = 0
        n_valid = 0
        hypothesis = []
        sources = []
        references = []
        for i, batch in enumerate(
                self.get_iterator(
                    data_set, lang1, lang2 if lang2 != lang1 else None, span=span
                )
        ):
            spans = None
            assert len(batch) >= 2
            if len(batch) == 2:
                if lang1 == lang2:
                    x2, len2 = batch
                    x1, len1 = add_noise(
                        x2,
                        len2,
                        self.params,
                        len(self.data["dico"]) - 1,
                        rng,
                        torch_rng,
                    )
                else:
                    (x1, len1, ids1, len_ids1), (x2, len2, ids2, len_ids2) = batch
                    if deobfuscate:
                        (x1, len1, x2, len2) = self.trainer.deobfuscate_by_variable(
                            x1, x2, deobfuscation_proba, params.roberta_mode, rng
                        )
                        if x1 is None:
                            continue
            else:
                assert len(batch) == 3
                (
                    (x1, len1, ids1, len_ids1),
                    (x2, len2, ids2, len_ids2),
                    (spans, len_spans, _, _),
                ) = batch

            langs1 = x1.clone().fill_(lang1_id)
            langs2 = x2.clone().fill_(lang2_id)

            # target words to predict
            alen = torch.arange(len2.max(), dtype=torch.long, device=len2.device)
            pred_mask = (
                    alen[:, None] < len2[None] - 1
            )  # do not predict anything given the last target word
            y = x2[1:].masked_select(pred_mask[:-1])
            assert len(y) == (len2 - 1).sum().item()

            # cuda
            x1, len1, langs1, x2, len2, langs2, y, spans = to_cuda(
                x1, len1, langs1, x2, len2, langs2, y, spans
            )
            # encode source sentence
            enc1 = encoder(
                "fwd", x=x1, lengths=len1, langs=langs1, causal=False, spans=spans
            )
            enc1 = enc1.transpose(0, 1)
            enc1 = enc1.half() if params.fp16 else enc1
            # decode target sentence
            dec2 = decoder(
                "fwd",
                x=x2,
                lengths=len2,
                langs=langs2,
                causal=True,
                src_enc=enc1,
                src_len=len1,
                spans=spans,
            )

            # loss
            word_scores, loss = decoder(
                "predict", tensor=dec2, pred_mask=pred_mask, y=y, get_scores=True
            )

            # update stats
            n_words += y.size(0)
            xe_loss += loss.item() * len(y)
            n_valid += (word_scores.max(1)[1] == y).sum().item()

            # generate translation - translate / convert to text
            if (
                    eval_bleu or eval_computation or eval_subtoken_score
            ) and data_set in datasets_for_bleu:
                len_v = (3 * len1 + 10).clamp(max=params.max_len)
                if params.beam_size == 1:
                    if params.number_samples > 1:
                        assert params.eval_temperature is not None
                        generated, lengths = decoder.generate(
                            enc1.repeat_interleave(params.number_samples, dim=0),
                            len1.repeat_interleave(params.number_samples, dim=0),
                            lang2_id,
                            max_len=len_v.repeat_interleave(
                                params.number_samples, dim=0
                            ),
                            sample_temperature=params.eval_temperature,
                        )
                        generated = generated.T.reshape(
                            -1, params.number_samples, generated.shape[0]
                        ).T
                        lengths, _ = lengths.reshape(-1, params.number_samples).max(
                            dim=1
                        )
                    else:
                        generated, lengths = decoder.generate(
                            enc1, len1, lang2_id, max_len=len_v
                        )
                    # print(f'path 1: {generated.shape}')

                else:
                    assert params.number_samples == 1
                    generated, lengths, _ = decoder.generate_beam(
                        enc1,
                        len1,
                        lang2_id,
                        beam_size=params.beam_size,
                        length_penalty=params.length_penalty,
                        early_stopping=params.early_stopping,
                        max_len=len_v,
                    )
                    # print(f'path 2: {generated.shape}')
                if i == 0:
                    # show 1 evaluation example and the corresponding model generation
                    show_batch(
                        logger,
                        [
                            ("source", x1.transpose(0, 1)),
                            ("target", x2.transpose(0, 1)),
                            (
                                "gen",
                                generated.transpose(0, 1)
                                if len(generated.shape) == 2
                                else generated[:, 0, :].transpose(0, 1),
                            ),
                        ],
                        self.data["dico"],
                        self.params.roberta_mode,
                        f"{data_set} {lang1}-{lang2}",
                    )

                hypothesis.extend(
                    convert_to_text(
                        generated,
                        lengths,
                        self.dico,
                        params,
                        generate_several_reps=True,
                    )
                )
                references.extend(convert_to_text(x2, len2, self.dico, params))
                sources.extend(convert_to_text(x1, len1, self.dico, params))
        return hypothesis, n_valid, n_words, references, sources, xe_loss

    def decode_and_test_data_parallel_worker(self,data_set, datasets_for_bleu, decoder, deobfuscate, deobfuscation_proba,
                                      encoder, eval_bleu, eval_computation, eval_subtoken_score, lang1, lang1_id,
                                      lang2, lang2_id, params, rng, span, torch_rng, worker_id, device_no):

        #### Reload Model ####
        reloaded = torch.load(os.path.join(self.params.dump_path, (self.eval_model_name + ".pth")), map_location="cpu")
        # change params of the reloaded model so that it will
        # relaod its own weights and not the MLM or DOBF pretrained model
        reloaded["params"]["reload_model"] = ",".join([model_path] * 2)
        reloaded["params"]["lgs_mapping"] = ""
        reloaded["params"]["reload_encoder_for_decoder"] = False
        self.reloaded_params = AttrDict(reloaded["params"])

        # build dictionary / update parameters
        self.dico = Dictionary(
            reloaded["dico_id2word"], reloaded["dico_word2id"], reloaded["dico_counts"]
        )

        assert self.reloaded_params.n_words == len(self.dico)
        assert self.reloaded_params.bos_index == self.dico.index(BOS_WORD)
        assert self.reloaded_params.eos_index == self.dico.index(EOS_WORD)
        assert self.reloaded_params.pad_index == self.dico.index(PAD_WORD)
        assert self.reloaded_params.unk_index == self.dico.index(UNK_WORD)
        assert self.reloaded_params.mask_index == self.dico.index(MASK_WORD)

        # build model / reload weights (in the build_model method)
        print("Building model with device number {}...".format(device_no), flush=True, file=sys.stderr)
        gpu = torch.cuda_is_available() if device_no is None else False
        encoder, decoder = build_model(self.reloaded_params, self.dico, gpu=gpu,
                                       device_no=device_no)

        encoder = encoder[0]
        decoder = decoder[0]

        #### Model Reloadad, Decode ####

        n_words = 0
        xe_loss = 0
        n_valid = 0
        hypothesis = []
        sources = []
        references = []
        for i, batch in enumerate(
                self.get_iterator(
                    data_set, lang1, lang2 if lang2 != lang1 else None, span=span
                )
        ):
            spans = None
            assert len(batch) >= 2
            if len(batch) == 2:
                raise NotImplementedError("not supported in parallel mode yet")
                if lang1 == lang2:
                    x2, len2 = batch
                    x1, len1 = add_noise(
                        x2,
                        len2,
                        self.params,
                        len(self.data["dico"]) - 1,
                        rng,
                        torch_rng,
                    )
                else:
                    (x1, len1, ids1, len_ids1), (x2, len2, ids2, len_ids2) = batch
                    if deobfuscate:
                        (x1, len1, x2, len2) = self.trainer.deobfuscate_by_variable(
                            x1, x2, deobfuscation_proba, params.roberta_mode, rng
                        )
                        if x1 is None:
                            continue
            else:
                assert len(batch) == 3
                (
                    (x1, len1, ids1, len_ids1),
                    (x2, len2, ids2, len_ids2),
                    (spans, len_spans, _, _),
                ) = batch

            langs1 = x1.clone().fill_(lang1_id)
            langs2 = x2.clone().fill_(lang2_id)

            # target words to predict
            alen = torch.arange(len2.max(), dtype=torch.long, device=len2.device)
            pred_mask = (
                    alen[:, None] < len2[None] - 1
            )  # do not predict anything given the last target word
            y = x2[1:].masked_select(pred_mask[:-1])
            assert len(y) == (len2 - 1).sum().item()

            # cuda
            device_id = f"cuda:{device_no}"
            x1, len1, langs1, x2, len2, langs2, y, spans = to_device(
                x1, len1, langs1, x2, len2, langs2, y, spans, device=device_id
            )
            # encode source sentence
            enc1 = encoder(
                "fwd", x=x1, lengths=len1, langs=langs1, causal=False, spans=spans
            )
            enc1 = enc1.transpose(0, 1)
            enc1 = enc1.half() if params.fp16 else enc1
            # decode target sentence
            dec2 = decoder(
                "fwd",
                x=x2,
                lengths=len2,
                langs=langs2,
                causal=True,
                src_enc=enc1,
                src_len=len1,
                spans=spans,
            )

            # loss
            word_scores, loss = decoder(
                "predict", tensor=dec2, pred_mask=pred_mask, y=y, get_scores=True
            )

            # update stats
            n_words += y.size(0)
            xe_loss += loss.item() * len(y)
            n_valid += (word_scores.max(1)[1] == y).sum().item()

            # generate translation - translate / convert to text
            if (
                    eval_bleu or eval_computation or eval_subtoken_score
            ) and data_set in datasets_for_bleu:
                len_v = (3 * len1 + 10).clamp(max=params.max_len)
                if params.beam_size == 1:
                    if params.number_samples > 1:
                        assert params.eval_temperature is not None
                        generated, lengths = decoder.generate(
                            enc1.repeat_interleave(params.number_samples, dim=0),
                            len1.repeat_interleave(params.number_samples, dim=0),
                            lang2_id,
                            max_len=len_v.repeat_interleave(
                                params.number_samples, dim=0
                            ),
                            sample_temperature=params.eval_temperature,
                        )
                        generated = generated.T.reshape(
                            -1, params.number_samples, generated.shape[0]
                        ).T
                        lengths, _ = lengths.reshape(-1, params.number_samples).max(
                            dim=1
                        )
                    else:
                        generated, lengths = decoder.generate(
                            enc1, len1, lang2_id, max_len=len_v
                        )
                    # print(f'path 1: {generated.shape}')

                else:
                    assert params.number_samples == 1
                    generated, lengths, _ = decoder.generate_beam(
                        enc1,
                        len1,
                        lang2_id,
                        beam_size=params.beam_size,
                        length_penalty=params.length_penalty,
                        early_stopping=params.early_stopping,
                        max_len=len_v,
                    )
                    # print(f'path 2: {generated.shape}')
                if i == 0:
                    # show 1 evaluation example and the corresponding model generation
                    show_batch(
                        logger,
                        [
                            ("source", x1.transpose(0, 1)),
                            ("target", x2.transpose(0, 1)),
                            (
                                "gen",
                                generated.transpose(0, 1)
                                if len(generated.shape) == 2
                                else generated[:, 0, :].transpose(0, 1),
                            ),
                        ],
                        self.data["dico"],
                        self.params.roberta_mode,
                        f"{data_set} {lang1}-{lang2}",
                    )

                hypothesis.extend(
                    convert_to_text(
                        generated,
                        lengths,
                        self.dico,
                        params,
                        generate_several_reps=True,
                    )
                )
                references.extend(convert_to_text(x2, len2, self.dico, params))
                sources.extend(convert_to_text(x1, len1, self.dico, params))
        # clear memory
        del encoder
        del decoder
        return hypothesis, n_valid, n_words, references, sources, xe_loss

    def decode_and_test_data_parallel_driver(self, data_set, datasets_for_bleu, decoder, deobfuscate, deobfuscation_proba,
                                      encoder, eval_bleu, eval_computation, eval_subtoken_score, lang1, lang1_id,
                                      lang2, lang2_id, params, rng, span, torch_rng, n_jobs, n_workers):

        n_words = 0
        xe_loss = 0
        n_valid = 0
        hypothesis = []
        sources = []
        references = []
        jobs = []
        for i in range(n_jobs):
            device_no = (i & n_workers)
            # do not pass down the encoder or the decoder
            jobs.append((
                    data_set,
                    datasets_for_bleu,
                    None, #decoder.to("cpu"),
                    deobfuscate,
                    deobfuscation_proba,
                    None, #encoder.to("cpu"),
                    eval_bleu,
                    eval_computation,
                    eval_subtoken_score,
                    lang1,
                    lang1_id,
                    lang2,
                    lang2_id,
                    params,
                    rng,
                    span,
                    torch_rng,
                    i,
                    device_no))

        with mp.Pool(n_jobs) as pool:
            for job in tqdm(pool.starmap_async(self.decode_and_test_data_parallel_worker, jobs).get()):
                hypothesis.extend(job[0])
                n_valid += job[1]
                n_words += job[2]
                references.extend(job[3])
                sources.extend(job[4])
                xe_loss += job[5]

        return hypothesis, n_valid, n_words, references, sources, xe_loss

    @staticmethod
    def write_hypo_ref_src(
        data_set,
        hypothesis,
        lang1,
        lang2,
        params,
        references,
        scores,
        sources=None,
        deobfuscation_proba=None,
    ):
        # hypothesis / reference paths
        hyp_paths = []
        ref_name = "ref.{0}.{1}.txt".format(
            get_l1l2_string(lang1, lang2, deobfuscation_proba), data_set
        )
        ref_path = os.path.join(params.hyp_path, ref_name)
        # remove any prefix in reference if present before saving to hypothesis file (for dataset prefix runs)
        hypothesis = remove_prefix(hypothesis)

        # export sentences to hypothesis file / restore BPE segmentation
        for beam_number in range(len(hypothesis[0])):
            hyp_name = "hyp{0}.{1}.{2}_beam{3}.txt".format(
                scores["epoch"],
                get_l1l2_string(lang1, lang2, deobfuscation_proba),
                data_set,
                beam_number,
            )
            hyp_path = os.path.join(params.hyp_path, hyp_name)
            hyp_paths.append(hyp_path)
            print(f"outputing hypotheses in {hyp_path}")
            with open(hyp_path, "w", encoding="utf-8") as f:
                f.write("\n".join([hyp[beam_number] for hyp in hypothesis]) + "\n")
            restore_segmentation(
                hyp_path, roberta_mode=params.roberta_mode, single_line=True
            )
        # export reference to ref file / restore BPE segmentation
        with open(ref_path, "w", encoding="utf-8") as f:
            f.write("\n".join([ref for ref in references]) + "\n")
        restore_segmentation(
            ref_path, roberta_mode=params.roberta_mode, single_line=True
        )
        if sources:
            src_path = ref_path.replace("ref.", "src.")
            with open(src_path, "w", encoding="utf-8") as f:
                f.write("\n".join([src for src in sources]) + "\n")
            restore_segmentation(
                src_path, roberta_mode=params.roberta_mode, single_line=True
            )
            return hyp_paths, ref_path, src_path
        return hyp_paths, ref_path, None

    def compute_comp_acc(
        self,
        data_set,
        hyp_paths,
        hypothesis,
        lang1,
        lang2,
        params,
        ref_path,
        scores,
        roberta_mode=False,
        evosuite_functions=False,
    ):
        assert self.evosuite_tests_dico is not None or not evosuite_functions
        eval_script_dir = EVAL_SCRIPT_FOLDER[data_set]
        if params.eval_fixed_data:
            eval_script_dir = Path(str(eval_script_dir)+"_fixed")
        func_run_stats, func_run_out = eval_function_output(
            ref_path,
            hyp_paths,
            params.id_paths[(lang1, lang2, data_set)],
            lang2,
            params.eval_scripts_folders[(lang1, lang2, data_set)],
            eval_script_dir,
            params.retry_mismatching_types,
            roberta_mode,
            evosuite_functions=evosuite_functions,
            evosuite_tests=self.evosuite_tests_dico,
        )
        out_paths = []
        success_for_beam_number = [0 for _ in range(len(hypothesis[0]))]
        prefix = "st" if evosuite_functions else ""
        for beam_number in range(len(hypothesis[0])):
            out_name = prefix + "hyp{0}.{1}-{2}.{3}_beam{4}.out.txt".format(
                scores["epoch"], lang1, lang2, data_set, beam_number
            )
            out_path = os.path.join(params.hyp_path, out_name)
            out_paths.append(out_path)
            with open(out_path, "w", encoding="utf-8") as f:
                for results_list in func_run_out:
                    result_for_beam = (
                        results_list[beam_number]
                        if beam_number < len(results_list)
                        else ""
                    )
                    if result_for_beam.startswith("success"):
                        success_for_beam_number[beam_number] += 1
                    f.write(result_for_beam + "\n")
                f.write("\n")
        vizualize_translated_files(
            lang1,
            lang2,
            params.ref_paths[(lang2, lang1, data_set)],
            hyp_paths,
            params.id_paths[(lang1, lang2, data_set)],
            ref_path,
            out_paths,
        )
        logger.info(
            prefix
            + "Computation res %s %s %s : %s"
            % (data_set, lang1, lang2, json.dumps(func_run_stats))
        )
        total_evaluated = func_run_stats["total_evaluated"] if func_run_stats["total_evaluated"] else 1
        scores[
            "%s_%s-%s_mt_comp_acc" % (data_set + prefix, lang1, lang2)
        ] = func_run_stats["success"] / total_evaluated
        scores["%s_%s-%s_mt_failure" % (data_set + prefix, lang1, lang2)] = func_run_stats["failure"] / total_evaluated
        scores["%s_%s-%s_mt_error" % (data_set + prefix, lang1, lang2)] = func_run_stats["error"] / total_evaluated
        scores["%s_%s-%s_mt_gold" % (data_set + prefix, lang1, lang2)] = func_run_stats["identical_gold"] / total_evaluated
        scores["%s_%s-%s_mt_timeout" % (data_set + prefix, lang1, lang2)] = func_run_stats["timeout"] / total_evaluated
        for beam_number, success_for_beam in enumerate(success_for_beam_number):
            scores[
                "%s_%s-%smt_comp_acc_contrib_beam_%i"
                % (data_set + prefix, lang1, lang2, beam_number)
            ] = (
                success_for_beam
                / (
                    func_run_stats["total_evaluated"]
                    if func_run_stats["total_evaluated"]
                    else 1
                )
            )
        for out_path in out_paths:
            Path(out_path).unlink()

def get_l1l2_string(lang1, lang2, deobfuscation_proba):
    l1l2 = [lang1, lang2]
    if deobfuscation_proba is not None:
        l1l2.append(f"obf_proba_{1 - deobfuscation_proba}")

    l1l2 = "-".join(l1l2)
    return l1l2

def eval_moses_bleu(ref, hyp):
    """
    Given a file of hypothesis and reference files,
    evaluate the BLEU score using Moses scripts.
    """
    assert os.path.isfile(hyp)
    assert os.path.isfile(ref) or os.path.isfile(ref + "0")
    assert os.path.isfile(BLEU_SCRIPT_PATH)
    command = BLEU_SCRIPT_PATH + " %s < %s"
    p = subprocess.Popen(command % (ref, hyp), stdout=subprocess.PIPE, shell=True)
    result = p.communicate()[0].decode("utf-8")
    if result.startswith("BLEU"):
        return float(result[7 : result.index(",")])
    else:
        logger.warning('Impossible to parse BLEU score! "%s"' % result)
        return -1


for _ in range(int(input())):
    S = int(input())
    L = []
    for s in range(S):
        L.append(list(map(int, input().split())))
    i = S-1
    while(i > 0):
        for j in range(i):
            L[i-1][j] += max(L[i][j], L[i][j+1])
        i -= 1
    print(L[0][0])


# encoding: utf-8

from __future__ import unicode_literals

from base64 import b64encode
from zlib import compress

def redelta_encode(numbers):
	"""Encode a series of line numbers as the difference from line to line (deltas) to reduce entropy.

	The delta is stored as a single signed byte per line, meaning no two consecutive lines may vary by more than +/-
	127 lines within the original source material. The resulting bytestring is then zlib compressed and b64 encoded.

	Lines without line numbers (none or unexpected zeros) will inherit the last known line number after decoding.
	"""

	def inner():
		lines = iter(numbers)
		prev = next(lines)

		for line in lines:
			delta = (line or prev) - prev  # Handle the "no line number given" case.
			if delta < 0: delta = -1 * delta + 127  # Store "signed" values.
			prev = (line or prev)  # Track our line number, or the last known good one.
			yield delta  # Store the delta.

	return b64encode(compress(bytes(bytearray(inner())))).decode('latin1')

def redelta_decode(source):
	"""Decode a series of line numbers encoded as the difference from line to line.

	The literal reverse of `redelta_encode`.
	"""

	pass

def chunk(line, mapping={None: 'text', '${': 'escape', '#{': 'bless', '&{': 'args', '%{': 'format', '@{': 'json'}):
	"""Chunkify and "tag" a block of text into plain text and code sections.

	The first delimeter is blank to represent text sections, and keep the indexes aligned with the tags.

	Values are yielded in the form (tag, text).
	"""

	skipping = 0  # How many closing parenthesis will we need to skip?
	start = None  # Starting position of current match.
	last = 0

	i = 0

	text = line.line

	while i < len(text):
		if start is not None:
			if text[i] == '{':
				skipping += 1

			elif text[i] == '}':
				if skipping:
					skipping -= 1
				else:
					yield line.clone(kind=mapping[text[start-2:start]], line=text[start:i])
					start = None
					last = i = i + 1
					continue

		elif text[i:i+2] in mapping:
			if last is not None and last != i:
				yield line.clone(kind=mapping[None], line=text[last:i])
				last = None

			start = i = i + 2
			continue

		i += 1

	if last < len(text):
		yield line.clone(kind=mapping[None], line=text[last:])


from irrhound.shared import WhoisProxy

def test_whois_expand_as():

     asn = 137

     objects = WhoisProxy.expand_as(asn)

     for route in objects:
         print("P: {} from source {}".format(route.cidr, route.source))

     assert (len(objects) > 0)


from components import Direction
from events import GameAction
import input
import logging
import utils

class BehaviourManager:
    """The Bahavoiur manager stores component behevoiurs."""
    logger = logging.getLogger('behaviour.BehaviourManager')

    def __init__(self, event_manager, entitiy_manager, config=None):
        self.event_manager = event_manager
        self.entity_manager = entitiy_manager
        self.behaviours = {}
        if config is not None:
            self.configure(config)

    def configure(self, config):
        """Configure behaviourse base on a configuration.

        Ars:
            config: {'component_behaviours': [{component_type, name: ''}]}

        """
        if 'component_behaviours' in config:
            for comp_behaviour_config in config['component_behaviours']:
                self.add_comp_behaviour_from_config(comp_behaviour_config)
        if 'min_turn_time' in config:
            self.min_turn_time = config['min_turn_time']

    def add_comp_behaviour_from_config(self, config):
        """Associates the behaviour specified be the name with the specified
        compoenent_type.

        Args:
            component_type (str): The component type associated with the
                behaviour
            behaviour_impl (str): The implementation of the behaviour.

        """
        impl = config['impl']
        behaviour_class = utils.class_for_name(impl[0], impl[1])
        behaviour = behaviour_class(self.event_manager, self.entity_manager)
        self.add_component_behaviour(config['component_type'], behaviour)

    def add_component_behaviour(self, component_type, behaviour):
        """Register the given behaviour with the specified component type."""
        self.logger.debug("Register %s for component type %s",
                          component_type, behaviour)
        self.behaviours.update({component_type: behaviour})

    def update(self, round, delta_time):
        """Updates all behaviours for all entitites with a associated
        components."""
        for component_type, behaviour in self.behaviours.items():
            components = self.entity_manager.get_all(component_type)
            for entity, component in components.items():
                behaviour.entity = entity
                behaviour.component = component
                behaviour.update(round, delta_time)

class TurnBehaviourManager(BehaviourManager):

    logger = logging.getLogger('behaviour.TurnBehaviourManager')

    def __init__(self, event_manager, entitiy_manager, config=None):
        BehaviourManager.__init__(self, event_manager,
                                  entitiy_manager, config)
        self.locked_entities = []

    def update(self, round, delta_time):
        """Updates all behaviours for all entitites with a associated
        component, but only if there are no locked entities."""

        # TODO: Make the method prettier ^^
        for component_type, behaviour in self.behaviours.items():
            components = self.entity_manager.get_all(component_type)

            for entity, component in components.items():
                s = self.locked_entities
                turn = self.entity_manager.get(entity, "Turn")
                if (turn.ticks == 0 and turn.delta >= turn.min_turn_time):
                    # It's the entity's Turn!
                    if (turn.locking and entity not in s):
                        # Lock it
                        self.locked_entities.append(entity)
                    if (entity in s or len(s) == 0):
                        behaviour.entity = entity
                        behaviour.component = component
                        ticks = behaviour.update(round, delta_time)
                        if (ticks is not None):
                            # The behaviour has made an action
                            if (entity in s):
                                self.locked_entities.remove(entity)
                            self._add_ticks(entity, ticks)
                else:
                    turn.delta += delta_time
        self._normalize()

    def _add_ticks(self, entity, ticks):
        """ Count up the Tick Counter of entity, and reset its delta time """
        turn_comp = self.entity_manager.get(entity, "Turn")
        turn_comp.ticks += ticks
        turn_comp.delta = 0

    def _normalize(self):
        """ Reduce the Tick counter of all turn components by the minimum
        Tick count, to prevent tick counters from increasing indefinitely.
        This assures an only relative tick count """
        min_ticks = None
        for component_type, behaviour in self.behaviours.items():
            components = self.entity_manager.get_all(component_type)
            ticks = [self.entity_manager.get(entity, "Turn").ticks
                     for entity in components.keys()]
            if (len(ticks) > 0):
                minimum_of_type = min(ticks)
                if (min_ticks is None or min_ticks > minimum_of_type):
                    min_ticks = minimum_of_type
        if min_ticks is not None:
            for component_type, behaviour in self.behaviours.items():
                components = self.entity_manager.get_all(component_type)
                for turn_comp in [self.entity_manager.get(entity, "Turn")
                                  for entity in components.keys()]:
                    turn_comp.ticks -= min_ticks

class EntityComponentBehaviour(object):
    """Implements logic for entities with specific components."""

    def __init__(self, event_manager, entitiy_manager):
        self.entity_manager = entitiy_manager
        self.event_manager = event_manager
        self.entity = None
        self.component = None

    def update(self, round, delta_time):
        pass

class InputBehaviour(EntityComponentBehaviour):
    """Implements User Input. Controls all entites with an InputComponent."""

    logger = logging.getLogger('behaviours.InputBehaviour')

    def __init__(self, event_manager, entitiy_manager):
        super(InputBehaviour, self).__init__(event_manager, entitiy_manager)
        # TODO make configurable
        self.isometric = True

    def update(self, round, delta_time):
        """Converts input to an appropriate InputAction."""
        # TODO Implement gamestatus aware behaviour manager to keep the turn
        # based logic as central as possible and possibly enable switching
        # between realtime and turn based. (A realtime behaviour would not
        # check for a state
        speed = self.move()

        if input.is_pressed(input.K_ENTER):
            # TODO: Implement Targeting Inputs (combined input or
            # sequential?)
            # TODISCUSS: Context actions, autotargeting
            speed = self.use(0, 0)
        return speed

    def set_input_direction(self):
        direction = self.component.direction
        if input.is_pressed(
            input.K_LEFT) or input.is_pressed(
            input.K_KP1) or input.is_pressed(
            input.K_KP4) or input.is_pressed(
                input.K_KP7):
            direction.set(Direction.D_LEFT)
        else:
            direction.set(Direction.D_LEFT, 0)
        if input.is_pressed(
            input.K_RIGHT) or input.is_pressed(
            input.K_KP3) or input.is_pressed(
            input.K_KP6) or input.is_pressed(
                input.K_KP9):
            direction.set(Direction.D_RIGHT)
        else:
            direction.set(Direction.D_RIGHT, 0)
        if input.is_pressed(
            input.K_DOWN) or input.is_pressed(
            input.K_KP1) or input.is_pressed(
            input.K_KP2) or input.is_pressed(
                input.K_KP3):
            direction.set(Direction.D_DOWN)
        else:
            direction.set(Direction.D_DOWN, 0)
        if input.is_pressed(
            input.K_UP) or input.is_pressed(
            input.K_KP7) or input.is_pressed(
            input.K_KP8) or input.is_pressed(
                input.K_KP9):
            direction.set(Direction.D_UP)
        else:
            direction.set(Direction.D_UP, 0)

    def set_iso_input_direction(self):
        direction = self.component.direction
        if input.is_pressed(
            input.K_LEFT) or input.is_pressed(
            input.K_UP) or input.is_pressed(
            input.K_KP4) or input.is_pressed(
            input.K_KP7) or input.is_pressed(
                input.K_KP8):
            direction.set(Direction.D_LEFT)
        else:
            direction.set(Direction.D_LEFT, 0)
        if input.is_pressed(
            input.K_RIGHT) or input.is_pressed(
            input.K_DOWN) or input.is_pressed(
            input.K_KP2) or input.is_pressed(
            input.K_KP3) or input.is_pressed(
                input.K_KP6):
            direction.set(Direction.D_RIGHT)
        else:
            direction.set(Direction.D_RIGHT, 0)
        if input.is_pressed(
            input.K_LEFT) or input.is_pressed(
            input.K_DOWN) or input.is_pressed(
            input.K_KP4) or input.is_pressed(
            input.K_KP1) or input.is_pressed(
                input.K_KP2):
            direction.set(Direction.D_DOWN)
        else:
            direction.set(Direction.D_DOWN, 0)
        if input.is_pressed(
            input.K_RIGHT) or input.is_pressed(
            input.K_UP) or input.is_pressed(
            input.K_KP6) or input.is_pressed(
            input.K_KP8) or input.is_pressed(
                input.K_KP9):
            direction.set(Direction.D_UP)
        else:
            direction.set(Direction.D_UP, 0)

    def move(self):
        """Throws a MoveAction."""
        if self.isometric:
            self.set_iso_input_direction()
        else:
            self.set_input_direction()
        return 1 if self.component.direction.direction > 0 else None

    def use(self, dx, dy):
        """ Throws a UseEntity Event """
        event = self.event_manager.create(GameAction.UseEntityAction)
        event.user = self.entity
        self.event_manager.throw(event)
        return 2


import contextlib
import logging
import os.path as osp

import numpy as np

@contextlib.contextmanager
def np_print_options(*args, **kwargs):
    original = np.get_printoptions()
    np.set_printoptions(*args, **kwargs)
    yield
    np.set_printoptions(**original)

def set_logger(output_dir=None, log_file=None, debug=False):
    head = '%(asctime)-15s Host %(message)s'
    logger_level = logging.INFO if not debug else logging.DEBUG
    if all((output_dir, log_file)) and len(log_file) > 0:
        logger = logging.getLogger('crosscityadap')
        log_path = osp.join(output_dir, log_file)
        handler = logging.FileHandler(log_path)
        formatter = logging.Formatter(head)
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        handler = logging.StreamHandler()
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        logger.setLevel(logger_level)
    else:
        logging.basicConfig(level=logger_level, format=head)
        logger = logging.getLogger('crosscityadap')
    return logger


#Importing the required libraries
import nltk
import random
from nltk.tokenize import word_tokenize
import sys
import re
import math

# Initializing n-gram dictionaries upto n=4.
unigram = {}
bigram = {}
trigram = {}
fourgram = {}

# Cleaning up text suitably and tokenizing it
def preprocess(text):
    # the function prpocess a single line
    # it takes the single line and returns a list of tokens for that particular line
    # make text lower
    cleaned_text = text.lower()
    # remove non-ASCII characters
    #cleaned_text = re.sub(r'[^\x00-\x7F]+',' ', cleaned_text)
    # remove URLS
    cleaned_text = re.sub(r"http\S+", "<URL>", cleaned_text)
    # remove HTs
    cleaned_text = re.sub(r"#[A-Za-z0-9_]+", "<HASHTAG>", cleaned_text)
    # remove Mentions
    cleaned_text = re.sub(r"@[A-Za-z0-9_]+", "<MENTION>", cleaned_text)
    # replace percentage quantities with tags
    cleaned_text = re.sub(r'(\d+(\.\d+)?%)', "<PERCENT>", cleaned_text)
    # replace numbers with tags
    cleaned_text = re.sub("^\d+\s|\s\d+\s|\s\d+$", " <NUM> ", cleaned_text)
    # cleaned_text = re.sub(r'[0-9]', " <NUM> ", cleaned_text)
    # hypenated words are accounted for by joining them/merging them together
    cleaned_text = re.sub(r'\w+(?:-\w+)+', '', cleaned_text)
    # Substitue for punctuations
    cleaned_text = re.sub(r"(\'t)", " not", cleaned_text)
    cleaned_text = re.sub(r'(i\'m)', "i am", cleaned_text)
    cleaned_text = re.sub(r'(ain\'t)', "am not", cleaned_text)
    cleaned_text = re.sub(r'(\'ll)', " will", cleaned_text)
    cleaned_text = re.sub(r'(\'ve)', " have", cleaned_text)
    cleaned_text = re.sub(r'(\'re)', " are", cleaned_text)
    cleaned_text = re.sub(r'(\'s)', " is", cleaned_text)
    cleaned_text = re.sub(r'(\'re)', " are", cleaned_text)
    # removing repetetive spam
    cleaned_text = re.sub('\!\!+', '!', cleaned_text)
    cleaned_text = re.sub('\*\*+', '*', cleaned_text)
    cleaned_text = re.sub('\>\>+', '>', cleaned_text)
    cleaned_text = re.sub('\<\<+', '<', cleaned_text)
    cleaned_text = re.sub('\?\?+', '?', cleaned_text)
    cleaned_text = re.sub('\!\!+', '!', cleaned_text)
    cleaned_text = re.sub('\.\.+', '.', cleaned_text)
    cleaned_text = re.sub('\,\,+', ',', cleaned_text)
    cleaned_text = re.sub('\:\:+', ':', cleaned_text)
    cleaned_text = re.sub('\;\;+', ';', cleaned_text)
    # matching punctuation characters at end of sentences and padding them
    cleaned_text = re.sub('([;:.,!?()])', r' \1 ', cleaned_text)
    # removing multiple spaces finally
    cleaned_text = re.sub('\s{2,}', ' ', cleaned_text)
    # remove trailing white spaces
    # important to get rid of empty tokens at the end of list
    cleaned_text = re.sub(r'\s+$', '', cleaned_text)
    # tokenization based on spaces for each line
    spaces = r"\s+"
    tokenized_sent = re.split(spaces, cleaned_text)
    return tokenized_sent

#how to randomly choose from a list
def random_choice(list):
    return random.choice(list)

def convert_some_to_unk(tokenized_sentences):
    vocab = {}
    for sent in tokenized_sentences:
        for token in sent:
            if token in vocab.keys():
                vocab[token] += 1
            else:
                vocab[token] = 1
    #to_be_changed = int(ratio * len(vocab))
    vocab = {k: v for k, v in sorted(vocab.items(), key=lambda item: item[1])}
    to_be_changed = 0
    for key in vocab.keys():
        if vocab[key] == 1:
            to_be_changed += 1
    words_to_be_changed = list(vocab.keys())[0:to_be_changed]
    # temp_sent = []
    # temp_tokenized_sents = []
    # for sent in tokenized_sentences:
    #     for word in sent:
    #         if word in words_to_be_changed:
    #             temp_sent.append("<UNK>")
    #         else:
    #             temp_sent.append(word)
    #     temp_tokenized_sents.append(temp_sent)
    # tokenized_sentences = temp_tokenized_sents
    # return tokenized_sentences
    for i,sent in enumerate(tokenized_sentences):
        for j,token in enumerate(sent):
            if token in words_to_be_changed:
                tokenized_sentences[i][j] = '<UNK>'
    return tokenized_sentences

#Creating n-gram dictionaries for each n
def ngram_dict(data):
    token_sent = []
    for sent in data:
        token_sent.append(preprocess(sent))
    token_sent = convert_some_to_unk(token_sent)

    for sent in token_sent:
        for token in sent:
            if token in unigram.keys():
                unigram[token] += 1
            else:
                unigram[token] = 1

    for sent in token_sent:
        for i in range(len(sent)-1):
            if sent[i] not in bigram.keys():
                # If a key isnt present, u make that into a dictionary to work with as well
                bigram[sent[i]] = {}
            if sent[i+1] not in bigram[sent[i]].keys():
                # Now that a bigram dict is also made, u append the same way as u did for unigrams
                bigram[sent[i]][sent[i+1]] = 1
            else:
                bigram[sent[i]][sent[i+1]] += 1

    for sent in token_sent:
        for i in range(len(sent)-2):
            if sent[i] not in trigram.keys():
                # If a key isnt present, u make that into a dictionary to work with as well
                trigram[sent[i]] = {}
            if sent[i+1] not in trigram[sent[i]].keys():
                # Now that a bigram dict is also made, u append the same way as u did for unigrams
                trigram[sent[i]][sent[i+1]] = {}
            if sent[i+2] not in trigram[sent[i]][sent[i+1]].keys():
                trigram[sent[i]][sent[i+1]][sent[i+2]] = 1
            else:
                trigram[sent[i]][sent[i+1]][sent[i+2]] += 1

    for sent in token_sent:
        for i in range(len(sent)-3):
            if sent[i] not in fourgram.keys():
                # If a key isnt present, u make that into a dictionary to work with as well
                fourgram[sent[i]] = {}
            if sent[i+1] not in fourgram[sent[i]].keys():
                # Now that a bigram dict is also made, u append the same way as u did for unigrams
                fourgram[sent[i]][sent[i+1]] = {}
            if sent[i+2] not in fourgram[sent[i]][sent[i+1]].keys():
                fourgram[sent[i]][sent[i+1]][sent[i+2]] = {}
            if sent[i+3] not in fourgram[sent[i]][sent[i+1]][sent[i+2]].keys():
                fourgram[sent[i]][sent[i+1]][sent[i+2]][sent[i+3]] = 1
            else:
                fourgram[sent[i]][sent[i+1]][sent[i+2]][sent[i+3]] += 1

#Implementing Kneyser-Ney for n-grams
## Implementing for different discount values and accounting for high order as well
def kneyser_ney(n, n_gram, high_ord=True):
    if n == 1:
        d_1 = 0.5
        if high_ord:
            if (unigram[n_gram[0]]) > 0:
                # since it is higher order, discount factor = 0
                numer = unigram[n_gram[0]]
            else:
                numer = 0
            # Total frequency count of the all unigrams
            denom = sum([item[1] for item in unigram.items()])
            return numer/denom  # terms after lambda are zero, so no P_cont here
        else:
            # now we include pcont, so here:
            # num = different string types before the final word
            # denom = the number of different possible n-gram types
            p_cont = len(set([item[0] for item in bigram.items() if n_gram[0] in item[1].keys()]))
            denom_cont = len(bigram.keys())
            return p_cont/denom_cont
    if n == 2:
        d_2 = 0.75
        lambda_2 = (d_2 * len(bigram[n_gram[0]]))/(sum([item[1] for item in bigram[n_gram[0]].items()]))
        if high_ord:
            if n_gram[1] in bigram[n_gram[0]].keys():
                num_2 = bigram[n_gram[0]][n_gram[1]]
            else:
                num_2 = 0
            # total frequency count
            return ((max(num_2-d_2, 0))/sum([item[1] for item in bigram[n_gram[0]].items()])) + (lambda_2*kneyser_ney(1, n_gram[1:], False))
        else:
            num_2 = 0
            for first in trigram.keys():
                if n_gram[0] in trigram[first].keys():
                    if n_gram[1] in trigram[first][n_gram[0]].keys():
                        num_2 += 1
            return (max(num_2-d_2, 0)/len(trigram.keys()))+(lambda_2*kneyser_ney(1, n_gram[1:], False))
    if n==3:
        d_3 = 0.9
        try:
            lambd = (d_3 * len(trigram[n_gram[0]][n_gram[1]]))/(sum([item[1] for item in trigram[n_gram[0]][n_gram[1]].items()]))
        except:
            return d_3*kneyser_ney(2,n_gram[1:],False)
        count = trigram[n_gram[0]][n_gram[1]][n_gram[2]] if n_gram[2] in trigram[n_gram[0]][n_gram[1]].keys() else 0
        among = sum([item[1] for item in trigram[n_gram[0]][n_gram[1]].items()])
        return (max(0,count-d_3)/among)+(lambd*kneyser_ney(2,n_gram[1:],False))
    if n==4:
        d_4 = 0.9
        try:
            lambd = (d_4 * len(fourgram[n_gram[0]][n_gram[1]][n_gram[2]]))/(sum([item[1] for item in fourgram[n_gram[0]][n_gram[1]][n_gram[2]].items()]))
        except:
            return d_4*kneyser_ney(3,n_gram[1:],False)
        count = fourgram[n_gram[0]][n_gram[1]][n_gram[2]][n_gram[3]] if n_gram[3] in fourgram[n_gram[0]][n_gram[1]][n_gram[2]].keys() else 0
        among = sum([item[1] for item in fourgram[n_gram[0]][n_gram[1]][n_gram[2]].items()])
        return (max(0,count-d_4)/among)+(lambd*kneyser_ney(3,n_gram[1:],False))

#Implementing Witten Bell algorithm
def witten_bell(n, n_gram, high_ord= True):
    if n == 1:
        #We calculate the standard probility of the n_gram, called "original prob"
        num = unigram[n_gram[0]]
        denom = sum([item[1] for item in unigram.items()])
        return (num/denom)
    if n == 2:
        no_tokens = len(bigram[n_gram[0]])
        types = sum([item[1] for item in bigram[n_gram[0]].items()])
        plus_lambda = (types/(no_tokens+types))
        minus_lambda = (no_tokens/(no_tokens+types))
        if n_gram[1] in bigram[n_gram[0]].keys(): #if seen in higher order model, we can use it
            return (plus_lambda)*(bigram[n_gram[0]][n_gram[1]]/types) + (minus_lambda)*(witten_bell(1, n_gram[1:]))
        else:
            return minus_lambda*witten_bell(1, n_gram[1:])
    if n == 3:
        # if n_gram[1] in trigram[n_gram[0]].keys():
        try:
            no_tokens = len(trigram[n_gram[0]][n_gram[1]])
            types = sum([item[1] for item in trigram[n_gram[0]][n_gram[1]].items()])
            plus_lambda = (types/(no_tokens+types))
            minus_lambda = (no_tokens/(no_tokens+types))
            if n_gram[2] in trigram[n_gram[0]][n_gram[1]].keys():
                return (plus_lambda)*(trigram[n_gram[0]][n_gram[1]][n_gram[2]]/types) + (minus_lambda)*(witten_bell(2, n_gram[1:]))
            else:
                return minus_lambda*witten_bell(2, n_gram[1:])
        # else:
        except:
            return 0.9*witten_bell(2, n_gram[1:])
    if n == 4:
        # if n_gram[1] in fourgram[n_gram[0]].keys():
        try:
            try:
                # if n_gram[2] in fourgram[n_gram[0]][n_gram[1]].keys():
                    no_tokens = len(fourgram[n_gram[0]][n_gram[1]][n_gram[2]])
                    types = sum([item[1] for item in fourgram[n_gram[0]][n_gram[1]][n_gram[2]].items()])
                    plus_lambda = (types/(no_tokens+types))
                    minus_lambda = (no_tokens/(no_tokens+types))
                    if n_gram[3] in fourgram[n_gram[0]][n_gram[1]][n_gram[2]].keys():
                        return (plus_lambda)*(fourgram[n_gram[0]][n_gram[1]][n_gram[2]][n_gram[3]]/types) + (minus_lambda)*(witten_bell(3, n_gram[1:]))
                    else:
                        return minus_lambda*witten_bell(3, n_gram[1:])
            except KeyError:
                return 0.9*witten_bell(3, n_gram[1:])
        # else:
        except KeyError:
            return 0.9*witten_bell(3, n_gram[1:])

#calculating perplexity using an array of probabilities
def perplexity(prob_array):
    log_prob = 0
    for i in range(len(prob_array)):
        log_prob += math.log(prob_array[i])
    return math.exp(-log_prob/len(prob_array))

### INTERACTIVE PART OF CODE ###
# check if all arguments are given
if len(sys.argv) != 4:
    print("please provide all the arguments")
    exit()

# obtain n, model type and path from the first argument
num = int(sys.argv[1])
model = sys.argv[2]
path = sys.argv[3]

# Making sure that n is valid else we set it to 4(default)
if num > 4:
    n = 4
else:
    n = num

# Choosing the model type based on the input
if model == "k":
    smooth_model = kneyser_ney
elif model == "w":
    smooth_model = witten_bell
else:
    print("Please provide a valid model type")
    exit()

# # reading file and storing the data
# with open(path) as file:
#     fdata = file.readlines()

#####################################################################333333333333333##########
#                        TRAIN - TEST CODE
#####################################################################333333333333333##########
#shuffling data in a list
# with open(path) as file:
#     fdata = file.readlines()
# random.shuffle(fdata)
# # Choosing training and testing data accordingly
# training_data = fdata[:len(fdata)-1000]
# testing_data = fdata[len(fdata)-1000:]
# ngram_dict(training_data)
# final_per = []
# final_sent = []
# count = 0
# print("Training dict is ready")
# for sentence in training_data:
#     print(sentence)
#     print(count)
#     count += 1
#     real_sent =  re.sub(r'\n', '', sentence)
#     sentence = preprocess(sentence)
#     if len(sentence) < 4:
#         for i in range(4-len(sentence)):
#             sentence.append("<UNK>")
#     sen_temp = []
#     for i in sentence:
#         if i in unigram.keys():
#             sen_temp.append(i)
#         else:
#             sen_temp.append("<UNK>")
#     sentence = sen_temp
#     perplex = []
#     for i in range(len(sentence)-n+1):
#         prob = smooth_model(n,sentence[i:i+n])
#         perplex.append(prob)
#     # print("Final perplexity:", perplexity(perplex))
#     final_sent.append(real_sent)
#     final_per.append(perplexity(perplex))
# #python code to find average of a list
# def mean(numbers):
#     return float(sum(numbers)) / max(len(numbers), 1)

# file1 = open("2020114017_LM1_train-perplexity.txt","w")#write mode
# file1.write(str(mean(final_per)) + "\n")
# for i in range(len(final_sent)):
#     file1.write(final_sent[i]+'\t'+str(final_per[i])+'\n')
# file1.close()

# # testing the model
# final_per = []
# final_sent = []
# count = 0
# for sentence in testing_data:
#     print(sentence)
#     print(count)
#     count += 1
#     real_sent =  re.sub(r'\n', '', sentence)
#     sentence = preprocess(sentence)
#     if len(sentence) < 4:
#         for i in range(4-len(sentence)):
#             sentence.append("<UNK>")
#     sen_temp = []
#     for i in sentence:
#         if i in unigram.keys():
#             sen_temp.append(i)
#         else:
#             sen_temp.append("<UNK>")
#     sentence = sen_temp
#     perplex = []
#     for i in range(len(sentence)-n+1):
#         prob = smooth_model(n,sentence[i:i+n])
#         perplex.append(prob)
#     # print("Final perplexity:", perplexity(perplex))
#     final_sent.append(real_sent)
#     final_per.append(perplexity(perplex))
# #python code to find average of a list
# def mean(numbers):
#     return float(sum(numbers)) / max(len(numbers), 1)

# file1 = open("t1.txt","w")#write mode
# file1.write(str(mean(final_per)) + "\n")
# for i in range(len(final_sent)):
#     file1.write(final_sent[i]+'\t'+str(final_per[i])+'\n')
# file1.close()
# print("Over")

#####################################################################333333333333333##########
#                                SINGLE SENTENCE
#####################################################################333333333333333##########

# # generating the uni,bi,tri,four gram dictionaries
with open(path) as file:
    fdata = file.readlines()
ngram_dict(fdata)
print("Trained")

sentence = input("input sentence: ")
sentence = preprocess(sentence)
sen_temp = []
for i in sentence:
    if i in unigram.keys():
        sen_temp.append(i)
    else:
        sen_temp.append("<UNK>")
sentence = sen_temp

#Calculating perplexity of the Input sentence.
perplex = []
prob_final = 1
for i in range(len(sentence)-n+1):
    prob = smooth_model(n,sentence[i:i+n])
    print("For",sentence[i:i+n],'probability is:',prob)
    perplex.append(prob)
    prob_final *= prob
print("Final Probability:", prob_final)
print("Final Perplexity:", perplexity(perplex))


import torch
import numpy
from dataset import PromiseDataset
from data import Data
from model import VNet
import torch.backends.cudnn as cudnn
import SimpleITK as sitk
import os

def write_result(result, path, new_spacing = [1, 1, 1.5], volume_size = [128, 128, 64]):
    original_image = sitk.ReadImage(path)
    to_write = sitk.Image(original_image.GetSize()[0], original_image.GetSize()[1], original_image.GetSize()[2], sitk.sitkFloat32)
    factor = numpy.asarray(original_image.GetSpacing()) / [new_spacing[0], new_spacing[1], new_spacing[2]]
    factor_size = numpy.asarray(original_image.GetSize() * factor, dtype=float)
    new_size = numpy.max([factor_size, volume_size], axis = 0)
    new_size = new_size.astype(dtype=int)

    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(original_image)
    resampler.SetOutputSpacing(new_spacing)
    resampler.SetSize(new_size.tolist())
    resampler.SetInterpolator(sitk.sitkNearestNeighbor)

    to_write = resampler.Execute(to_write)
    image_centroid = numpy.asarray(new_size, dtype = float) / 2.0
    image_start_pixel = (image_centroid - numpy.asarray(volume_size) / 2.0).astype(dtype = int)

    for dst_x, src_x in zip(range(0, result.shape[0]), range(image_start_pixel[0], int(image_start_pixel[0] + volume_size[0]))):
        for dst_y, src_y in zip(range(0, result.shape[1]), range(image_start_pixel[1], int(image_start_pixel[1] + volume_size[1]))):
            for dst_z, src_z in zip(range(0, result.shape[2]), range(image_start_pixel[2], int(image_start_pixel[2] + volume_size[2]))):
                to_write.SetPixel(int(src_x), int(src_y), int(src_z), float(result[dst_x, dst_y, dst_z]))

    resampler.SetOutputSpacing([original_image.GetSpacing()[0], original_image.GetSpacing()[1], original_image.GetSpacing()[2]])
    resampler.SetSize(original_image.GetSize())
    to_write = resampler.Execute(to_write)
    threshold_filter = sitk.BinaryThresholdImageFilter()
    threshold_filter.SetInsideValue(1)
    threshold_filter.SetOutsideValue(0)
    threshold_filter.SetLowerThreshold(0.5)
    to_write = threshold_filter.Execute(to_write)

    connected_component = sitk.ConnectedComponentImageFilter()
    to_write_connected_component = connected_component.Execute(sitk.Cast(to_write, sitk.sitkUInt8))
    connected_component_array = numpy.transpose(sitk.GetArrayFromImage(to_write_connected_component).astype(dtype = float), [2, 1, 0])
    lab = numpy.zeros(int(numpy.max(connected_component_array) + 1), dtype = float)
    for i in range(1, int(numpy.max(connected_component_array) + 1)):
        lab[i] = numpy.sum(connected_component_array == i)
    active_lab = numpy.argmax(lab)
    to_write = (to_write_connected_component == active_lab)
    to_write = sitk.Cast(to_write, sitk.sitkUInt8)
    writer = sitk.ImageFileWriter()
    patient_name = path.split("/")[-1]
    writer.SetFileName(os.path.join("result", patient_name))
    writer.Execute(to_write)

def inference(model):
    model.eval()
    promise_dataset = PromiseDataset(is_train = False)
    test_loader = torch.utils.data.DataLoader(dataset = promise_dataset, batch_size = 1)
    with torch.no_grad():
        for batch_index, data in enumerate(test_loader):
            data, path = data
            pred = model(data)
            _, _, z, y, x = data.shape
            result = model(data)
            #tuple to string
            path = "".join(path)
            write_result(result, path)

if __name__ == "__main__":
    net = VNet()
    net = torch.nn.DataParallel(net, device_ids=[0, 1, 2, 3])
    net.cuda()
    cudnn.benchmark = True
    net.load_state_dict(torch.load("weights/promise12_weight.pth"))
    inference(net)


from openerp.osv import fields
from openerp.osv.osv import Model

class Character(Model):
    _name = 'ro.character'
    _description = 'Ragnarok Online Base Character'

    _columns = {
        'name': fields.char('Name of the character'),
        'type': fields.selection([
            ('npc', 'NPC'),
            ('monster', 'Monster'),
            ('player', 'Player')
        ], 'Character type'),
    }