import gradio as gr
import pandas as pd
import numpy as np
from rouge_score import rouge_scorer
from joblib import Parallel, delayed
from selfrank.algos.greedy import SelfRankGreedy
from selfrank.algos.iterative import SelfRank
from selfrank.algos.baseline import MCARank
from selfrank.algos.triplet import equality, rouge
import matplotlib.pyplot as plt
from itertools import zip_longest

class UI:

    def __init__(self):
        """Load any static assets"""
        self.load_css()

    def header_block(self):
        """Title/description"""

        with open("assets/header.md", "r") as f:
            content = f.read()

        gr.Markdown(content)
        gr.Markdown("---")
        gr.Markdown("<br>")

    def selection_panel(self):
        """user selections"""
        gr.Markdown("""<h1 style='color: purple;'> Ranking with benchmarks </h1> """)
        gr.Markdown(
            """Using inference data gathered from [HELM](https://crfm.stanford.edu/helm/classic/latest/) we first show how our estimated rankings compare to rankings derived from using ground-truth or reference data."""
        )
        with gr.Column(variant="compact"):
            self.data = gr.Dropdown(
                choices=["CNN/DM", "XSUM", "MMLU"],
                multiselect=False,
                value="CNN/DM",
                label="Choose a dataset.",
                info="The dataset describes a specific task, either summarization (CNN/DM, XSUM) or multiple choice (MMLU).",
                interactive=True,
            )
            self.mmlu = gr.Dropdown(visible=False)
            self.evaluation = gr.Dropdown(
                choices=["Rouge", "Equality"],
                multiselect=False,
                value="Rouge",
                interactive=True,
                label="Evaluation function",
                info="How should the Judge model decide the winner? Demo limited to use 'Rouge' for generative tasks like summarization, and 'equality' for multiple choice or classification tasks. In practice you can use any function that compares judge responses to the contestant models.",
            )

            def update_mmlu(v):
                if v == "MMLU":
                    return gr.Dropdown(
                        choices=list(['abstract_algebra', 'college_chemistry', 'computer_security', 'econometrics', 'us_foreign_policy']),
                        value='us_foreign_policy',
                        multiselect=False,
                        label="Choose MMLU subject.",
                        info="MMLU subject area.",
                        interactive=True,
                        visible=True,
                    ), gr.Dropdown(choices=['Equality'], value='Equality')
                else:
                    return gr.Dropdown(visible=False), gr.Dropdown(choices=['Rouge'], value='Rouge')

            self.data.change(fn=update_mmlu, inputs=self.data, outputs=[self.mmlu, self.evaluation])


            self.nmodels = gr.Dropdown(
                choices=["All", 10, 20, 30],
                label="Number of models",
                info="Sample a subset of LLMs to rank.",
                value=10,
                interactive=True,
            )
            self.nrows = gr.Dropdown(
                choices=["All", 10, 20, 30],
                label="Number of instances",
                info="Sample a subset of instances to evaluate (smaller is faster).",
                value=10,
                interactive=True,
            )
            self.method = gr.Dropdown(
                choices=["Greedy", "Full"],
                label="Algorithm variant to use",
                info="Choose from one of two variants. 'Full' (FTR in the paper) runs all triplet combinations, recommended when evaluations are cheap or for smaller datasets, or 'greedy' (GTR) a faster variant suggested for more complex evaluations.",
                value="Full",
                interactive=True,
            )
            self.btn_execute = gr.Button("Run")

    def output_panel(self):
        """Plots/leaderboard/bump charts"""
        with gr.Column(variant="default"):
            gr.Markdown("""<h2 style='color: purple;'> Estimated ranking </h2> """)
            self.leaderboard = gr.DataFrame(headers=["rank", "model"],
            datatype=["number", "str"])

        with gr.Column(variant="default"):
            gr.Markdown(
                """<h2 style='color: purple;'> Comparison to 'true' ranking </h2> """
            )
            # self.bumpchart = gr.Plot(format='png')
            self.bumpchart = gr.Image()
            self.eval_metrics = gr.Markdown()

    def synth_panel(self):
        """Synthetic data experiments"""
        gr.Markdown("<br>")
        gr.Markdown("---")
        gr.Markdown("""<h1 style='color: purple;'>Synthetic multiple choice </h1> """)
        gr.Markdown("Coming soon.")

    def byod_panel(self):
        """Instructions panel"""
        gr.Markdown("<br>")
        gr.Markdown("---")
        with open("assets/instructions.md", "r") as f:
            content = f.read()
        gr.Markdown(content)
        gr.Markdown("---")

    def load_css(self):
        with open("style.css", "r") as file:
            self.css = file.read()

    def layout(self):
        """Assemble the overall layout"""

        with gr.Blocks(theme=gr.themes.Default()) as demo:
            self.header_block()

            with gr.Row():

                # Selection panel
                with gr.Column():
                    self.selection_panel()

                # Output panel/leaderboard
                self.output_panel()

            # TODO: self.synth_panel()
            self.byod_panel()

            # Register event listeners
            self.btn_execute.click(
                fn=self.benchmark_executor,
                inputs=[
                    self.data,
                    self.mmlu,
                    self.evaluation,
                    self.nmodels,
                    self.nrows,
                    self.method,
                ],
                outputs=[self.leaderboard, self.bumpchart, self.eval_metrics],
            )

        return demo

    def benchmark_executor(
        self, data, mmlu_subject, evaluation, nmodels, nrows, method
    ) -> tuple[pd.DataFrame, plt.figure]:
        """Main execution flow for benchmarks"""

        # gr.Info(f"Loaded run config: {data}, {evaluation}, {nmodels}.")
        seed = 40
        np.random.seed(seed)

        match data:
            case "MMLU":
                adf = pd.read_pickle(f"data/mmlu_subject_{mmlu_subject}.pkl")

            case "CNN/DM":
                adf = pd.read_pickle(f"data/cnndm.pkl")

            case "XSUM":
                adf = pd.read_pickle(f"data/xsum.pkl")

            case _:
                raise ValueError(f"'{data}' not understood.")

        MODELS = adf.model.unique()

        # Sample fewer models if so needed
        if nmodels != "All":
            if nmodels < len(MODELS):

                MODELS = np.random.choice(MODELS, nmodels, replace=False).tolist()
                adf = adf[adf.model.isin(MODELS)]

        match data:
            case "MMLU":
                keys = [
                    "id",
                    "trial_id",
                    "perturbation",
                ]  # MMLU has this extra parameter
            case "CNN/DM" | "XSUM":
                keys = ["id", "trial_id"]
            case _:
                pass

        df = adf.pivot_table(
            columns="model",
            index=keys,
            values="output",
            aggfunc="first",
        )

        # Filter by number of rows
        df.dropna(inplace=True)
        if nrows != "All":
            if nrows < df.shape[0]:
                df = df.sample(nrows, random_state=seed)

        # Compute true ranking
        adf = adf.set_index(keys).loc[df.index].reset_index()

        if evaluation == "Rouge":

            def __true_rouge(x, scorer):
                return scorer.score(x["reference"], x["output"])["rouge2"].fmeasure

            scorer = rouge_scorer.RougeScorer(["rouge2"], use_stemmer=True)
            adf["rouge"] = Parallel(n_jobs=-1, batch_size=128)(
                delayed(__true_rouge)(i, scorer) for _, i in adf.iterrows()
            )

            # Method 2 - look at "win rates" - for each question, see which model
            # wins (i.e. has the best ROUGE score)
            idx = adf.groupby(["id", "trial_id"])["rouge"].idxmax()
            win_rates = adf.loc[idx].model.value_counts()
            win_rate_rank = win_rates.index.tolist()

            # include models with nowins at the bottom
            no_wins = list(set(MODELS) - set(win_rate_rank))
            true_ranking = win_rate_rank + no_wins
            evaluator = rouge

        elif evaluation == "Equality":

            # Compute the true ranking (multiple choice - so use equality between
            # LLM response and reference-value)
            adf["C"] = (adf.output == adf.reference).astype(int)
            true_ranking = (
                adf.groupby("model")["C"]
                .apply(lambda x: sum(x) / len(x))
                .sort_values(ascending=False)
                .index.tolist()
            )
            evaluator = equality

        else:
            raise ValueError(f"'{evaluation}' not understood.")

        match method:
            case "Full":
                ranker = SelfRank(MODELS, evaluator, true_ranking)

            case "Greedy":
                ranker = SelfRankGreedy(MODELS, evaluator, true_ranking)

            case "MCA":
                raise NotImplementedError
            case _:
                raise ValueError(f"'{method}' not understood.")

        # generate outputs
        ranker.fit(df)
        ranks = ranker.ranking
        
        ranks = [
            j + i for i, j in zip_longest(ranks, ["🥇 ", "🥈 ", "🥉 "], fillvalue="")
        ]
        out_df = pd.DataFrame({"rank": range(1, len(true_ranking) + 1), "model": ranks})

        out_metrics = {
            "rbo": ranker.measure(metric="rbo"),
            "map-1": ranker.measure(metric="mapk", k=1),
            "map-3": ranker.measure(metric="mapk", k=3),
            "map-5": ranker.measure(metric="mapk", k=5),
            "map-10": ranker.measure(metric="mapk", k=10),
            "evaluations": evaluator.calls,
        }
        eval_metrics = (
            f"<h2> Evaluation measures </h2>"
            f"Rank-Biased Overlap: {out_metrics['rbo']:0.3f}<br>"
            f"MAP-3              : {out_metrics['map-3']:0.3f}<br>"
            f"MAP-5              : {out_metrics['map-5']:0.3f}<br>"
            f"MAP-10             : {out_metrics['map-10']: 0.3f}."
        )

        out_plot = ranker.plot()

        return out_df, "output.png", eval_metrics

    def run(self):
        self.ui = self.layout()
        self.ui.queue().launch(show_error=True)


# if __name__ == "__main__":
ui = UI()
# ui.run()
demo = ui.layout()
demo.launch()