import os
import argparse
import sys
import cltl.dialogue_evaluation.utils.scenario_check as check
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from emissor.persistence import ScenarioStorage
from emissor.representation.scenario import Modality

import cltl.dialogue_evaluation.utils.text_signal as text_util
from cltl.dialogue_evaluation.api import BasicEvaluator
from cltl.dialogue_evaluation.metrics.utterance_likelihood import MLM


class LikelihoodEvaluator(BasicEvaluator):
    def __init__(self, model, model_name, max_context=300, len_top_tokens=20):
        """Creates an evaluator that will use USR Masked Language Model scoring to approximate the quality of a conversation, across turns.

        We use the Roberta model that was pretrained with the TopicalChat data by the USR team
        as a model for gettting the averaged token likelihood of the target sentence.
        The function *sentence_likelihood* also returns the most likely sentence according to the model
        and the averaged score for the mostly likely tokens.
        We can thus compare the actual response in a turn with the response that would be generated by the LM.

        The LM will return a number of results for the masked token with probability scores.
        We compare the target token with the results to get the score for the target token.
        If the target token is not in the results, we set the probability to "0".

        You can set the number of results returned by the model.
        The more results, the more likely the target gets a score, albeit a very low score.
        params
        model_path_mlm: one of ['google-bert/bert-base-multilingual-cased', 'adamlin/usr-topicalchat-roberta_ft', 'xlm-roberta-base', 'roberta-base']
        returns: None
        """
        super(LikelihoodEvaluator, self).__init__()
        self.model_path_mlm = model
        self.max_context = max_context
        self.len_top_tokens = len_top_tokens
        self.model_name= model_name

        # Create MLM
        self.model_mlm = MLM(path=self.model_path_mlm, top_results=self.len_top_tokens)

        self._log.debug(f"Likelihood Evaluator ready")

    def evaluate_conversation(self, scenario_folder, scenario_id, metrics_to_plot=None):
        # Create the scenario folder, the json files and a scenarioStorage and scenario in memory
        scenario_storage = ScenarioStorage(scenario_folder)
        scenario_ctrl = scenario_storage.load_scenario(scenario_id)
        signals = scenario_ctrl.get_signals(Modality.TEXT)
        ids, turns, speakers = text_util.get_utterances_with_context_from_signals(signals, self.max_context)

        print(f'----------SCENARIO:{scenario_folder}, EVALUATION:likelihood metrics---------')
        print('Nr of turns:', len(turns), 'Speakers:', speakers, 'Max context:', self.max_context)

        # Get likelihood scored
        speaker_mlm_scores = {k: [] for k in speakers}
        speaker_mlm_max_scores = {k: [] for k in speakers}
        speaker_turns = {k: [] for k in speakers}

        df = self._calculate_metrics(self.model_mlm, turns, speaker_mlm_scores, speaker_mlm_max_scores, speaker_turns)
        avg_df = self._average_metrics(speakers, turns, speaker_mlm_scores, speaker_mlm_max_scores)

        # Save
        evaluation_folder = os.path.join(scenario_folder, scenario_id, 'evaluation')
        if not os.path.exists(evaluation_folder):
            os.mkdir(evaluation_folder)
        self._save(df, avg_df, evaluation_folder)
        #
        if metrics_to_plot:
            self.plot_metrics_progression(metrics_to_plot, [df], evaluation_folder)

    @staticmethod
    def _calculate_metrics(model_mlm, turns, speaker_mlm_scores, speaker_mlm_max_scores, speaker_turns):
        # Iterate turns
        rows = []
        for index, turn in enumerate(turns):
            print(f"Processing turn {index}/{len(turns) - 1}")
            context = turn[0]
            target = turn[1]
            cue = turn[2]
            speaker = turn[3]
         #   print(f"\tCalculating likelihood scores")
            llh, best_sentence, max_score = model_mlm.sentence_likelihood(context, target)
            rows.append({"Turn": index, "Speaker": speaker, "Cue": cue, "Response": target, "Context": context,
                         "MLM response": best_sentence, "System llh": llh, "MLM llh": max_score})

            if speaker:
                speaker_turns[speaker].append(index)
                speaker_mlm_scores[speaker].append(llh)
                speaker_mlm_max_scores[speaker].append(max_score)

        return pd.DataFrame(rows)

    @staticmethod
    def _average_metrics(speakers, turns, speaker_mlm_scores, speaker_mlm_max_scores):
        # Iterate turns
        print(f"\n\tCalculating average likelihood scores")
        overall_rows = []
        for speaker in speakers:
            print(f"\t\tProcessing speaker {speaker}/{len(speakers) - 1}")
            mlm_scores = speaker_mlm_scores[speaker]
            mlm_average_score = sum(mlm_scores) / len(mlm_scores)
            mlm_max_scores = speaker_mlm_max_scores[speaker]
            mlm_average_max_score = sum(mlm_max_scores) / len(mlm_max_scores)
            overall_rows.append({'Speaker': speaker, 'Nr. turns': len(turns),
                                 'MLM': mlm_average_score, 'MLM avg': mlm_average_score,
                                 'MLM max': mlm_max_scores, 'MLM avg max': mlm_average_max_score})

        # Save
        return pd.DataFrame(overall_rows)

    def _save(self, df, avg_df, evaluation_folder):
        filename = "likelihood_evaluation_" + self.model_name + "_context" + str(self.max_context) + ".csv"
        path = os.path.join(evaluation_folder, filename)
        df.to_csv(path, index=False)
        print(f"\n\tSaved to file: {path}")

        filename = "likelihood_evaluation_" + self.model_name + "_context" + str(self.max_context) + "_overall.csv"
        path = os.path.join(evaluation_folder, filename)
        avg_df.to_csv(path, index=False)
        print(f"\n\tSaved to file: {path}")

    def plot_metrics_progression(self, metrics, convo_dfs, evaluation_folder):
        # Plot metrics progression per conversation
        for metric in metrics:
            metric_df = pd.DataFrame()

            # Iterate conversations
            for idx, convo_df in enumerate(convo_dfs):
                conversation_id = f'Conversation {idx}'
                convo_df = convo_df.set_index('Turn')

                # Add into a dataframe
                if len(metric_df) == 0:
                    metric_df[conversation_id] = convo_df[metric]
                else:
                    metric_df = pd.concat([metric_df, convo_df[metric]], axis=1)
                    metric_df.rename(columns={metric: conversation_id}, inplace=True)

            # Cutoff and plot
            self.plot_progression(metric_df, metric, evaluation_folder)

    @staticmethod
    def plot_progression(df_to_plot, xlabel, evaluation_folder):
        df_to_plot = df_to_plot.reset_index().melt('Turn', var_name='cols', value_name=xlabel)

        g = sns.relplot(x="Turn", y=xlabel, hue='cols', data=df_to_plot, kind='line')

        ax = plt.gca()
        plt.xlim(0)
        plt.xticks(ax.get_xticks()[::5], rotation=45)

        plot_file = os.path.join(evaluation_folder, f"{xlabel}.png")
        g.figure.savefig(plot_file, dpi=300, transparent=True, bbox_inches='tight')
        plt.close()
        print(f"\tSaved to file: {plot_file}")



    def process_all_scenarios(self, emissor:str, scenarios:[]):
        for scenario in scenarios:
            if not scenario.startswith("."):
                scenario_path = os.path.join(emissor, scenario)
                has_scenario, has_text, has_image, has_rdf = check.check_scenario_data(scenario_path, scenario)
                check_message = "Scenario:" + scenario + "\n"
                check_message += "\tScenario JSON:" + str(has_scenario) + "\n"
                check_message += "\tText JSON:" + str(has_text) + "\n"
                check_message += "\tImage JSON:" + str(has_image) + "\n"
                check_message += "\tRDF :" + str(has_rdf) + "\n"
                print(check_message)
                if not has_scenario:
                    print("No scenario JSON found. Skipping:", scenario_path)
                elif not has_text:
                    print("No text JSON found. Skipping:", scenario_path)
                else:
                    self.evaluate_conversation(emissor, scenario)

def main(emissor_path:str, scenario:str,  model, model_name, max_context=300, len_top_tokens=20):
   # model = "/Users/piek/Desktop/t-MA-Combots-2024/code/ma-communicative-robots/leolani_text_to_ekg/resources/usr-topicalchat-roberta_ft"
   # model_name = "USR"
    model = "xlm-roberta-base"
    model_name = "XLM-RoBERTa"
  #  emissor_path = "/Users/piek/Desktop/t-MA-Combots-2024/assignments/assignment-1/leolani_local/emissor"
  #  emissor_path = "/Users/piek/Desktop/t-MA-Combots-2024/assignments/assignment-1/leolani_text_to_ekg_restrained/emissor"
    emissor_path = "/Users/piek/Desktop/t-MA-Combots-2024/assignments/assignment-1/leolani_text_to_ekg_wild/emissor"
    scenario=""
    max_context=200
    len_top_tokens=10
    evaluator = LikelihoodEvaluator(model=model, model_name=model_name, max_context=max_context,
                                    len_top_tokens=len_top_tokens)

   # emissor_path = "/Users/piek/Desktop/t-MA-Combots-2024/assignments/emissor-test"
   # scenario = ""

    folders = []
    if not scenario:
        folders = os.listdir(emissor_path)
    else:
        folders = [scenario]

    evaluator.process_all_scenarios(emissor_path, folders)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Statistical evaluation emissor scenario')
    parser.add_argument('--emissor-path', type=str, required=False, help="Path to the emissor folder", default='')
    parser.add_argument('--scenario', type=str, required=False, help="Identifier of the scenario", default='')
    parser.add_argument('--model', type=str, required=False, help="Path to the model or huggingface URL", default="google-bert/bert-base-multilingual-cased")
    parser.add_argument('--model_name', type=str, required=False, help="Name of the modelL", default="mBert")
    parser.add_argument('--context', type=int, required=False, help="Maximum character length of the context" , default=300)
    parser.add_argument('--top_results', type=int, required=False, help="Maximum number of MASKED results considered" , default=20)

    args, _ = parser.parse_known_args()
    print('Input arguments', sys.argv)
    main(emissor_path=args.emissor_path,
         scenario=args.scenario,
         model=args.model,
         model_name=args.model_name,
         max_context=args.context,
         len_top_tokens=args.top_results)

