from pathlib import Path

import pytest
import torch
from transformers import PreTrainedModel, PreTrainedTokenizerBase

from deception_detection.activations import Activations, PairedActivations
from deception_detection.data import GTCities, RepeHonestDishonestFacts
from deception_detection.detectors import (
    CovarianceAdjustedMMSDetector,
    DirectionDetector,
    LATDetector,
    LogisticRegressionDetector,
    MLPDetector,
    MMSDetector,
)
from deception_detection.metrics import get_auroc
from deception_detection.tokenized_data import TokenizedDataset


@pytest.fixture(scope="module")
def detector_test_data(gemma_model: PreTrainedModel, gemma_tokenizer: PreTrainedTokenizerBase):
    fit_data = RepeHonestDishonestFacts(variant="plain", shuffle_upon_init=False)
    n_train = int(len(fit_data) * 0.7)
    train_data, test_data = fit_data[:n_train], fit_data[n_train:]

    train_split_acts = PairedActivations.from_dataset(
        train_data, gemma_tokenizer, gemma_model, enforce_equal_detect_lengths=True
    )
    test_split_acts = PairedActivations.from_dataset(
        test_data, gemma_tokenizer, gemma_model, enforce_equal_detect_lengths=True
    )

    honest_train_acts, deceptive_train_acts = train_split_acts.get_all_acts()
    honest_test_acts, deceptive_test_acts = test_split_acts.get_all_acts()

    return {
        "honest_train_acts": honest_train_acts,
        "deceptive_train_acts": deceptive_train_acts,
        "honest_test_acts": honest_test_acts,
        "deceptive_test_acts": deceptive_test_acts,
    }


@pytest.mark.gpu
@pytest.mark.parametrize(
    "detector_class",
    [
        MMSDetector,
        LATDetector,
        LogisticRegressionDetector,
        MLPDetector,
        CovarianceAdjustedMMSDetector,
    ],
)
def test_test_classification_accuracy(
    detector_test_data: dict[str, Activations],
    detector_class: type[DirectionDetector],
):
    """Takes the dataset we use to train detectors, splits into train and test, and ensures we get
    high accuracy on the test set when fitting to the train set."""

    detector = detector_class(layers=list(range(4, 12)))
    detector.fit(
        detector_test_data["honest_train_acts"], detector_test_data["deceptive_train_acts"]
    )

    test_honest_scores = detector.score(detector_test_data["honest_test_acts"])
    test_dishonest_scores = detector.score(detector_test_data["deceptive_test_acts"])

    assert not any([score_row.isnan().any() for score_row in test_honest_scores.scores])
    assert not any([score_row.isnan().any() for score_row in test_dishonest_scores.scores])

    mean_honest_scores = torch.tensor(
        [score.mean() for score in test_honest_scores.scores if len(score) > 0]
    )
    mean_dishonest_scores = torch.tensor(
        [score.mean() for score in test_dishonest_scores.scores if len(score) > 0]
    )

    accuracy = (mean_honest_scores > mean_dishonest_scores).float().mean()
    assert accuracy > 0.99, f"Accuracy is only {accuracy}"


@pytest.mark.gpu
def test_generalization(
    gemma_model: PreTrainedModel,
    gemma_tokenizer: PreTrainedTokenizerBase,
):
    # Prepare RepeHonestDishonestFacts (train set)
    train_data = RepeHonestDishonestFacts()
    split_acts = PairedActivations.from_dataset(
        train_data, gemma_tokenizer, gemma_model, layers=list(range(9, 15))
    )
    honest_train_acts = split_acts.honest_acts
    deceptive_train_acts = split_acts.deceptive_acts

    # Prepare Cities dataset (test set)
    test_data = GTCities()
    split_test_acts = PairedActivations.from_dataset(
        test_data, gemma_tokenizer, gemma_model, layers=list(range(9, 15))
    )
    honest_test_acts = split_test_acts.honest_acts
    deceptive_test_acts = split_test_acts.deceptive_acts

    # Train and evaluate detector
    detector = MMSDetector(layers=list(range(9, 15)))
    detector.fit(deceptive_train_acts, honest_train_acts)

    test_honest_scores = detector.score(honest_test_acts)
    mean_honest_scores = torch.tensor([score.mean() for score in test_honest_scores.scores])
    test_dishonest_scores = detector.score(deceptive_test_acts)
    mean_dishonest_scores = torch.tensor([score.mean() for score in test_dishonest_scores.scores])

    auroc = get_auroc(mean_honest_scores, mean_dishonest_scores)

    assert auroc > 0.95, f"AUROC is only {auroc}"


@pytest.mark.parametrize("detector_class", [MMSDetector, LATDetector, LogisticRegressionDetector])
def test_save_load_detector(tmp_path: Path, detector_class: type[DirectionDetector]):
    # Create a detector with random directions
    detector = detector_class(layers=[0, 1, 2])
    detector.directions = torch.rand(3, 768)  # Assuming 768 is the embedding size

    # Save the detector
    save_path = tmp_path / "test_detector.pkl"
    detector.save(save_path)

    # Load the detector
    loaded_detector = detector_class.load(save_path)

    # Check if the loaded detector has the same attributes
    assert detector.layers == loaded_detector.layers
    assert loaded_detector.directions is not None

    # Ensure both tensors are on the same device before comparison
    assert detector.directions.device == loaded_detector.directions.device
    assert torch.allclose(detector.directions, loaded_detector.directions)


@pytest.fixture(scope="module")
def random_test_data(gemma_tokenizer: PreTrainedTokenizerBase) -> dict[str, Activations]:
    """Alternative to detector_test_data that doesn't need a gpu."""
    layers = [4]
    emb_size = 24
    fit_data = RepeHonestDishonestFacts()[:50]
    tokenized_dataset = TokenizedDataset.from_dataset(fit_data, gemma_tokenizer)

    def get_acts() -> Activations:
        shape = tokenized_dataset.tokens.shape + (len(layers), emb_size)
        return Activations(
            all_acts=torch.randn(shape),
            tokenized_dataset=tokenized_dataset,
            layers=layers,
        )

    acts = get_acts()
    return {
        "honest_train_acts": acts,
        "deceptive_train_acts": acts,
        "honest_test_acts": acts,
        "deceptive_test_acts": acts,
    }


@pytest.mark.parametrize("detector_class", [MMSDetector, LogisticRegressionDetector])
def test_fit_detector_with_random_acts(
    detector_class: type[DirectionDetector],
    random_test_data: dict[str, Activations],
):
    detector = detector_class(layers=[4])
    detector.fit(random_test_data["honest_train_acts"], random_test_data["deceptive_train_acts"])
    assert detector.directions is not None
    emb_size = random_test_data["honest_train_acts"].all_acts.shape[-1]
    assert detector.directions.shape == (len(detector.layers), emb_size)

    detector.score(random_test_data["honest_test_acts"])


@pytest.mark.parametrize("detector_class", [MMSDetector, LogisticRegressionDetector])
def test_detector_with_sparse_activations(
    random_test_data: dict[str, Activations], detector_class: type[DirectionDetector]
):
    """Tests that detectors can score activations with sparse tensors"""

    # Train detector on dense activations
    if detector_class == LogisticRegressionDetector:
        detector = LogisticRegressionDetector(layers=[4], normalize=False)
    else:
        detector = detector_class(layers=[4])
    detector.fit(random_test_data["honest_train_acts"], random_test_data["deceptive_train_acts"])

    # Save test activations in sparse format
    dense_acts = random_test_data["honest_test_acts"]
    sparse_acts = Activations(
        all_acts=dense_acts.all_acts.to_sparse(),
        tokenized_dataset=dense_acts.tokenized_dataset,
        layers=dense_acts.layers,
    )

    # Score using sparse activations
    sparse_scores = detector.score(sparse_acts)
    dense_scores = detector.score(dense_acts)

    # Compare scores between sparse and dense versions
    for sparse_score, dense_score in zip(sparse_scores.scores, dense_scores.scores, strict=True):
        assert torch.allclose(sparse_score, dense_score, atol=1e-5)
