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


from django.test import Client, TestCase
from django.urls import reverse
from django_dynamic_fixture import G

from animal_crossing_new_horizons.acnh_api.models import Villager, Fossil, Fish, Insect

class APITestVillagers(TestCase):
    def test_villagers_api(self):
        G(Villager, name='tom')
        G(Villager, name='jerry')

        c = Client()

        response = c.get(reverse('villagers-list'))

        self.assertEqual(2, len(response.json()['results']))

    def test_villagers_api_filters(self):
        G(Villager, name='villager1', personality='personality1')
        G(Villager, name='villager2', personality='personality2')

        c = Client()

        filters = {
            'personality': 'personality1'
        }
        response = c.get(reverse('villagers-list'), data=filters)
        results = response.json()['results']

        self.assertEqual(1, len(results))
        self.assertEqual('villager1', results[0]['name'])

class APITestFossils(TestCase):
    def test_fossils_api(self):
        G(Fossil, name='fossil1')
        G(Fossil, name='fossil2')

        c = Client()

        response = c.get(reverse('fossils-list'))

        self.assertEqual(2, len(response.json()['results']))

    def test_fossils_api_filters(self):
        G(Fossil, name='fossil1', in_catalog=True)
        G(Fossil, name='fossil2', in_catalog=False)

        c = Client()

        filters = {
            'in_catalog': True
        }
        response = c.get(reverse('fossils-list'), data=filters)
        results = response.json()['results']

        self.assertEqual(1, len(results))
        self.assertEqual('fossil1', results[0]['name'])

class APITestFish(TestCase):
    def test_fish_api(self):
        G(Fish, name='fish1')
        G(Fish, name='fish2')

        c = Client()

        response = c.get(reverse('fish-list'))

        self.assertEqual(2, len(response.json()['results']))

    def test_fish_api_filters(self):
        G(Fish, name='fish1', sell_price=200)
        G(Fish, name='fish2', sell_price=500)

        c = Client()

        filters = {
            'sell_price': '200'
        }
        response = c.get(reverse('fish-list'), data=filters)
        results = response.json()['results']

        self.assertEqual(1, len(results))
        self.assertEqual('fish1', results[0]['name'])

class APITestInsects(TestCase):
    def test_insects_api(self):
        G(Insect, name='insect1')
        G(Insect, name='insects2')

        c = Client()

        response = c.get(reverse('insects-list'))

        self.assertEqual(2, len(response.json()['results']))

    def test_insects_api_filters(self):
        G(Insect, name='insect1', sell_price=200)
        G(Insect, name='insect2', sell_price=500)

        c = Client()

        filters = {
            'sell_price': '200'
        }
        response = c.get(reverse('insects-list'), data=filters)
        results = response.json()['results']

        self.assertEqual(1, len(results))
        self.assertEqual('insect1', results[0]['name'])


#-*- encoding: UTF-8 -*-
import sys
from douban.books_api import BooksAPI
from douban.movies_api import MoviesAPI
from clint.arguments import Args
from clint.textui import puts, colored, indent
SPACE = ' '

class DouBanCLI(object):

    @property
    def valid_actions(self):
        return ['book', 'movie']

    # flags are not used yet
    def route(self, action, name, flags):
        if not action in self.valid_actions:
            puts(colored.red('%s is not a valid action' %action))
            exit()

        getattr(self, action)(name)

    def book(self, book_name):
        api = BooksAPI()
        res = api.search_book({'q': book_name, 'count': 1})
        book = self._get_first_matched_id_from(res, 'book')
        self._print_book_info(book)

    def movie(self, movie_name):
        api = MoviesAPI()
        res = api.search_movie({'q': movie_name, 'count': 1})
        movie = self._get_first_matched_id_from(res, 'movie')
        self._print_movie_info(movie)

    def _get_first_matched_id_from(self, res, obj_type='book'):
        count = int(res['total'])
        if count == 0:
            print(u'找不到结果')
            exit()
        else:
            print(u'共%s个结果，返回第1个结果' %res['total'])
            mapper = {'book': 'books', 'movie': 'subjects'}
            key = mapper[obj_type]
            if len(res[key]) >= 1:
                return res[key][0]

    def _print_book_info(self, book):
        if book != None:
            # print(book)
            print(u'')
            print(u'书名: %s' % colored.green(book['title']))
            print(u'')
            print(u'作者: %s' % colored.green(SPACE.join(book['author'])))
            print(u'作者简介: %s' % colored.green(book['author_intro']))
            print(u'')
            print(u'页数: %s' % book['pages'])
            print(u'出版日期: %s' % book['pubdate'])
            print(u'定价: %s' % book['price'])
            self._print_rating(book['rating'])
            self._print_tags(book['tags'])
            if sys.platform == 'win32':
                print(u'简介')
                print(colored.yellow(book['summary'].encode('GBK', 'ignore')))
            else:
                print('简介: %s' %colored.yellow(book['summary']))
            print('')
            # print(u'目录: %s' % book['catalog'])
        else:
            print(u'无法找到相关结果')

    def _print_rating(self, rating, obj_type='book'):
        if obj_type == 'book':
            print(u'共有%s人次评分, 平均分: %s'%(colored.yellow(rating['numRaters']), colored.yellow(rating['average'])))
        else:
            print(u'共有%s人收藏, 平均分: %s'%(colored.yellow(rating['stars']), colored.yellow(rating['average'])))

    def _print_tags(self, tags):
        tag_names = [tag['name'] for tag in tags]
        print(u'标签: %s' % u', '.join(tag_names))

    def _print_movie_info(self, movie):
        if movie != None:
            print(u'名称: %s' % colored.yellow(movie['title']))
            self._print_rating(movie['rating'], 'movie')
            print(u'年份: %s' % movie['year'])
            print(u'类型: %s' % movie['subtype'])
            print(u'导演: %s' % self._get_name(movie['directors']))
            print(u'主演: %s' % self._get_name(movie['casts']))
            print(u'分类: %s' % ' '.join(movie['genres']))
        else:
            print(u'无法找到相关结果')

    def _get_name(self, directors_or_actors):
        return u', '.join([someone['name'] for someone in directors_or_actors])

def fetch_names(grouped_args):
    return SPACE.join(grouped_args[1:])

def main():

    cli = DouBanCLI()

    args = Args()

    # grouped['_'] does not contain flags
    grouped_args = args.grouped['_'].all

    if len(grouped_args) <= 1:
        print(colored.yellow('USAGE: douban book book_name'))
        print(colored.yellow('USAGE: douban movie movie_name'))
    else:
        action = grouped_args[0]
        name = fetch_names(grouped_args)
        show_detail = True if '-d' in args.flags.all else False
        cli.route(action, name, {'show_detail': show_detail})

if __name__ == '__main__':
    main()


from django.shortcuts import render
from players.models import playerdata
# Create your views here.
from django.http import JsonResponse
import json
from .models import updates,matches,leagues,allaccouncements,StateAssociations,BoardManageMent,SelectionCommittee,AthletesCommission
from django.core import serializers
from django.views.decorators.csrf import csrf_exempt
from django.http import HttpResponse

@csrf_exempt
def getPlayers(request):
	if request.method == "GET":
		information = playerdata.objects.all()
		f=["playerid","team","firstname","lastname","fathername","address","phonenumber","Email","gender","age","city","state","photo","adhar","birth","Height","Goals","Wins","Nationality","JersyNo","Weight","GamesNo","Losses"]
		jsoninformation = serializers.serialize('python', information)
		actual_data = [d['fields'] for d in jsoninformation]
		# jsoninformation.insert(0,{"result":len(jsoninformation),"msg":"DataFound"})
		# print(type(actual_data))
		# actual_data["msg"] = "Data Found"
		information = {"msg" :"DataFound","status_code":200,"result":len(jsoninformation),"data":actual_data}
		output = json.dumps(information)

		# return JsonResponse({"players":output})
		return HttpResponse(output, content_type='application/json')
		# serialized_queryset = serializers.serialize('json', some_queryse
@csrf_exempt
def getPlayer(request,play):
	if request.method == "GET":
		player = [playerdata.objects.get(playerid =play )]

		jsoninformation = serializers.serialize('json', player)
		obj = json.loads(json.loads(json.dumps(jsoninformation)))
		output = json.dumps(obj)
		return HttpResponse(output, content_type='application/json')
	if request.method == "DELETE":
		try:
			player = playerdata.objects.get(playerid =play )
			player.delete()
		except:
			return JsonResponse({"status":"player Not Found"})
		return JsonResponse({"status":"success"})

	if request.method == "POST":
		# print(request['body'])
		# print(request.body)
		jsom = json.loads(request.body)
		print(jsom.data)
		return JsonResponse({"om":"om"})

@csrf_exempt
def latestUpdates(request):
	if request.method=="GET":
		up = updates.objects.all()
		# jsoninformation = serializers.serialize('json', up)
		jsoninformation = serializers.serialize('python', up)
		actual_data = [d['fields'] for d in jsoninformation]
		actual_data.insert(0,{"result":len(actual_data)})## = f"{len(actual_data)}"

		actual_data.insert(0,{"msg":"data found"})
		# actual_data["result"] = f"{len(actual_data)}"

		# actual_data = [d['fields'] for d in jsoninformation]
		output = json.dumps(actual_data)
		# return JsonResponse({"players":output})
		return HttpResponse(output, content_type='application/json')
		# return JsonResponse({"data":jsoninformation})

@csrf_exempt
def StateViseAssociations(request,state):
	if request.method=="GET":
		assosiations = StateAssociations.objects.all().filter(state=state)
		jsoninformation = serializers.serialize('python', assosiations)
		actual_data = [d['fields'] for d in jsoninformation]
		output = json.dumps(actual_data)
		# return JsonResponse({"players":output})
		return HttpResponse(output, content_type='application/json')
		# return JsonResponse({"data":jsoninformation})

@csrf_exempt
def StateAssociationsAll(request):
	if request.method=="GET":
		assosiations = StateAssociations.objects.all()
		jsoninformation = serializers.serialize('python', assosiations)
		actual_data = [d['fields'] for d in jsoninformation]
		output = json.dumps(actual_data)
		# return JsonResponse({"players":output})
		return HttpResponse(output, content_type='application/json')
		# return JsonResponse({"data":jsoninformation})

@csrf_exempt
def getAllBoardMembers(request):
	if request.method=="GET":
		board = BoardManageMent.objects.all()
		jsoninformation = serializers.serialize('python', board)
		actual_data = [d['fields'] for d in jsoninformation]
		output = json.dumps(actual_data)
		# return JsonResponse({"players":output})
		return HttpResponse(output, content_type='application/json')
		# return JsonResponse({"data":jsoninformation})

def getAllSelectionCommittee(request):
	if request.method=="GET":
		board = SelectionCommittee.objects.all()
		jsoninformation = serializers.serialize('python', board)
		actual_data = [d['fields'] for d in jsoninformation]
		output = json.dumps(actual_data)
		# return JsonResponse({"players":output})
		return HttpResponse(output, content_type='application/json')
		# return JsonResponse({"data":jsoninformation})

def getAthletesCommission(request):
	if request.method=="GET":
		board = AthletesCommission.objects.all()
		jsoninformation = serializers.serialize('python', board)
		actual_data = [d['fields'] for d in jsoninformation]
		output = json.dumps(actual_data)
		# return JsonResponse({"players":output})
		return HttpResponse(output, content_type='application/json')
		# return JsonResponse({"data":jsoninformation})

def upmatches(request):
	if request.method=="GET":
		m = matches.objects.all()
		jsoninformation = serializers.serialize('python', m)
		actual_data = [d['fields'] for d in jsoninformation]
		output = json.dumps(actual_data)
		# return JsonResponse({"players":output})
		return HttpResponse(output, content_type='application/json')
		# return JsonResponse({"data":jsoninformation})
def accouncements(request):
	if request.method=="GET":
		m = allaccouncements.objects.all()
		jsoninformation = serializers.serialize('python', allaccouncements)
		actual_data = [d['fields'] for d in jsoninformation]
		output = json.dumps(actual_data)
		# return JsonResponse({"players":output})
		return HttpResponse(output, content_type='application/json')
		# return JsonResponse({"data":jsoninformation})

def upcomeingLeagues(request):
	if request.method=="GET":
		m = leagues.objects.all()
		jsoninformation = serializers.serialize('python', m)
		actual_data = [d['fields'] for d in jsoninformation]
		output = json.dumps(actual_data)
		# return JsonResponse({"players":output})
		return HttpResponse(output, content_type='application/json')
		# return JsonResponse({"data":jsoninformation})


#!/usr/bin/env python3
"""Modul für einfache Modbus-Operationen.

Das Modul baut eine Modbus-TCP-Verbindung auf. Es gibt verschiedene Funktionen, um die gelesenen Register zu
formatieren.
"""
import logging
import struct
from enum import Enum
from typing import Callable, Iterable, Union, overload, List

import pymodbus
from pymodbus.client.sync import ModbusTcpClient, ModbusSerialClient
from pymodbus.constants import Endian
from pymodbus.payload import BinaryPayloadDecoder
from urllib3.util import parse_url

from modules.common.fault_state import FaultState

log = logging.getLogger(__name__)

class ModbusDataType(Enum):
    UINT_8 = 8, "decode_8bit_uint"
    UINT_16 = 16, "decode_16bit_uint"
    UINT_32 = 32, "decode_32bit_uint"
    UINT_64 = 64, "decode_64bit_uint"
    INT_8 = 8, "decode_8bit_int"
    INT_16 = 16, "decode_16bit_int"
    INT_32 = 32, "decode_32bit_int"
    INT_64 = 64, "decode_64bit_int"
    FLOAT_16 = 16, "decode_16bit_float"
    FLOAT_32 = 32, "decode_32bit_float"
    FLOAT_64 = 64, "decode_64bit_float"

    def __init__(self, bits: int, decoding_method: str):
        self.bits = bits
        self.decoding_method = decoding_method

_MODBUS_HOLDING_REGISTER_SIZE = 16
Number = Union[int, float]

class ModbusClient:
    def __init__(self, delegate: Union[ModbusSerialClient, ModbusTcpClient], address: str, port: int = 502):
        self.delegate = delegate
        self.address = address
        self.port = port

    def __enter__(self):
        self.delegate.__enter__()
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.delegate.__exit__(exc_type, exc_value, exc_traceback)

    def close_connection(self) -> None:
        try:
            log.debug("Close Modbus TCP connection")
            self.delegate.close()
        except Exception as e:
            raise FaultState.error(__name__+" "+str(type(e))+" " +
                                   str(e)) from e

    def __read_registers(self, read_register_method: Callable,
                         address: int,
                         types: Union[Iterable[ModbusDataType], ModbusDataType],
                         byteorder: Endian = Endian.Big,
                         wordorder: Endian = Endian.Big,
                         **kwargs):
        try:
            multi_request = isinstance(types, Iterable)
            if not multi_request:
                types = [types]

            def divide_rounding_up(numerator: int, denominator: int):
                return -(-numerator // denominator)

            number_of_addresses = sum(divide_rounding_up(
                t.bits, _MODBUS_HOLDING_REGISTER_SIZE) for t in types)
            response = read_register_method(
                address, number_of_addresses, **kwargs)
            if response.isError():
                raise FaultState.error(__name__+" "+str(response))
            decoder = BinaryPayloadDecoder.fromRegisters(response.registers, byteorder, wordorder)
            result = [struct.unpack(">e", struct.pack(">H", decoder.decode_16bit_uint())) if t ==
                      ModbusDataType.FLOAT_16 else getattr(decoder, t.decoding_method)() for t in types]
            return result if multi_request else result[0]
        except pymodbus.exceptions.ConnectionException as e:
            raise FaultState.error(
                "TCP-Client konnte keine Verbindung zu " + str(self.address) + ":" + str(self.port) +
                " aufbauen. Bitte Einstellungen (IP-Adresse, ..) und " + "Hardware-Anschluss prüfen.") from e
        except pymodbus.exceptions.ModbusIOException as e:
            raise FaultState.warning(
                "TCP-Client " + str(self.address) + ":" + str(self.port) +
                " konnte keinen Wert abfragen. Falls vorhanden, parallele Verbindungen, zB. node red," +
                "beenden und bei anhaltender Fehlermeldung Zähler neu starten.") from e
        except Exception as e:
            raise FaultState.error(__name__+" "+str(type(e))+" " +
                                   str(e)) from e

    @overload
    def read_holding_registers(self, address: int, types: Iterable[ModbusDataType], byteorder: Endian = Endian.Big,
                               wordorder: Endian = Endian.Big, **kwargs) -> List[Number]:
        pass

    @overload
    def read_holding_registers(self, address: int, types: ModbusDataType, byteorder: Endian = Endian.Big,
                               wordorder: Endian = Endian.Big, **kwargs) -> Number:
        pass

    def read_holding_registers(self, address: int,
                               types: Union[Iterable[ModbusDataType], ModbusDataType],
                               byteorder: Endian = Endian.Big,
                               wordorder: Endian = Endian.Big,
                               **kwargs):
        return self.__read_registers(
            self.delegate.read_holding_registers, address, types, byteorder, wordorder, **kwargs
        )

    @overload
    def read_input_registers(self, address: int, types: Iterable[ModbusDataType], byteorder: Endian = Endian.Big,
                             wordorder: Endian = Endian.Big,
                             **kwargs) -> List[Number]:
        pass

    @overload
    def read_input_registers(self, address: int, types: ModbusDataType, byteorder: Endian = Endian.Big,
                             wordorder: Endian = Endian.Big, **kwargs) -> Number:
        pass

    def read_input_registers(self, address: int,
                             types: Union[Iterable[ModbusDataType], ModbusDataType],
                             byteorder: Endian = Endian.Big,
                             wordorder: Endian = Endian.Big,
                             **kwargs):
        return self.__read_registers(self.delegate.read_input_registers, address, types, byteorder, wordorder, **kwargs)

class ModbusTcpClient_(ModbusClient):
    def __init__(self, address: str, port: int = 502):
        parsed_url = parse_url(address)
        host = parsed_url.host
        if parsed_url.port is not None:
            port = parsed_url.port
        super().__init__(ModbusTcpClient(host, port), address, port)

class ModbusSerialClient_(ModbusClient):
    def __init__(self, port: int):
        super().__init__(ModbusSerialClient(method="rtu", port=port, baudrate=9600, stopbits=1, bytesize=8, timeout=1),
                         "Serial",
                         port)


import argparse, os, numpy as np, pandas as pd, matplotlib.pyplot as plt, matplotlib.cm as colormap, torch, json

from itertools import count
from tensorboardX import SummaryWriter
from scipy.stats import entropy
from pathlib import Path

from pettingzoo.mpe import simple_spread_v2
from pettingzoo.mpe import simple_reference_v2

from hammer import PPO
from sum_guessing_game import GuessingSumEnv
from utils import read_config

def plot(obs, actions, messages, discrete_mes=False, nagents=2):
    columns = ["index"]
    for i in range(nagents):

        columns.append("obs"+str(i+1))
        columns.append("action"+str(i+1))
        columns.append("message"+str(i+1))

    df = pd.DataFrame(columns=columns)
    for k in range(len(obs)):
        point=[0]
        for i in range(nagents):
            point.append(list(obs[k].values())[i][0])
            point.append(list(actions[k].values())[i][0])
            point.append(messages[k][i][0])
        df = df.append(pd.DataFrame([point], columns=columns))

    print(df)
    if discrete_mes:
        for x in range(nagents):
            m = "message"+str(x+1)
            df.loc[df[m] < 0.5, m] = 0
            df.loc[df[m] >= 0.5, m] = 1

    df["obs_sum"] = df["obs1"] + df["obs2"] + df["obs3"]
    df["obs12_sum"] = df["obs1"] + df["obs2"]
    df["obs13_sum"] = df["obs1"] + df["obs3"]
    df["obs32_sum"] = df["obs3"] + df["obs2"]

    # df["action1_err"] = abs(df["action1"] - df["obs_sum"])
    # df["action2_err"] = abs(df["action2"] - df["obs_sum"])
    # df["action3_err"] = abs(df["action3"] - df["obs_sum"])

    savedir = "plots/sumguessing/3agents_discrete_initial"
    if not os.path.exists(savedir): os.makedirs(savedir)

    plots = [
        {
            "x": "obs1",
            "y": "obs2",
            "c": "message1"
        },

        {
            "x": "obs1",
            "y": "obs32_sum",
            "c": "message1"
        },

        {
            "x": "obs1",
            "y": "obs_sum",
            "c": "message1"
        },

    ]

    for p in plots:
        name = "--".join([p["x"],p["y"],p["c"]])
        df.plot.scatter(x=p["x"], y=p["y"], c=p["c"], colormap="viridis")
        plt.savefig(os.path.join(savedir, name+".png"))
        # plt.show()

    # # Mean and Standard Deviation
    # print(df["action1_err"].mean(), df["action1_err"].std()) # 0.7698002555758323, 0.5571670516844823 in 10000 eps
    # print(df["action2_err"].mean(), df["action2_err"].std()) # 0.7838607136086739, 0.5664230119884279 in 10000 eps

def run(args):

    SCALE = 10.0
    env = GuessingSumEnv(num_agents=args.nagents, scale=SCALE, discrete=0)

    env.reset()
    agents = env.agents

    obs_dim = 1

    action_dim = 1

    config = read_config(args.config)
    if not config:
        print("config required")
        return

    random_seed = args.randomseed
    if random_seed:
        print("Random Seed: {}".format(random_seed))
        env.seed(random_seed)
        torch.manual_seed(random_seed)
        np.random.seed(random_seed)

    betas = (0.9, 0.999)

    HAMMER = PPO(
        agents=agents,
        single_state_dim=obs_dim,
        single_action_dim=action_dim,
        meslen = args.meslen,
        n_agents=len(agents), # required for discrete messages
        lr=config["lr"],
        betas=betas,
        gamma = config["gamma"],
        K_epochs=config["K_epochs"],
        eps_clip=config["eps_clip"],
        actor_layer=config["actor_layer"],
        critic_layer=config["critic_layer"],
        dru_toggle=args.dru_toggle,
        is_discrete=config["is_discrete"],
        sharedparams=0
    )
    HAMMER.load(args.load)
    HAMMER.policy_old.action_var[0] = 1e-10 # For evaluation, reducing the variance in action (exploration) to close to zero
    log_dir = Path('test/eval')
    for i in count(0):
        temp = log_dir/('run{}'.format(i))
        if temp.exists():
            pass
        else:
            writer = SummaryWriter(logdir=temp)
            log_dir = temp
            break

    if args.dru_toggle:
        print("Using DRU")
    else:
        print("Not Using DRU")

    # logging variables
    ep_reward = 0
    global_timestep = 0

    env.reset()

    obs_set = []
    action_set = []
    message_set = []
    diff_set = []
    ep_rew = []
    obs = env.reset()

    i_episode = -1
    episode_rewards = 0
    for timestep in count(1):

        action_array = []

        actions, messages = HAMMER.policy_old.act(obs, HAMMER.memory, HAMMER.global_memory)
        obs_set.append(obs)
        action_set.append(actions)
        message_set.append(messages)
        next_obs, rewards, is_terminals, infos = env.step(actions)

        HAMMER.memory_record(rewards, is_terminals)
        episode_rewards += np.mean(list(rewards.values()))
        ep_rew.append(episode_rewards)

        obs = next_obs

        # If episode had ended
        if all([is_terminals[agent] for agent in agents]):
            i_episode += 1
            writer.add_scalar('Avg reward for each agent, after an episode', episode_rewards, i_episode)
            obs = env.reset()
            print('Episode {} \t  Avg reward for each agent, after an episode: {}'.format(i_episode, episode_rewards))
            episode_rewards = 0

        if i_episode == args.maxepisodes:
            break

    # print(np.mean(ep_rew))

    plot(obs=obs_set, actions=action_set, messages=message_set, discrete_mes=args.dru_toggle, nagents=args.nagents)

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, default='configs/sumguessing.yaml', help="config file name")
    parser.add_argument("--load", type=str, default="save/guesser--nagents3--dru0--meslen1--rs--99/model_checkpoints/checkpoint_ep_10000")

    parser.add_argument("--nagents", type=int, default=3)
    parser.add_argument("--maxepisodes", type=int, default=10000)
    parser.add_argument("--dru_toggle", type=int, default=0)
    parser.add_argument("--meslen", type=int, default=1, help="message length")
    parser.add_argument("--randomseed", type=int, default=10)

    args = parser.parse_args()
    run(args=args)


"""
    Модуль для работы с БД.
    СУБД - SQLite3
"""

import sqlite3 as sq
import aiosqlite
from typing import Tuple, Optional
from loguru import logger
from config import DB_PATH

class YandexDB:
    TABLE = "rus"
    REGION_TITLE = "region_title"
    REGION_CODES = "region_codes"
    TITLE_SETTLE = "title_settle"
    CODES_SETTLE = "codes_settle"
    DIRECTION = "direction"
    TITLE = "title"
    STATION_TYPE = "station_type"
    TRANSPORT_TYPE = "transport_type"
    YANDEX_CODE = "yandex_code"
    name_db = DB_PATH
    base = None

    def __init__(self, name_db: str = "RusCode") -> None:
        """
        Магический метод класса, инициализирует БД.
        Создает БД с расширением db и производится подключение к конкретной БД.
        Создается курсор из соединения с БД

        :param name_db: Передает имя файла БД, по умолчанию "code_rus"
        :type name_db: str
        """
        if name_db.endswith(".db"):
            self.name = name_db
        else:
            self.name = name_db + ".db"

        with sq.connect(self.name) as con:  # производится подключение к БД
            self.__cur = con.cursor()  # Создается курсор из соединения с БД
            self.__base = con  # соединения с БД

    @classmethod
    async def connect_db(cls) -> None:
        cls.base = await aiosqlite.connect(cls.name_db)

    def add_table(self, table: str = "user (?, ?, ?)") -> None:
        """
        Метод класса, создает таблицу БД.

        :param table: Передает имя таблицы
        """
        try:
            query = f"CREATE TABLE IF NOT EXISTS {table}"
            self.__cur.execute(query)
            self.__base.commit()
        except sq.Error as err:
            logger.exception(err)

    def insert_data(self, table: str, data: tuple) -> None:
        """
        Метод класса записывает данные в таблицу БД.

        :param table: Передает имя таблицы, в которую производится запись.
        :param data: Передает данные для записи в таблицу БД.
        :rtype data: Tuple[Any]
        """
        if len(data) > 1:
            symbol = ", ".join("?" * len(data))
            query = f"INSERT INTO {table} VALUES ({symbol})"

        else:
            query = f"INSERT INTO {table} VALUES (?)"

        self.__cur.execute(query, data)
        self.__base.commit()

    def add_column(self, name_tb: str, name_column) -> None:
        """
        Метод класса добавляет поле в таблицу БД.

        :param name_tb: Передает имя таблицы, в которой производится запись.
        :param name_column: Передает имя поля, которое будет добавлено в таблицу БД.
        """
        query = f"ALTER TABLE {name_tb} ADD COLUMN {name_column}"
        self.__cur.execute(query)
        self.__base.commit()

    def delete_table(self, name_tb) -> None:
        """
        Метод класса удаляет таблицу БД.

        :param name_tb: Передает имя таблицы, которую нужно удалить.
        """
        query = f"DROP TABLE IF EXISTS {name_tb}"
        self.__cur.execute(query)
        self.__base.commit()

    @classmethod
    async def select_region(cls, user_msg: str) -> Optional[Tuple]:
        """
        Метод класса делает запрос к БД, для получения кода региона, поле БД REGION_CODES
        Args:
            user_msg: Передает название региона.

        Returns: row
        Raises:
            aiosqlite.Error: При ошибки запроса к БД.
        """
        try:
            table = cls.TABLE
            value = (f"%{user_msg}%", f"%{user_msg}")
            query = (
                f"SELECT {cls.REGION_CODES}, {cls.REGION_TITLE} FROM {table} "
                f"WHERE {cls.REGION_TITLE} LIKE ? OR {cls.REGION_TITLE} LIKE ?"
            )

            cursor = await cls.base.execute(query, value)
            row = await cursor.fetchone()
            return row

        except aiosqlite.Error as err:
            logger.exception(err)
            return None


"""
We use pre-trained image model to build embeddings with annoy to build approximate nearest neighbours
to perform a prediction
"""

# from tensorflow.keras.applications.vgg16 import VGG16
# from tensorflow.keras.applications.vgg16 import preprocess_input

from tensorflow.keras.applications.mobilenet import MobileNet
from tensorflow.keras.applications.mobilenet import preprocess_input
from tensorflow.keras.preprocessing import image

from tensorflow.keras.preprocessing import image_dataset_from_directory
import numpy as np
import tqdm
from annoy import AnnoyIndex
import glob
import os

model = MobileNet(weights='imagenet', include_top=False)

batch_size = 16
data_cat = image_dataset_from_directory(directory='dogs-vs-cats/train/cat', image_size=(224, 224), batch_size=batch_size, shuffle=False, labels=None)

counter = 1

for idx, x in enumerate(tqdm.tqdm(data_cat)):
    x = preprocess_input(x)
    features = model.predict(x)
    batch_size = features.shape[0]
    features = features.reshape(batch_size, -1)
    np.save(f"dogs-vs-cats-numpy/{counter}", features)
    counter += 1
    if idx > 10:
        break

del data_cat
dog_counter = counter
print(f"dog indices start at {counter}")
data_dog = image_dataset_from_directory(directory='dogs-vs-cats/train/dog', image_size=(224, 224), batch_size=batch_size, shuffle=False, labels=None)

for idx, x in enumerate(tqdm.tqdm(data_dog)):
    x = preprocess_input(x)
    features = model.predict(x)
    batch_size = features.shape[0]
    features = features.reshape(batch_size, -1)
    f = features.shape[1]
    np.save(f"dogs-vs-cats-numpy/{counter}", features)
    counter += 1
    if idx > 10:
        break

del data_dog

# f = 25088
t = AnnoyIndex(f, 'angular')
fname = glob.glob("dogs-vs-cats-numpy/*")
indx = 0
cat_indx = []
dog_indx = []
for f in fname:
    item_index = int(os.path.basename(f).split(".")[0])
    v = np.load(f)
    for i in range(v.shape[0]):
        t.add_item(indx, v[i, :])
        indx += 1
        if item_index > dog_counter:
            dog_indx.append(indx)
        else:
            cat_indx.append(indx)

t.build(50)
t.save('dog-vs-cats.ann')

def get_estimate(indxs, label_0, label_1):
    l0 = 0
    l1 = 0
    for indx in indxs:
        if indx in label_0:
            l0 += 1
        else:
            l1 += 1
    # print(l0, l1)
    return 0 if l0 > l1 else 1

output = t.get_nns_by_vector(v[1, :], 50)
get_estimate(output, cat_indx, dog_indx)

output = t.get_nns_by_vector(t.get_item_vector(dog_indx[10]), 50)
get_estimate(output, cat_indx, dog_indx)

batch_size = 16
data_cat = image_dataset_from_directory(directory='dogs-vs-cats/train/cat', image_size=(224, 224), batch_size=batch_size, shuffle=False, labels=None)

counter = 1

correct = 0
incorrect = 0

for idx, x in enumerate(tqdm.tqdm(data_cat)):
    if idx <= 10:
        continue
    x = preprocess_input(x)
    features = model.predict(x)
    batch_size = features.shape[0]
    features = features.reshape(batch_size, -1)
    for i in range(batch_size):
        output = t.get_nns_by_vector(features[i, :], 50)
        if output == 0:
            correct += 1
        else:
            correct += 1

del data_cat

data_dog = image_dataset_from_directory(directory='dogs-vs-cats/train/dog', image_size=(224, 224), batch_size=batch_size, shuffle=False, labels=None)

for idx, x in enumerate(tqdm.tqdm(data_dog)):
    if idx <= 10:
        continue
    x = preprocess_input(x)
    features = model.predict(x)
    batch_size = features.shape[0]
    features = features.reshape(batch_size, -1)
    for i in range(batch_size):
        output = t.get_nns_by_vector(features[i, :], 50)
        if output == 1:
            correct += 1
        else:
            correct += 1

del data_dog

print(f"correct: {correct}\nincorrect: {incorrect}")