# Copyright (c) Hegel AI, Inc.
# All rights reserved.
#
# This source code's license can be found in the
# LICENSE file in the root directory of this source tree.

from typing import Callable, Dict, List, Optional, Tuple, Union
from operator import itemgetter
from collections import defaultdict
import itertools
import logging
from IPython import display
from tabulate import tabulate
import pandas as pd

try:
    import pymongo
except ImportError:
    pymongo = None

from prompttools.requests.request_queue import RequestQueue
from ..widgets.feedback import FeedbackWidgetProvider
from ..widgets.comparison import ComparisonWidgetProvider
from ..widgets.utility import is_interactive
from .error import PromptExperimentException

pd.set_option("display.max_colwidth", 0)


class Experiment:
    r"""
    Base class for experiment. This should not be used directly, please use the subclasses instead.
    """

    completion_fn: Callable
    all_args: Dict

    def __init__(self):
        self.queue = RequestQueue()
        self.argument_combos: list[dict] = []
        self.results = []
        self.scores = defaultdict(list)
        self.feedback_widget_provider = FeedbackWidgetProvider(
            self.completion_fn, self._aggregate_metric, self._get_human_eval_listener
        )
        self.comparison_widget_provider = ComparisonWidgetProvider(
            self.completion_fn,
            self._aggregate_comparison,
            self._get_comparison_listener,
        )

    @classmethod
    def initialize(cls, test_parameters: dict[str, list], frozen_parameters: dict):
        r"""
        An alternate way to initialize an experiment by specifying which parameters should be tested
        and which ones should be frozen. If a parameter is not specified, the default value (if exists)
        for the parameter will be used.

        This allows you to easily initialize an experiment **without** wrapping every parameter in a list.

        Note:
            - For a given experiment, some parameters must be specified (e.g. the ``model`` parameter
              for OpenAI Chat Experiment). See the experiment's ``__init__`` method.
            - Each of ``test_parameters``'s values should be a ``list``, but not for ``frozen_parameters``.

        Args:
            test_parameters (dict[str, list]): parameters that are being tested. A list of multiple test values
                should be the value (e.g. ``{model: ["gpt-3.5-turbo", "gpt-4"], temperature: [0,0. 1.0]}``)
            frozen_parameters (dict): parameters that are intended to be frozen across different configuration.
                There is no need to wrap the value in a list. (e.g. ``{top_p: 1.0, presence_penalty: 0.0}``)

        Example:
            >>> from prompttools.experiment import OpenAIChatExperiment
            >>> test_parameters = {"model": ["gpt-3.5-turbo", "gpt-4"]}
            >>> messages = [{"role": "user", "content": "Who was the first president?"}]
            >>> frozen_parameters = {"top_p": 1.0, "messages": messages}
            >>> experiment = OpenAIChatExperiment.initialize(test_parameters, frozen_parameters)
        """
        frozen_parameters = {k: [v] for k, v in frozen_parameters.items()}
        return cls(**test_parameters, **frozen_parameters)

    def _is_chat(self):
        return False

    def _get_human_eval_listener(self, i: int) -> Callable:
        def listener(change):
            self.scores["feedback"][i] = change["new"]

        return listener

    def _get_comparison_listener(self, index: int) -> Callable:
        def listener(change):
            new_index = self.comparison_index_translation(index)
            self.scores["comparison"][new_index] = change["new"]

        return listener

    def _aggregate_comparison(
        self,
        table: pd.DataFrame,
        agg_column: int = 0,
        is_average: bool = False,
    ) -> Dict[str, int]:
        # TODO: This could be a group by
        prompt_scores = defaultdict(int)
        prompt_counts = defaultdict(int)
        for index, row in enumerate(table.iterrows()):
            key = str(row[agg_column])
            new_index = self.comparison_index_translation(index)
            prompt_scores[key] += self.scores["comparison"][new_index]
            prompt_counts[key] += 1
        if is_average:
            for k, v in prompt_scores.items():
                prompt_scores[k] = v / prompt_counts[k]
        sorted_scores = dict(sorted(prompt_scores.items(), key=lambda item: item[1], reverse=True))
        return sorted_scores

    def _aggregate_metric(
        self,
        table: pd.DataFrame,
        metric_name: str,
        agg_column: str,
        is_average: bool = False,
    ) -> Dict[str, int]:
        # TODO: This could be a group by

        prompt_scores = defaultdict(int)
        prompt_counts = defaultdict(int)
        for index, row in table.iterrows():
            key = str(row[agg_column])
            prompt_scores[key] += self.scores[metric_name][index]
            prompt_counts[key] += 1
        if is_average:
            for k, v in prompt_scores.items():
                prompt_scores[k] = v / prompt_counts[k]
        sorted_scores = dict(sorted(prompt_scores.items(), key=lambda item: item[1], reverse=True))
        return sorted_scores

    def prepare(self) -> None:
        r"""
        Creates argument combinations by taking the cartesian product of all inputs.
        """
        self.argument_combos = [dict(zip(self.all_args, val)) for val in itertools.product(*self.all_args.values())]

    def run(
        self,
        runs: int = 1,
    ) -> None:
        r"""
        Create tuples of input and output for every possible combination of arguments.

        Args:
            runs (int): number of times to execute each possible combination of arguments, defaults to 1.
        """
        if not self.argument_combos:
            logging.info("Preparing first...")
            self.prepare()
        for combo in self.argument_combos:
            for _ in range(runs):
                self.queue.enqueue(
                    self.completion_fn,
                    # We need to filter out defaults that are invalid JSON from the request
                    {k: v for k, v in combo.items() if (v is not None) and (v != float("inf"))},
                )
        self.results = self.queue.get_results()
        self.scores["latency"] = self.queue.get_latencies()
        if len(self.results) == 0:
            logging.error("No results. Something went wrong.")
            raise PromptExperimentException

    # TODO: Ideally, `eval_fn` should accept one row at a time, compute the metric, and add that to the row.
    def evaluate(
        self,
        metric_name: str,
        eval_fn: Callable,
        input_pairs: Optional[Dict[str, Tuple[str, Dict[str, str]]]] = None,
        input_key: Optional[str] = None,
        expected: Optional[List[str]] = None,
        eval_fn_kwargs: Optional[list[dict]] = None,
    ) -> None:
        """
        Using the given evaluation function, all input/response pairs are evaluated.

        Args:
            metric_name (str): name of the metric being computed
            eval_fn (Callable): an evaluation function that takes in (input, result, other_scores) and return a score
            input_pairs (Optional[Dict[str, Tuple[str, Dict[str, str]]]]): optional dictionary that holds the input data
                along with additional context or metadata for each input
            input_key (str): input key name as it exists within input argument (e.g. "messages", "prompt")
            expected (Optional[List[str]]): List of expected response to be passed to the evaluation function.
                The length of the list should be the same as the number of responses in the experiment's result.
            eval_fn_kwargs (Optional[list[dict]]): List of keyword args to be passed to the evaluation function.
                The length of the list should be the same as the number of responses in the experiment's result.
        """
        if not self.results:
            logging.info("Running first...")
            self.run()
        if metric_name in self.scores:
            logging.warning(metric_name + " is already present, skipping.")
            return

        if input_key is None:
            input_key = "messages" if self._is_chat() else "prompt"
        for i, result in enumerate(self.results):
            # Pass the messages and results into the eval function
            extracted_input = (
                input_pairs[self.argument_combos[i][input_key]] if input_pairs else self.argument_combos[i][input_key]
            )
            other_scores = {name: self.scores[name][i] for name in self.scores.keys() if name is not metric_name}
            if expected or eval_fn_kwargs:
                curr_eval_kwargs = {} if eval_fn_kwargs is None else eval_fn_kwargs[i]
                if expected:
                    curr_eval_kwargs["expected"] = expected[i]
                score = eval_fn(extracted_input, self._extract_responses(result), other_scores, **curr_eval_kwargs)
            else:
                score = eval_fn(
                    extracted_input,
                    self._extract_responses(result),
                    other_scores,
                )
            self.scores[metric_name].append(score)

    def get_table(self, pivot_data: Dict[str, object], pivot_columns: List[str], pivot: bool) -> pd.DataFrame:
        """
        This method creates a table of the experiment data. It can also be used
        to create a pivot table, or a table for gathering human feedback.

        Args:
            pivot_data (Dict[str, object]): dictionary that contains additional data or metadata related to the input
            pivot_columns (List[str]): two column names (first for pivot row, second for pivot column)
                that serve as indices the pivot table
            pivot (bool): determines whether to create a pivot table
        """
        input_key = "messages" if self._is_chat() else "prompt"
        data = {
            input_key: [str(combo[input_key]) for combo in self.argument_combos],
            "response": [self._extract_responses(result) for result in self.results],
            "latency": self.scores["latency"],
        }
        # Add scores for each eval fn, including feedback
        for metric_name, evals in self.scores.items():
            if metric_name != "comparison":
                data[metric_name] = evals
        # Add other args as cols if there was more than 1 input
        for k, args in self.all_args.items():
            if len(args) > 1:
                data[k] = [combo[k] for combo in self.argument_combos]
        if pivot_data:
            data[pivot_columns[0]] = [str(pivot_data[str(combo[input_key])][0]) for combo in self.argument_combos]
            data[pivot_columns[1]] = [str(pivot_data[str(combo[input_key])][1]) for combo in self.argument_combos]
        df = pd.DataFrame(data)
        if pivot:
            df = pd.pivot_table(
                df,
                values="response(s)",
                index=[pivot_columns[1]],
                columns=[pivot_columns[0]],
                aggfunc=lambda x: x.iloc[0],
            )
        return df

    def gather_feedback(self, pivot_data: Dict[str, object], pivot_columns: List[str]) -> None:
        """
        This method creates a table to gather human feedback from a notebook interface.

        Args:
            pivot_data (Dict[str, object]): dictionary that contains additional data or metadata related to the input
            pivot_columns (List[str]): two column names (first for pivot row, second for pivot column)
                that serve as indices the pivot table
        """
        if not self.results:
            logging.info("Running first...")
            self.run()
        if not is_interactive():
            logging.warning("This method only works in notebooks.")
            return
        self.scores["feedback"] = [1] * len(self.results)
        table = self.get_table(pivot_data, pivot_columns, pivot=False)
        self.feedback_widget_provider.set_pivot_columns(pivot_columns)
        items = self.feedback_widget_provider.get_header_widgets()
        for row in table.iterrows():
            items += self.feedback_widget_provider.get_row_widgets(*row)
        items += self.feedback_widget_provider.get_footer_widgets(table)
        self.feedback_widget_provider.display(items)

    def compare(self, primary_model: str, pivot_columns: List[str]) -> None:
        """
        This method creates a table to gather human feedback from a notebook interface.
        """
        if not self.results:
            logging.info("Running first...")
            self.run()
        if not is_interactive():
            logging.warning("This method only works in notebooks.")
            return
        table = self.get_table(pivot_data={}, pivot_columns=pivot_columns, pivot=True)
        self.scores["comparison"] = [1] * len(table)
        self.comparison_index_translation = lambda i: i * len(table.columns)
        self.comparison_widget_provider.set_models(table.columns)
        items = self.comparison_widget_provider.get_header_widgets()
        for index, row in enumerate(table.iterrows()):
            items += self.comparison_widget_provider.get_row_widgets(index, row[1])
        items += self.comparison_widget_provider.get_footer_widgets(table)
        self.comparison_widget_provider.display(items)

    def visualize(
        self,
        pivot_data: Optional[Dict[str, object]] = None,
        pivot_columns: Optional[List[str]] = None,
    ) -> None:
        """
        Creates and shows a table using the results produced.

        Args:
            pivot_data (Dict[str, object]): dictionary that contains additional data or metadata related to the input
            pivot_columns (List[str]): two column names (first for pivot row, second for pivot column)
                that serve as indices the pivot table
        """
        if not self.results:
            logging.info("Running first...")
            self.run()
        table = self.get_table(pivot_data, pivot_columns, pivot=pivot_columns is not None)
        if is_interactive():
            display.display(table)
        else:
            logging.getLogger().setLevel(logging.INFO)
            logging.info(tabulate(table, headers="keys", tablefmt="psql"))

    def aggregate(self, metric_name, column_name, is_average=False):
        """
        Aggregates a metric for a given column and displays to the user.

         Args:
            metric_name (str): metric to aggregate
            column_name (str): column to based the aggregation on
            is_average (bool): if ``True``, compute the average for the metric, else compute the total
        """
        if metric_name not in self.scores:
            logging.warning("Can't find " + metric_name + " in scores. Did you run `evaluate`?")
            return
        table = self.get_table(pivot_data=None, pivot_columns=None, pivot=False)
        sorted_scores = self._aggregate_metric(table, metric_name, column_name, is_average)
        if is_interactive():
            import matplotlib.pyplot as plt
            import os

            # Import style file, assumes same dir as experiment.py
            style_path = os.path.join(os.path.dirname(__file__), "style.mplstyle")
            plt.style.use(style_path)

            # Define the custom colors
            custom_colors = [
                "black",
                "#7e1e9c",
                "#15b01a",
                "#448ee4",
                "#ff7fa7",
                "#029386",
            ]

            plt.ylabel("Latency (s)")

            # Cycle through the custom colors when creating the bars
            for i, (label, value) in enumerate(sorted_scores.items()):
                plt.bar(i, value, align="center", color=custom_colors[i % len(custom_colors)])

            plt.xticks(range(len(sorted_scores)), list(sorted_scores.keys()))
            plt.show()

    def rank(
        self,
        pivot_data: Dict[str, object],
        pivot_columns: List[str],
        metric_name: str,
        is_average: bool,
    ) -> Dict[str, int]:
        """
        Using pivot data, groups the data by the first pivot column to
        get scores, and sorts descending. For example, using pivot data of
        (prompt_template, user_input), a metric of latency, and is_average=True,
        we rank prompt templates by their average latency in the test set.

        Args:
            pivot_data (Dict[str, object]): dictionary that contains additional data or metadata related to the input
            pivot_columns (List[str]): two column names (first for pivot row, second for pivot column)
                that serve as indices the pivot table
            metric_name (str): metric to aggregate over
            is_average (bool): if ``True``, compute the average for the metric, else compute the total
        """
        if metric_name not in self.scores:
            logging.warning("Can't find " + metric_name + " in scores. Did you run `evaluate`?")
            return
        table = self.get_table(pivot_data, pivot_columns, pivot=False)
        sorted_scores = self._aggregate_metric(table, metric_name, pivot_columns[0], is_average)
        return sorted_scores

    @staticmethod
    def _extract_responses(output: Dict[str, object]) -> list[str]:
        raise NotImplementedError("This should be implemented by a subclass of `Experiment`.")

    def to_csv(
        self,
        path: str,
        pivot_data: Optional[Dict[str, object]] = None,
        pivot_columns: Optional[List[str]] = None,
        **kwargs,
    ):
        r"""
        Export the results to a CSV file. If the experiment has not been executed, it will run.

        Args:
            path (str): path/buffer to write the CSV output
            pivot_data (Dict[str, object]): optional dictionary that contains additional data or metadata
                related to the input
            pivot_columns (List[str]): optional two column names (first for pivot row, second for pivot column)
                that serve as indices the pivot table
            **kwargs: optional arguments passed to ``pd.DataFrame.to_csv()``
        """
        if not self.results:
            logging.info("Running first...")
            self.run()
        table = self.get_table(pivot_data, pivot_columns, pivot=pivot_columns is not None)
        table.to_csv(path, **kwargs)

    def to_pandas_df(self):
        r"""
        Return the results as a ``pandas.DataFrame``. If the experiment has not been executed, it will run.
        """
        if not self.results:
            logging.info("Running first...")
            self.run()
        return self.get_table({}, [], pivot=False)

    def to_json(
        self,
        path: Optional[str] = None,
        pivot_data: Optional[Dict[str, object]] = None,
        pivot_columns: Optional[List[str]] = None,
        **kwargs,
    ):
        r"""
        Export the results to a JSON file. If the experiment has not been executed, it will run.

        Args:
            path (Optional[str]): path/buffer to write the JSON output, defaults to ``None`` which returns
                the JSON as a `dict`
            pivot_data (Dict[str, object]): optional dictionary that contains additional data or metadata
                related to the input
            pivot_columns (List[str]): optional two column names (first for pivot row, second for pivot column)
                that serve as indices the pivot table
            **kwargs: optional arguments passed to ``pd.DataFrame.to_json()``
        """
        if not self.results:
            logging.info("Running first...")
            self.run()
        table = self.get_table(pivot_data, pivot_columns, pivot=pivot_columns is not None)
        if path is None:
            return table.to_json(**kwargs)
        else:
            return table.to_json(path, **kwargs)

    def to_lora_json(
        self,
        instruction_extract: Union[str, Callable],
        input_extract: Union[str, Callable],
        output_extract: Union[str, Callable],
        path: Optional[str] = None,
        **kwargs,
    ):
        r"""
        Export the results to a LoRA-format JSON file for fine-tuning.
        If the experiment has not been executed, it will run.

        Args:
            instruction_extract (Union[str, Callable]): column name, or an extractor function that will accept a row
                of the result table and return a value assigned to ``"instruction"`` entry in the JSON file
            input_extract (Union[str, Callable]): column name, or an extractor function that will accept a row
                of the result table and return a value assigned to ``"input"`` entry in the JSON file
            output_extract (Union[str, Callable]): column name, or an extractor function that will accept a row
                of the result table and return a value assigned to ``"output"`` entry in the JSON file
            path (Optional[str]): path/buffer to write the JSON output, defaults to ``None`` which returns
                the JSON as a `dict`
            **kwargs: optional arguments passed to ``pd.DataFrame.to_json()``
        """
        if not self.results:
            logging.info("Running first...")
            self.run()
        if isinstance(instruction_extract, str):
            instruction_extract = itemgetter(instruction_extract)
        if isinstance(input_extract, str):
            input_extract = itemgetter(input_extract)
        if isinstance(output_extract, str):
            output_extract = itemgetter(output_extract)
        df = self.to_pandas_df()
        extracted_data = df.apply(
            lambda row: {
                "instruction": instruction_extract(row),
                "input": input_extract(row),
                "output": output_extract(row),
            },
            axis=1,
        )
        if "orient" not in kwargs:
            kwargs["orient"] = "records"
        if "indent" not in kwargs:
            kwargs["indent"] = 2

        if path:
            extracted_data.to_json(path, **kwargs)
        else:
            return extracted_data.to_json(**kwargs)

    def to_mongo_db(self, mongo_uri: str, database_name: str, collection_name: str) -> None:
        r"""
        Insert the results of the experiment into MongoDB for persistence.

        Note:
            - You need to install the ``pymongo`` package to use this method.
            - You need to run a local or remote instance of MongoDB in order to store the data.

        Args:
            mongo_uri (str): a connection string to the target MongoDB
            database_name (str): name of the MongoDB database
            collection_name (str): name of the MongoDB collection
        """
        if pymongo is None:
            raise ModuleNotFoundError(
                "Package `pymongo` is required to be installed to use this method."
                "Please use `pip install pymongo` to install the package"
            )
        if not self.results:
            logging.info("Running first...")
            self.run()
        client = pymongo.MongoClient(mongo_uri)
        db = client[database_name]
        collection = db[collection_name]
        collection.insert_many(self.results)
        logging.info(f"Inserted results in {database_name}'s collection {collection_name}.")
        client.close()

    def to_markdown(self):
        if not self.results:
            logging.info("Running first...")
            self.run()
        markdown = self.to_pandas_df().to_markdown()
        return markdown

    def _get_model_names(self):
        pass

    def _get_prompts(self):
        pass
