from typing import Literal

import pytest
import torch

from deception_detection.experiment import Experiment, Scores
from deception_detection.prompt_scorer import PromptScorer, ReluPromptScorer, create_prompt_scorer
from deception_detection.types import Label


@pytest.fixture
def mock_experiment() -> Experiment:
    torch.manual_seed(0)
    random_scores = [torch.rand(i) for i in range(50, 200)]

    class MockExperiment(Experiment):
        def __init__(self):
            pass

        def get_control_scores(self, all_tokens: bool = False):
            return {"alpaca__plain": Scores(scores=random_scores)}

    return MockExperiment()


@pytest.mark.parametrize("method", ["mean", "max", "relu"])
def test_prompt_scorer_initialization(
    mock_experiment: Experiment, method: Literal["mean", "max", "relu"]
):
    scorer = create_prompt_scorer(method=method, experiment=mock_experiment, fprs=[0.05, 0.01])
    assert isinstance(scorer.prompt_score_threshold(0.05), float)
    assert isinstance(scorer, PromptScorer)


def test_prompt_scorer_reduce_mean(mock_experiment: Experiment):
    scorer = create_prompt_scorer(method="mean", experiment=mock_experiment, fprs=[0.05, 0.01])
    scores = [torch.tensor([0.0, 0.2]) + x for x in torch.arange(0, 1, 0.1)]
    reduced = scorer.reduce(scores)
    assert torch.allclose(reduced, torch.arange(0, 1, 0.1) + 0.1)


def test_prompt_scorer_reduce_max(mock_experiment: Experiment):
    scorer = create_prompt_scorer(method="max", experiment=mock_experiment, fprs=[0.05, 0.01])
    scores = [
        torch.tensor([0.1, 0.3, 0.2]),
        torch.tensor([0.4, 0.2, 0.5]),
        torch.tensor([0.1, 0.6, 0.3, 100]),
    ]
    reduced = scorer.reduce(scores)
    assert torch.allclose(reduced, torch.tensor([0.3, 0.5, 100]))


def test_prompt_scorer_reduce_max_empty_tensor(mock_experiment: Experiment):
    scorer = create_prompt_scorer(method="max", experiment=mock_experiment, fprs=[0.05, 0.01])
    scores = [torch.tensor([]), torch.tensor([0.3, 0.4])]
    reduced = scorer.reduce(scores)
    assert torch.isnan(reduced[0])


def test_prompt_scorer_reduce_relu(mock_experiment: Experiment):
    scorer = create_prompt_scorer(method="relu", experiment=mock_experiment, fprs=[0.05, 0.1])
    assert isinstance(scorer, ReluPromptScorer)
    assert scorer.relu_threshold is not None
    scores = [
        scorer.relu_threshold + torch.rand(500),
        scorer.relu_threshold + torch.rand(500) - 1,
        scorer.relu_threshold + torch.rand(500) - 0.5,
    ]
    reduced = scorer.reduce(scores)
    assert reduced[0] == pytest.approx(0.5, abs=0.02)
    assert reduced[1] == pytest.approx(0, abs=0.02)
    assert reduced[2] == pytest.approx(0.125, abs=0.02)


def test_prompt_scorer_prompt_scores(mock_experiment: Experiment):
    scorer = create_prompt_scorer(method="mean", experiment=mock_experiment, fprs=[0.05, 0.1])
    scores = [torch.tensor([0.1, 0.2, 0.3]), torch.tensor([0.2, 0.3, 0.4])]
    prompt_scores = scorer.reduce(scores)
    assert prompt_scores.shape == torch.Size([2])


def test_prompt_scorer_positive_rate(mock_experiment: Experiment):
    scorer = create_prompt_scorer(method="mean", experiment=mock_experiment, fprs=[0.05, 0.01])
    scores = [torch.tensor([0.1, 0.2, 0.3]), torch.tensor([0.2, 0.3, 0.4])]
    positive_rate = scorer.positive_rate(scores)
    assert 0 <= positive_rate <= 1


def test_prompt_scorer_auroc(mock_experiment: Experiment):
    scorer = create_prompt_scorer(method="mean", experiment=mock_experiment, fprs=[0.05, 0.1])

    scores = Scores(scores=[[0.1, 0.2], [0.7]], labels=[Label.HONEST, Label.DECEPTIVE])
    assert scorer.auroc(scores) == 1

    scores = Scores(scores=[[1.0], [0.5]], labels=[Label.HONEST, Label.DECEPTIVE])
    assert scorer.auroc(scores) == 0

    honest_scores = Scores(scores=[torch.rand(10) for _ in range(100)], labels=[Label.HONEST] * 100)
    deceptive_scores = Scores(
        scores=[torch.rand(10) for _ in range(100)], labels=[Label.DECEPTIVE] * 100
    )
    scores = Scores.from_list([honest_scores, deceptive_scores])
    assert scorer.auroc(scores) == pytest.approx(0.5, 0.05)


def test_prompt_scorer_auroc_alpaca(mock_experiment: Experiment):
    scorer = create_prompt_scorer(method="mean", experiment=mock_experiment, fprs=[0.05, 0.1])

    random_scores = [torch.rand(i) for i in range(10, 200)]

    scores_0_1 = Scores(
        scores=[torch.tensor([3.0, 1.0])] + random_scores,
        labels=[Label.HONEST] + [Label.DECEPTIVE] * len(random_scores),
    )
    assert scorer.auroc_alpaca(scores_0_1) == pytest.approx(0.5, abs=0.05)

    scores_neg = Scores(
        scores=[s - 2 for s in random_scores], labels=[Label.DECEPTIVE] * len(random_scores)
    )
    assert scorer.auroc_alpaca(scores_neg) == 0

    scores_large = Scores(
        scores=[s + 2 for s in random_scores], labels=[Label.DECEPTIVE] * len(random_scores)
    )
    assert scorer.auroc_alpaca(scores_large) == 1


def test_prompt_scorer_normalize_scores(mock_experiment: Experiment):
    scorer = create_prompt_scorer(method="mean", experiment=mock_experiment, fprs=[0.05, 0.01])
    raw_scores = [torch.tensor([0.1, 0.2, 0.3]), torch.tensor([0.2, 0.3, 0.4])]
    scores = Scores(scores=raw_scores)
    normalized = scorer.normalize_scores(scores)
    assert len(normalized) == len(scores)
    assert all(torch.is_tensor(score) for score in normalized)


def test_prompt_scorer_invalid_method(mock_experiment: Experiment):
    with pytest.raises(ValueError):
        create_prompt_scorer(method="invalid", experiment=mock_experiment, fprs=[0.05, 0.1])  # type: ignore
