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


class UI:

    def __init__(self):
        """Load any static assets"""
        pass

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

        gr.Markdown(
            """<h1 style='text-align: center; color: black;'>🥇 Ranking LLMs without ground truth </h1>"""
        )
        gr.Markdown(
            "This space demonstrates reference-free ranking of large language models describe in our ACL Findings paper [Ranking Large Language Models without Ground Truth](https://arxiv.org/abs/2402.14860). <br>"
            "Inspired by real life where both an expert and a knowledgeable person can identify a novice the main idea is to consider triplets of models, where each one of them evaluates the other two, correctly identifying the worst model in the triplet with high probability. Iteratively performing such evaluations yields a estimated ranking that doesn't require ground truth/reference data which can be expensive to gather. The methods are a viable low-resource ranking mechanism for practical use.<br>"
            "[Source code](https://huggingface.co/spaces/ibm/llm-rank-themselves/tree/main).<br>"
        )
        gr.Markdown('---') 
        gr.Markdown('<br>')
        

    def selection_panel(self):
        """user selections"""
        gr.Markdown("""<h2 style='color: purple;'> Benchmark experiments </h2> """)
        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 task",
                interactive=True,
            )
            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.",
            )
            self.nmodels = gr.Dropdown(
                choices=[None, 10, 20, 30],
                label="Number of models",
                info="Sample a subset of LLMs to rank.",
                value=10,
                interactive=True,
            )
            self.nrows = gr.Dropdown(
                choices=[None, 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()

        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("""<h2 style='color: purple;'>Synthetic multiple choice </h2> """)

    def byod_panel(self):
        """ Synthetic data experiments """
        gr.Markdown('<br>')
        gr.Markdown('---')
        gr.Markdown("""<h2 style='color: purple;'>BYOD </h2> """)
        
            
    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()

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

        return demo 
    
    def benchmark_executor(self, data, evaluation, nmodels, nrows, method) -> tuple[pd.DataFrame, plt.figure]:
        """ Main execution flow for benchmarks """
        
        #gr.Info(f"Loaded run config: {data}, {evaluation}, {nmodels}.")
        
        match data:
            case 'MMLU':
                adf = pd.read_pickle(f"data/mmlu_subject_abstract_algebra.pkl")
                MODELS = adf.model.unique()
            
            case 'CNN/DM':
                adf = pd.read_pickle(f"data/cnndm.pkl")
                MODELS = adf.model.unique()

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

        # Sample fewer models if so needed
        if nmodels is not None:
            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':
                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 is not None:
            if nrows < df.shape[0]:
                df = df.sample(nrows)
        
        # 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)
        out_df = pd.DataFrame({'rank': range(1, len(true_ranking)+1), 'model': ranker.ranking})

        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"Evaluation measures: <br>"
                        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()