from pathlib import Path
from typing import cast

import pytest
import torch
from transformers import PreTrainedModel, PreTrainedTokenizerBase

from deception_detection.activations import Activations, PairedActivations
from deception_detection.data import DialogueDataset
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.types import Dialogue, Label, Message


@pytest.fixture
def tokenized_dataset(gemma_tokenizer: PreTrainedTokenizerBase) -> TokenizedDataset:
    dialogue = Dialogue(
        [
            Message("user", "How are you?", False),
            Message("assistant", "Very well!", True),
        ]
    )
    return TokenizedDataset.from_dialogue_list(
        [dialogue], gemma_tokenizer, padding={"gemma": {"left": 4, "right": 4}}
    )


@pytest.mark.gpu
def test_activations(gemma_model: PreTrainedModel, tokenized_dataset: TokenizedDataset) -> None:
    activations = Activations.from_model(gemma_model, tokenized_dataset, save_logits=True)

    n_batch: int = len(tokenized_dataset)
    seq_len: int = tokenized_dataset.tokens.shape[1]
    n_layers: int = gemma_model.config.num_hidden_layers
    d_emb: int = gemma_model.config.hidden_size
    n_vocab: int = gemma_model.config.vocab_size

    assert tokenized_dataset.detection_mask is not None
    n_detect_toks = tokenized_dataset.detection_mask.sum().item()

    assert activations.all_acts.shape == (n_batch, seq_len, n_layers + 1, d_emb)
    assert activations.get_masked_activations().shape == (n_detect_toks, n_layers + 1, d_emb)

    assert activations.logits is not None
    assert activations.logits.shape == (n_batch, seq_len, n_vocab)
    assert activations.get_masked_logits().shape == (n_detect_toks, n_vocab)

    assert not activations.all_acts.requires_grad
    assert not activations.logits.requires_grad

    assert not activations.all_acts.isnan().any()
    assert not activations.logits.isnan().any()


@pytest.mark.gpu
def test_subset_layers(gemma_model: PreTrainedModel, tokenized_dataset: TokenizedDataset) -> None:
    layers = [5, 10]
    activations = Activations.from_model(gemma_model, tokenized_dataset, layers=layers)

    n_batch: int = len(tokenized_dataset)
    seq_len: int = tokenized_dataset.tokens.shape[1]
    d_emb: int = gemma_model.config.hidden_size
    assert activations.layers == layers
    assert activations.all_acts.shape == (n_batch, seq_len, len(layers), d_emb)

    assert tokenized_dataset.detection_mask is not None
    n_detect_toks = tokenized_dataset.detection_mask.sum().item()
    assert activations.get_masked_activations().shape == (n_detect_toks, len(layers), d_emb)


@pytest.mark.gpu
def test_save_load_activations(
    gemma_model: PreTrainedModel, tokenized_dataset: TokenizedDataset, tmp_path: Path
) -> None:
    activations = Activations.from_model(
        gemma_model, tokenized_dataset, save_logits=True, layers=[5, 10]
    )

    act_file = tmp_path / "activations.pt"
    # Save the activations
    activations.save(act_file)

    # Load the activations
    loaded_activations = Activations.load(act_file)

    # Check if the loaded activations match the original
    assert torch.equal(activations.all_acts, loaded_activations.all_acts)
    assert activations.logits is not None and loaded_activations.logits is not None
    assert torch.equal(activations.logits, loaded_activations.logits)
    assert activations.layers == loaded_activations.layers

    # Check if the tokenized dataset is correctly loaded
    orig_dataset = activations.tokenized_dataset
    loaded_dataset = loaded_activations.tokenized_dataset
    assert orig_dataset.tokens.shape == loaded_dataset.tokens.shape
    assert torch.equal(orig_dataset.tokens, loaded_dataset.tokens)
    assert orig_dataset.detection_mask is not None
    assert loaded_dataset.detection_mask is not None
    assert torch.equal(orig_dataset.detection_mask, loaded_dataset.detection_mask)
    assert torch.equal(orig_dataset.attention_mask, loaded_dataset.attention_mask)

    # Check if the loaded activations produce the same results
    assert torch.equal(
        activations.get_masked_activations(), loaded_activations.get_masked_activations()
    )
    assert torch.equal(activations.get_masked_logits(), loaded_activations.get_masked_logits())


@pytest.fixture
def paired_dataset():
    honest_dialogues = [
        Dialogue(
            [
                Message("user", "What's the weather like?", False),
                Message("assistant", "It's sunny today.", False),
                Message("user", "Great, thanks!", False),
                Message("assistant", "You're welcome!", True),
            ]
        ),
        Dialogue(
            [
                Message("user", "How's the traffic?", False),
                Message("assistant", "The traffic is light.", True),
            ]
        ),
    ]
    deceptive_dialogues = [
        Dialogue(
            [
                Message("user", "What's the weather like?", False),
                Message("assistant", "It's raining heavily.", False),
                Message("user", "Oh, I see.", False),
                Message("assistant", "You're welcome!", True),
            ]
        ),
        Dialogue(
            [
                Message("user", "How's the traffic?", False),
                Message("assistant", "There's a major traffic jam.", True),
            ]
        ),
    ]
    honest_dataset = DialogueDataset(
        "plain", honest_dialogues, labels=[Label.HONEST] * len(honest_dialogues)
    )
    deceptive_dataset = DialogueDataset(
        "plain", deceptive_dialogues, labels=[Label.DECEPTIVE] * len(deceptive_dialogues)
    )
    dataset = honest_dataset + deceptive_dataset
    dataset.padding = {"gemma": {"left": 4, "right": 0}}
    return dataset


@pytest.mark.gpu
def test_split_paired_activations(
    gemma_model: PreTrainedModel,
    gemma_tokenizer: PreTrainedTokenizerBase,
    paired_dataset: DialogueDataset,
):
    # Test with default settings
    split_acts = PairedActivations.from_dataset(
        paired_dataset, gemma_tokenizer, gemma_model, layers=[15]
    )

    assert len(split_acts.honest_acts) == len(paired_dataset.get_honest())
    assert len(split_acts.deceptive_acts) == len(paired_dataset.get_deceptive())

    paired_dataset.padding = {"gemma": {"left": 2, "right": 0}}

    # Test detect_k_prior_tokens
    split_acts_k2 = PairedActivations.from_dataset(paired_dataset, gemma_tokenizer, gemma_model)

    assert split_acts_k2.honest_acts.tokenized_dataset.detection_mask is not None
    assert split_acts.honest_acts.tokenized_dataset.detection_mask is not None

    assert (
        split_acts_k2.honest_acts.tokenized_dataset.detection_mask.sum()
        < split_acts.honest_acts.tokenized_dataset.detection_mask.sum()
    )


@pytest.mark.gpu
def test_split_paired_detect_only_start_of_turn(
    gemma_model: PreTrainedModel,
    gemma_tokenizer: PreTrainedTokenizerBase,
    paired_dataset: DialogueDataset,
):
    # Test detect_only_start_of_turn
    paired_dataset.padding = {"gemma": {"left": 4, "right": 4}}  # Reset to default
    split_acts_start = PairedActivations.from_dataset(
        paired_dataset, gemma_tokenizer, gemma_model, detect_only_start_of_turn=True
    )

    # there's one start of turn token per dialogue on this dataset
    num_honest = len(paired_dataset.get_honest())
    assert split_acts_start.honest_acts.tokenized_dataset.detection_mask is not None
    assert split_acts_start.honest_acts.tokenized_dataset.detection_mask.sum() == num_honest

    # Check the tokens are actually start of turn tokens
    start_of_turn_id = cast(
        int, gemma_tokenizer("<start_of_turn>", add_special_tokens=False).input_ids[0]
    )
    masked_toks = split_acts_start.honest_acts.tokenized_dataset.tokens[
        split_acts_start.honest_acts.tokenized_dataset.detection_mask
    ]
    assert (masked_toks == start_of_turn_id).all()


@pytest.mark.gpu
def test_split_paired_layers(
    gemma_model: PreTrainedModel,
    gemma_tokenizer: PreTrainedTokenizerBase,
    paired_dataset: DialogueDataset,
):
    layers = [15]
    split_acts = PairedActivations.from_dataset(
        paired_dataset, gemma_tokenizer, gemma_model, layers=layers
    )
    honest_acts, deceptive_acts = split_acts.get_all_acts()
    assert honest_acts.layers == deceptive_acts.layers == layers


@pytest.mark.parametrize("sparse", [True, False])
def test_save_load(tokenized_dataset: TokenizedDataset, tmp_path: Path, sparse: bool):
    batch_size, seqpos = tokenized_dataset.tokens.shape
    layers = [0]
    d_emb = 256
    all_acts = torch.rand(batch_size, seqpos, len(layers), d_emb)
    all_acts = all_acts.to_sparse_coo() if sparse else all_acts

    acts = Activations(all_acts, tokenized_dataset, layers=layers)
    acts_path = tmp_path / "acts.pkl"
    acts.save(acts_path)
    loaded_acts = Activations.load(acts_path)

    assert loaded_acts.all_acts.is_sparse == sparse
    assert torch.allclose(loaded_acts.all_acts.to_dense(), acts.all_acts.to_dense())
    assert loaded_acts.layers == acts.layers
