# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/4.1_benchmark.traceability.ipynb (unless otherwise specified).

__all__ = ['logger', 'flatten', 'inferDoc2Vec', 'TraceabilityTestbed', 'TraceabilityProcedure', 'TraceabilityBenchmark']

# Cell
# Imports
import logging
import gensim

import matplotlib.pyplot as plt
import pandas as pd
import sentencepiece as sp
import numpy as np

from collections import Counter

# ds4se
from ..mgmnt.prep.traceability import *
from ..mgmnt.prep.nltk import *
from ..mgmnt.prep.bpe import *
from .i import *
from ..exp.info import *
from ..desc.stats import *
from ..vis import *

from pathlib import Path

# Cell
logger = logging.getLogger(__name__)
logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )

# Cell
flatten = lambda l: [item for sublist in l for item in sublist]

# Cell
def inferDoc2Vec(df, model, col_order):
    df_inferred = df.copy()
    df_inferred['src_vec'] =  [
        model.infer_vector(artifact.split(), steps=200) for artifact in df_inferred['src_doc'].values
    ]
    df_inferred['trgt_vec'] =  [
        model.infer_vector(artifact.split(), steps=200) for artifact in df_inferred['trgt_doc'].values
    ]

    df_inferred = df_inferred[col_order]
    return df_inferred

# Cell
class TraceabilityTestbed(Testbed):
    def __init__(self, traces, col_metric):
        Testbed.__init__(self, traces.copy())
        self.data = self.data.rename(columns={col_metric: 'link_metric'})

    def get_ground_truth(self):
        return self.data.loc[self.data.ground_truth == 1]

    def get_non_ground_truth(self):
        return self.data.loc[self.data.ground_truth == 0]

    def __getitem__(self, idx):
        return self.data.iloc[idx]

    def evaluate(self, inpt):
        pass

# Cell
class TraceabilityProcedure(Procedure):
    def __init__(self, metrics, vocab):
        self.vocab = vocab
        Procedure.__init__(self, metrics)

    def evaluate(self, inpt):
        self.metrics['src_entropies'] = get_entropies_from_docs(inpt.src_doc.values, self.vocab)
        self.metrics['trgt_entropies'] = get_entropies_from_docs(inpt.trgt_doc.values, self.vocab)
        logging.info("Finished calculating document level entropies")

        all_docs = inpt.src_doc.values + inpt.trgt_doc.values
        self.metrics['tot_entropy'] = get_entropy_from_docs(all_docs, self.vocab)
        logging.info("Finished calculating corpus level entropy")

        self.metrics['shared_entropy'] = shared_entropy_from_docs(
            (inpt.src_doc.values, inpt.trgt_doc.values), self.vocab
        )
        logging.info("Finished calculating corpus level shared entropy")
        self.metrics['shared_probs'] = [get_shared_probs_from_docs([[docs[0]], [docs[1]]], self.vocab)[0] for docs in zip(inpt.src_doc.values, inpt.trgt_doc.values)]

        get_shared_probs_from_docs(
            (inpt.src_doc.values, inpt.trgt_doc.values), self.vocab
        )[0]
        logging.info("Finished calculating corpus level shared relative probabilities")

        return self.metrics

# Cell
class TraceabilityBenchmark(Benchmark):
    def __init__(self, testbed, procedure):
        self.testbed = testbed
        self.ground_truth = testbed.get_ground_truth()
        self.non_ground_truth = testbed.get_non_ground_truth()
        self.procedure = procedure

        self.src_dtype = self.ground_truth.iloc[0].src_dtype
        self.trgt_dtype = self.ground_truth.iloc[0].trgt_dtype
        Benchmark.__init__(self, testbed, procedure)

    def visualize(self):
        min_len = min(len(self.ground_truth), len(self.non_ground_truth))
        src_dtype = self.ground_truth.iloc[0].src_dtype
        trgt_dtype = self.ground_truth.iloc[0].trgt_dtype

        gt_src_vecs = np.stack(self.ground_truth.src_vec.to_list())[:min_len]
        gt_trgt_vecs = np.stack(self.ground_truth.trgt_vec.to_list())[:min_len]
        gt_src_entropies = self.ground_truth_metrics['src_entropies'][:min_len]
        gt_trgt_entropies = self.ground_truth_metrics['trgt_entropies'][:min_len]
        gt_shared_probs = self.ground_truth_metrics['shared_probs'][:min_len]
        gt_link_metric = self.ground_truth.link_metric.values[:min_len]

        ngt_src_vecs = np.stack(self.non_ground_truth.src_vec.to_list())[:min_len]
        ngt_trgt_vecs = np.stack(self.non_ground_truth.trgt_vec.to_list())[:min_len]
        ngt_src_entropies = self.non_ground_truth_metrics['src_entropies'][:min_len]
        ngt_trgt_entropies = self.non_ground_truth_metrics['trgt_entropies'][:min_len]
        ngt_shared_probs = self.non_ground_truth_metrics['shared_probs'][:min_len]
        ngt_link_metric = self.non_ground_truth.link_metric.values[:min_len]
        logging.info(f'Clusterization for {self.src_dtype}')
        clusterize_w_entropy(gt_src_vecs, ngt_src_vecs, gt_src_entropies, ngt_src_entropies)
        logging.info(f'Clusterization for {self.trgt_dtype}')
        clusterize_w_entropy(gt_trgt_vecs, ngt_trgt_vecs, gt_trgt_entropies, ngt_trgt_entropies)
        visualize_events(gt_shared_probs, 'b', f'libest-{self.trgt_dtype}', 'groundtruth')
        visualize_events(ngt_shared_probs, 'r', f'libest-{self.trgt_dtype}', 'non-groundtruth')
        vis_3d(
            [gt_src_entropies, gt_trgt_entropies, gt_link_metric],
            [ngt_src_entropies, ngt_trgt_entropies, ngt_link_metric],
            src_dtype, trgt_dtype
        )

    def gen_report(self):
        print(f'Descriptive Statistics for Groundtruth {self.src_dtype}')
        report_stats(self.ground_truth_metrics['src_entropies'])
        print()
        print(f'Descriptive Statistics for Groundtruth {self.trgt_dtype}')
        report_stats(self.ground_truth_metrics['trgt_entropies'])
        print()

        print(f'Descriptive Statistics for Groundtruth {self.src_dtype}')
        report_stats(self.non_ground_truth_metrics['src_entropies'])
        print()
        print(f'Descriptive Statistics for Groundtruth {self.trgt_dtype}')
        report_stats(self.non_ground_truth_metrics['trgt_entropies'])
        print()
#         pass

    def evaluate(self):
        self.ground_truth_metrics = self.procedure.evaluate(self.ground_truth)
        self.non_ground_truth_metrics = self.procedure.evaluate(self.non_ground_truth)