import pandas as pd

from deception_detection.experiment import Experiment
from deception_detection.prompt_scorer import PromptScorerMethod, create_prompt_scorer
from deception_detection.repository import DatasetRepository
from deception_detection.results import token_html
from deception_detection.results.plot import (
    make_dataset_specific_plots,
    plot_roc_comparison,
    plot_roc_curve_with_auroc,
    violin_plot,
)
from deception_detection.types import Label

REPOSITORY = DatasetRepository()
FPRS = [0.01, 0.001]


class ExperimentPostprocessor:
    def __init__(
        self,
        experiment: Experiment,
        scoring_method: PromptScorerMethod = "mean",
        fprs: list[float] = FPRS,
        default_fpr: float = 0.01,
        include_black_box: bool = False,
        control_dataset_name: str = "alpaca__plain",
    ):
        """
        Postprocess an experiment. Produce the standard plots, etc.
        Add things here once they become useful.

        Args:
            experiment: The experiment to postprocess.
            scoring_method: The method to use to score the experiment.
            fprs: The false positive rates used to produce plots that require multiple FPRs.
            default_fpr: The default false positive rate used to produce plots that require a
            single FPR.
            include_black_box: Whether to include black box scores in the results table.
        """
        self.experiment = experiment
        self.cfg = experiment.cfg
        self.scorer = create_prompt_scorer(scoring_method, experiment, fprs, default_fpr, control_dataset_name=control_dataset_name)
        self.include_black_box = include_black_box
        self.control_dataset_name = control_dataset_name

    def run(self) -> None:
        """
        Run the postprocessing.
        """
        self.make_results_table()
        self.make_all_plots()

    def make_all_plots(self) -> None:
        """
        Make all the standard plots.
        """
        (self.experiment.cfg.folder / "plots").mkdir(parents=True, exist_ok=True)
        token_html.plot_control_token_scores(self.experiment, self.scorer, num_samples=20)
        plot_roc_curve_with_auroc(self.experiment, self.scorer, compare_to="honest", control_dataset_name=self.control_dataset_name)
        plot_roc_curve_with_auroc(self.experiment, self.scorer, compare_to="control", x_max=0.03, control_dataset_name=self.control_dataset_name)
        violin_plot(self.experiment, self.scorer)
        plot_roc_comparison(self.experiment, self.scorer, control_dataset_name=self.control_dataset_name)
        make_dataset_specific_plots(self.experiment, self.scorer)

    def make_results_table(self) -> None:
        """
        Make a results table of summary statistics
        - AUROCs between honest and deceptive dialogues
        - TPRs and FPRs of datasets at various FPRs on control datasets.
        """
        df = self.get_results_dataframe()
        new_columns = {f"TPR_{fpr:.4f}": f"TPR @ {fpr:.2%}" for fpr in self.scorer.fprs}
        new_columns.update({f"FPR_{fpr:.4f}": f"FPR @ {fpr:.2%}" for fpr in self.scorer.fprs})
        new_dataset_names: dict[str, str] = dict()  # {"geometry_of_truth_cities": "cities"}
        df.rename(index=new_dataset_names, columns=new_columns, inplace=True)
        print(
            df.to_markdown(
                floatfmt=(
                    "",
                    ".0f",
                    ".0f",
                    ".0f",
                    ".3f",
                    ".3f",
                    ".3f",
                    *[".1%" for _ in self.scorer.fprs * 2],
                )
            )
        )
        csv_file = self.experiment.cfg.folder / "results_table.csv"
        df.to_csv(csv_file)

    def get_results_dataframe(self) -> pd.DataFrame:
        assert self.scorer is not None, "Scorer must be initialized before getting results"
        eval_scores = self.experiment.get_eval_scores()

        table_contents = []
        for dataset_name, scores in eval_scores.items():
            split_scores = scores.get_scores_by_label()

            row = {
                "n_honest": len(split_scores[Label.HONEST]),
                "n_ambiguous": len(split_scores[Label.AMBIGUOUS]),
                "n_deceptive": len(split_scores[Label.DECEPTIVE]),
                # "n_skipped": len(split_scores[Label.SKIP]),
                "auroc": self.scorer.auroc(scores),
                "auroc_vs_alpaca": self.scorer.auroc_alpaca(scores),
                "balanced_accuracy_at_fpr": self.scorer.balanced_accuracy_at_fpr(scores),
                # "balanced_accuracy": self.scorer.balanced_accuracy(scores)

            }
            for threshold in self.scorer.fprs:
                fpr = self.scorer.positive_rate(split_scores[Label.HONEST].scores, threshold)
                fnr = 1 - self.scorer.positive_rate(split_scores[Label.DECEPTIVE].scores, threshold)
                row[f"recall {threshold:.1%}"] = 1 - fnr
                row[f"fpr {threshold:.1%}"] = fpr

            if self.include_black_box:
                bb_scores = self.experiment.get_black_box_scores()[dataset_name]
                row["black_box_auroc"] = self.scorer.auroc(bb_scores)
                row["black_box_auroc_vs_alpaca"] = float("nan")

            table_contents.append(row)

        return pd.DataFrame(table_contents, index=list(eval_scores.keys()))  # type: ignore
