import shutil
from collections.abc import Generator
from pathlib import Path
from typing import Literal

import matplotlib.pyplot as plt
import pandas as pd
import pytest
import torch

from deception_detection.detectors import DirectionDetector
from deception_detection.experiment import Experiment, ExperimentConfig, Scores
from deception_detection.models import ModelName
from deception_detection.postprocess import ExperimentPostprocessor


class TestExperiment:
    @pytest.fixture(scope="class", params=["mms", "lr"])
    def experiment(self, request: pytest.FixtureRequest) -> Generator[Experiment, None, None]:
        cfg = ExperimentConfig(
            method=request.param,
            model_name=ModelName.LLAMA_1B,
            train_data="repe_honesty__plain",
            eval_data=["ai_liar__original_with_answers"],
            eval_on_policy=False,
            train_on_policy=False,
            control_data=["alpaca__plain"],
            id="pytest",
            detect_only_start_of_turn=False,
            control_dataset_size=10,
            val_fraction=0.05,  # faster since less eval data in validation set
            detect_layers=[10, 12],  # faster, and can truncate model
        )

        exp = Experiment(cfg)
        exp.train_dataset = exp.train_dataset[:10]
        yield exp
        shutil.rmtree(exp.cfg.folder)

    def test_experiment_creation(self, experiment: Experiment) -> None:
        assert experiment is not None

    @pytest.mark.gpu
    def test_get_detector(self, experiment: Experiment) -> None:
        detector = experiment.get_detector()
        assert detector is not None
        assert Path(experiment.cfg.folder / "detector.pt").exists()

    @pytest.mark.gpu
    def test_get_eval_scores(self, experiment: Experiment) -> None:
        scores = experiment.get_eval_scores()
        assert len(scores) > 0
        assert Path(experiment.cfg.folder / "scores.json").exists()
        for _, dataset_scores in scores.items():
            assert isinstance(dataset_scores, Scores)
            assert len(dataset_scores.scores) > 0

    @pytest.mark.gpu
    def test_get_eval_scores_all_tokens(self, experiment: Experiment) -> None:
        all_token_scores = experiment.get_eval_scores(all_tokens=True)
        assert len(all_token_scores) > 0
        assert Path(experiment.cfg.folder / "scores_all_tokens.json").exists()
        for _, dataset_scores in all_token_scores.items():
            assert isinstance(dataset_scores, Scores)
            assert len(dataset_scores.scores) > 0

        scores = experiment.get_eval_scores()
        for dataset in all_token_scores:
            dataset_scores = scores[dataset]
            dataset_all_token_scores = all_token_scores[dataset]
            assert len(dataset_all_token_scores.scores) >= len(dataset_scores.scores)

    @pytest.mark.gpu
    def test_get_control_scores(self, experiment: Experiment) -> None:
        control_scores = experiment.get_control_scores()
        assert isinstance(control_scores, dict)
        assert Path(experiment.cfg.folder / "control_scores.json").exists()
        assert "alpaca__plain" in control_scores
        assert len(control_scores["alpaca__plain"].scores) == 10

    @pytest.mark.gpu
    def test_get_control_scores_all_tokens(self, experiment: Experiment) -> None:
        control_scores_all_tokens = experiment.get_control_scores(all_tokens=True)
        assert isinstance(control_scores_all_tokens, dict)
        assert Path(experiment.cfg.folder / "control_scores_all_tokens.json").exists()

    def test_get_black_box_scores(self, experiment: Experiment) -> None:
        black_box_scores = experiment.get_black_box_scores(mock=True)
        assert isinstance(black_box_scores, dict)
        assert Path(experiment.cfg.folder / "black_box_scores.json").exists()
        for dataset, scores in zip(
            experiment.eval_datasets, black_box_scores.values(), strict=True
        ):
            assert isinstance(scores, Scores)
            assert len(scores) == len(dataset)

    @pytest.mark.gpu
    @pytest.mark.parametrize("scoring_method", ["mean", "max", "relu"])
    def test_postprocessor(
        self, experiment: Experiment, scoring_method: Literal["mean", "max", "relu"]
    ) -> None:
        postprocessor = ExperimentPostprocessor(experiment, scoring_method=scoring_method)
        postprocessor.run()
        plt.close("all")

        # Check if results table was created
        assert Path(experiment.cfg.folder / "results_table.csv").exists()

        # Check if plots were created
        assert Path(experiment.cfg.folder / "plots" / f"violin_plot_{scoring_method}.png").exists()
        assert Path(experiment.cfg.folder / "token_scores" / "alpaca__plain.html").exists()

    @pytest.mark.gpu
    def test_results_loading(self, experiment: Experiment) -> None:
        cfg = ExperimentConfig.from_path(experiment.cfg.folder)
        new_experiment = Experiment(cfg)
        self.assert_scores_equal(experiment.get_eval_scores(), new_experiment.get_eval_scores())

    def assert_scores_equal(self, scores1: dict[str, Scores], scores2: dict[str, Scores]) -> None:
        assert scores1.keys() == scores2.keys()
        for dataset in scores1:
            assert all(
                [
                    torch.allclose(s1, s2)
                    for s1, s2 in zip(
                        scores1[dataset].scores, scores2[dataset].scores, strict=False
                    )
                ]
            )


@pytest.fixture(scope="function")
def sae_experiment() -> Generator[Experiment, None, None]:
    cfg = ExperimentConfig(
        method="mms",
        model_name=ModelName.GEMMA_9B,
        train_data="repe_honesty__plain",
        train_on_policy=False,
        eval_on_policy=False,
        eval_data=[],
        control_data=["alpaca__plain"],
        use_local_sae_acts=True,
        sae_latent_whitelist=[0, 1, 2, 3],
        detect_layers=[31],
        control_dataset_size=4,
    )
    exp = Experiment(cfg)
    exp.train_dataset = exp.train_dataset[:10]
    yield exp
    shutil.rmtree(exp.cfg.folder)


@pytest.mark.xfail(
    reason=(
        "Currently failing in plotting function for unclear reasons -- "
        "maybe df not populated sufficiently?"
    )
)
@pytest.mark.skip(reason="This test is very slow (1 min, >25% total time) and not really necessary")
@pytest.mark.gpu
def test_sae_experiment(sae_experiment: Experiment) -> None:
    detector = sae_experiment.get_detector()
    assert isinstance(detector, DirectionDetector)
    assert detector is not None
    assert detector.directions is not None
    assert sae_experiment.cfg.sae_latent_whitelist is not None
    n_sae_latents = len(sae_experiment.cfg.sae_latent_whitelist)
    n_layers = len(sae_experiment.cfg.detect_layers)
    assert isinstance(detector.directions, torch.Tensor)
    assert detector.directions.shape == (n_layers, n_sae_latents)

    # Test postprocessor with SAE experiment
    postprocessor = ExperimentPostprocessor(sae_experiment)
    postprocessor.run()

    # Check if results table was created
    results_table_path = sae_experiment.cfg.folder / "results_table.csv"
    assert results_table_path.exists()

    # Check contents of results table
    df = pd.read_csv(results_table_path, index_col=0)
    cols = [
        "n honest",
        "n deceptive",
        "AUROC vs. honest",
        "AUROC vs. alpaca",
        "TPR @ 1.00%",
        "FPR @ 1.00%",
    ]
    for col in cols:
        assert col in df.columns
