import pytest
import torch
from transformers import AutoConfig, AutoModelForCausalLM

from ludwig.constants import LOGITS, PREDICTIONS, PROBABILITIES
from ludwig.modules.training_hooks import NEFTuneHook
from ludwig.utils.llm_utils import (
    add_left_padding,
    create_attention_mask,
    FALLBACK_CONTEXT_LEN,
    find_last_matching_index,
    generate_merged_ids,
    get_context_len,
    get_realigned_target_and_prediction_tensors_for_inference,
    has_padding_token,
    pad_target_tensor_for_fine_tuning,
    remove_left_padding,
)
from ludwig.utils.tokenizers import HFTokenizer

pytestmark = [pytest.mark.llm]

# Pad token ID is 1 for OPT even though it uses the GPT2 tokenizer
# BOS token ID is 2
TEST_MODEL_NAME = "hf-internal-testing/tiny-random-OPTForCausalLM"


@pytest.fixture
def tokenizer():
    return HFTokenizer(TEST_MODEL_NAME).tokenizer


@pytest.fixture
def input_ids():
    # Provide sample input IDs tensor
    return torch.tensor([[3, 4, 5], [6, 7, 8]])


@pytest.fixture
def target_ids():
    # Provide sample target IDs tensor
    return torch.tensor([[9, 10, 11], [12, 13, 14]])


class TestSetContextLen:
    def test_max_sequence_length(self):
        # Test when 'max_sequence_length' is present in the model configuration
        config = AutoConfig.from_pretrained("huggyllama/llama-7b")
        assert get_context_len(config) == config.max_sequence_length

    def test_max_position_embeddings(self):
        # Test when 'max_position_embeddings' is present in the model configuration
        config = AutoConfig.from_pretrained("huggyllama/llama-7b")
        del config.max_sequence_length
        assert get_context_len(config) == config.max_position_embeddings

    def test_n_positions(self):
        # Test when 'n_positions' is present in the model configuration
        config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-GPTJForCausalLM")
        assert get_context_len(config) == config.n_positions

    def test_default_value(self):
        config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-GPTJForCausalLM")
        del config.n_positions
        assert get_context_len(config) == FALLBACK_CONTEXT_LEN


def test_has_padding_token_with_padding_tokens(tokenizer):
    input_sentence = "This is an example sentence."
    input_ids = tokenizer([input_sentence])
    input_ids["input_ids"] = torch.tensor(input_ids["input_ids"])
    padded_input_ids = torch.nn.functional.pad(input_ids["input_ids"], (10 - len(input_ids["input_ids"]), 1), value=1)

    assert has_padding_token(padded_input_ids, tokenizer)


def test_has_padding_token_without_padding_tokens(tokenizer):
    input_sentence = "This is an example sentence."
    input_ids = tokenizer([input_sentence])
    input_ids["input_ids"] = torch.tensor(input_ids["input_ids"])

    assert not has_padding_token(input_ids["input_ids"], tokenizer)


@pytest.mark.parametrize(
    "input_ids, expected",
    [
        # No padding
        (torch.tensor([5]), torch.tensor([5])),
        (torch.tensor([5, 3]), torch.tensor([5, 3])),
        # Padding
        (torch.tensor([1, 5, 5, 3]), torch.tensor([5, 5, 3])),
        # EOS token
        (torch.tensor([2, 5, 5, 3]), torch.tensor([2, 5, 5, 3])),
        # Padding + EOS token
        (torch.tensor([1, 2, 5, 5, 3]), torch.tensor([2, 5, 5, 3])),
    ],
)
def test_remove_left_padding(input_ids, expected, tokenizer):
    assert torch.equal(remove_left_padding(input_ids, tokenizer).squeeze(0), expected)


@pytest.mark.parametrize(
    "input_ids, max_length, pad_value, expected",
    [
        (torch.tensor([1, 2, 3]), 3, 0, torch.tensor([1, 2, 3])),
        (torch.tensor([1, 2, 3]), 5, 0, torch.tensor([0, 0, 1, 2, 3])),
        (torch.tensor([4, 5, 6, 7]), 6, 2, torch.tensor([2, 2, 4, 5, 6, 7])),
        (torch.tensor([8, 9]), 3, 1, torch.tensor([1, 8, 9])),
    ],
)
def test_add_left_padding(input_ids, max_length, pad_value, expected):
    padded = add_left_padding(input_ids, max_length, pad_value).squeeze(0)

    assert torch.equal(padded, expected)


def test_create_attention_mask_last_token_padding(tokenizer):
    input_ids = torch.tensor([3, 4, tokenizer.pad_token_id])
    attention_mask = create_attention_mask(input_ids, tokenizer)
    assert attention_mask[-1] == 1


@pytest.mark.parametrize(
    "input_ids, expected_output",
    [
        # No padding
        (torch.tensor([3, 4, 5]), torch.tensor([1, 1, 1])),
        (torch.tensor([1, 1, 4, 6, 8]), torch.tensor([0, 0, 1, 1, 1])),
        # All padding
        (torch.tensor([1, 1, 1]), torch.tensor([0, 0, 1])),
    ],
)
def test_create_attention_mask(input_ids, expected_output, tokenizer):
    attention_mask = create_attention_mask(input_ids, tokenizer)

    assert torch.equal(attention_mask, expected_output)


@pytest.mark.parametrize(
    "tensor_a, tensor_b, expected_index",
    [
        # Matching index at the end
        (torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]), torch.tensor([6, 7, 8]), 5),
        # No matching index
        (torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]), torch.tensor([9, 10]), -1),
        # Matching index in the middle. Fails because we're only checking the last X elements.
        (torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]), torch.tensor([4, 5, 6]), -1),
    ],
)
def test_find_last_matching_index(tensor_a, tensor_b, expected_index):
    last_matching_index = find_last_matching_index(tensor_a, tensor_b)
    assert last_matching_index == expected_index


def test_generate_merged_ids_with_target(tokenizer, input_ids, target_ids):
    # Test case when target_ids is not None
    merged_ids, attention_masks = generate_merged_ids(input_ids, target_ids, tokenizer)
    assert torch.equal(merged_ids, torch.tensor([[3, 4, 5, 9, 10, 11, 2], [6, 7, 8, 12, 13, 14, 2]]))
    assert merged_ids.shape == (2, 7)  # Check the shape of merged_ids
    assert attention_masks.shape == (2, 7)  # Check the shape of attention_masks


def test_generate_merged_ids_with_max_sequence_length(tokenizer, input_ids, target_ids):
    # Test case when max_sequence_length is provided
    max_sequence_length = 5
    merged_ids, attention_masks = generate_merged_ids(input_ids, target_ids, tokenizer, max_sequence_length)
    assert merged_ids.shape == (2, 5)  # Check the shape of merged_ids with truncation
    assert attention_masks.shape == (2, 5)  # Check the shape of attention_masks


def test_generate_merged_ids_padding_removal(tokenizer, input_ids, target_ids):
    # Test case to check removal of left padding from inputs and targets during merge
    padding_tokens = torch.tensor([tokenizer.pad_token_id, tokenizer.pad_token_id])

    # Adds 2 padding tokens to the left of input_ids and target_ids individually. Typically, if we just merged this
    # naively, we would expect to see [1, 1, 3, 4, 5, 1, 1, 9, 10, 11, 1], but we shouldn't see the padding tokens
    # except for the padding token at the end.
    input_ids_with_padding = torch.cat((padding_tokens.unsqueeze(0).expand(input_ids.size(0), -1), input_ids), dim=1)
    target_ids_with_padding = torch.cat((padding_tokens.unsqueeze(0).expand(target_ids.size(0), -1), target_ids), dim=1)

    merged_ids, attention_masks = generate_merged_ids(input_ids_with_padding, target_ids_with_padding, tokenizer)

    assert merged_ids.shape == (2, 7)  # Check the shape of merged_ids
    assert attention_masks.shape == (2, 7)  # Check the shape of attention_masks

    assert torch.equal(merged_ids[0][:3], input_ids[0])  # Check the input_ids part without padding
    assert torch.equal(merged_ids[0][3:-1], target_ids[0])  # Check the target_ids part without padding
    assert torch.equal(merged_ids[0][-1], torch.tensor(tokenizer.eos_token_id))  # Check the padding tokens

    assert torch.all(attention_masks == 1)


def test_generate_merged_ids_returns_tensor(tokenizer, input_ids, target_ids):
    # Test that the function returns torch.Tensor objects
    merged_ids, attention_masks = generate_merged_ids(input_ids, target_ids, tokenizer)
    assert isinstance(merged_ids, torch.Tensor)
    assert isinstance(attention_masks, torch.Tensor)


def test_pad_target_tensor_for_fine_tuning():
    of_name = "out_1"
    prediction = {
        of_name: {PREDICTIONS: torch.tensor([[764, 764, 764, 764, 764, 764, 764, 578, 619, 841, 182, 905, 483, 764]])}
    }

    # Scenario 1: Entire target tensor was passed into model inputs
    model_input = torch.tensor([[0, 0, 24, 52, 654, 529, 221, 78, 79, 504, 76, 397, 84, 0]])
    target = {of_name: torch.tensor([[78, 79, 504, 76, 397, 84, 0]])}
    expected_target = {of_name: torch.tensor([[-100, -100, -100, -100, -100, -100, -100, 78, 79, 504, 76, 397, 84, 0]])}
    updated_targets = pad_target_tensor_for_fine_tuning(target, prediction, model_input, of_name)
    assert torch.equal(expected_target[of_name], updated_targets[of_name])

    # Scenario 2: Entire target tensor was not passed into model inputs
    model_input = torch.tensor([[13, 24, 395, 13, 46, 57, 52, 41, 45, 37, 51, 14, 380, 435]])
    target = {of_name: torch.tensor([[78, 79, 504, 76, 397, 84, 0]])}
    expected_target = {
        of_name: torch.tensor([[-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100]])
    }
    updated_targets = pad_target_tensor_for_fine_tuning(target, prediction, model_input, of_name)
    assert torch.equal(expected_target[of_name], updated_targets[of_name])

    # Scenario 3: Partial target tensor was passed into model inputs
    model_input = torch.tensor([[0, 0, 24, 52, 654, 529, 221, 78, 79, 504, 76, 78, 79, 504]])
    target = {of_name: torch.tensor([[78, 79, 504, 76, 397, 84, 0]])}
    expected_target = {
        of_name: torch.tensor([[-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, 78, 79, 504]])
    }
    updated_targets = pad_target_tensor_for_fine_tuning(target, prediction, model_input, of_name)
    assert torch.equal(expected_target[of_name], updated_targets[of_name])


def test_get_realigned_target_and_prediction_tensors_for_inference(tokenizer):
    of_name = "out_1"
    vocab_size = 8

    # Scenario 1: Prediction and target tensors have the same length, so nothing should change
    targets = {of_name: torch.tensor([[78, 79, 504, 76, 397, 84, 0]])}
    predictions = {
        of_name: {
            PREDICTIONS: torch.tensor([[78, 79, 504, 76, 397, 84, 0]], dtype=torch.int64),
            PROBABILITIES: torch.randn(1, 7, vocab_size).to(torch.float32),
            LOGITS: torch.randn(1, 7, vocab_size).to(torch.float32),
        }
    }
    updated_targets, updated_predictions = get_realigned_target_and_prediction_tensors_for_inference(
        targets, predictions, of_name, tokenizer
    )

    assert targets == updated_targets
    assert predictions == updated_predictions
    assert predictions[of_name][PREDICTIONS].shape[1] == targets[of_name].shape[1]
    assert predictions[of_name][PROBABILITIES].shape[1] == targets[of_name].shape[1]
    assert predictions[of_name][LOGITS].shape[1] == targets[of_name].shape[1]

    # Scenario 2: Prediction length is longer than the target tensor, so we need to realign the target tensor
    targets = {of_name: torch.tensor([[78, 79, 504, 76, 397, 84, 0]])}
    predictions = {
        of_name: {
            PREDICTIONS: torch.tensor([[98, 47, 78, 79, 504, 76, 397, 84, 0]], dtype=torch.int64),
            PROBABILITIES: torch.randn(1, 9, vocab_size).to(torch.float32),
            LOGITS: torch.randn(1, 9, vocab_size).to(torch.float32),
        }
    }
    updated_targets, updated_predictions = get_realigned_target_and_prediction_tensors_for_inference(
        targets, predictions, of_name, tokenizer
    )

    for key in updated_predictions.keys():
        assert torch.equal(updated_predictions[key][PREDICTIONS], predictions[key][PREDICTIONS])
        assert torch.equal(updated_predictions[key][PROBABILITIES], predictions[key][PROBABILITIES])
        assert torch.equal(updated_predictions[key][LOGITS], predictions[key][LOGITS])

    assert torch.equal(updated_targets[of_name], torch.tensor([[78, 79, 504, 76, 397, 84, 0, 1, 1]]))

    # Scenario 3: Target length is longer than the prediction tensor, so we need to realign them
    targets = {of_name: torch.tensor([[98, 47, 78, 79, 504, 76, 397, 84, 0]])}
    predictions = {
        of_name: {
            PREDICTIONS: torch.tensor([[78, 79, 504, 76, 397, 84, 0]], dtype=torch.int64),
            PROBABILITIES: torch.randn(1, 7, vocab_size).to(torch.float32),
            LOGITS: torch.randn(1, 7, vocab_size).to(torch.float32),
        }
    }
    updated_targets, updated_predictions = get_realigned_target_and_prediction_tensors_for_inference(
        targets, predictions, of_name, tokenizer
    )

    assert torch.equal(updated_targets[of_name], targets[of_name])

    assert torch.equal(updated_predictions[of_name][PREDICTIONS], torch.tensor([[78, 79, 504, 76, 397, 84, 0, 1, 1]]))
    assert updated_predictions[of_name][PROBABILITIES].shape[1] == targets[of_name].shape[1]
    assert updated_predictions[of_name][LOGITS].shape[1] == targets[of_name].shape[1]

    assert torch.equal(updated_predictions[of_name][PROBABILITIES][0][-1], torch.zeros(vocab_size))
    assert torch.equal(updated_predictions[of_name][PROBABILITIES][0][-2], torch.zeros(vocab_size))
    assert not torch.equal(updated_predictions[of_name][PROBABILITIES][0][-3], torch.zeros(vocab_size))

    assert torch.equal(updated_predictions[of_name][LOGITS][0][-1], torch.zeros(vocab_size))
    assert torch.equal(updated_predictions[of_name][LOGITS][0][-2], torch.zeros(vocab_size))
    assert not torch.equal(updated_predictions[of_name][LOGITS][0][-3], torch.zeros(vocab_size))


def _setup_models_for_neftune():
    module_without_hook = AutoModelForCausalLM.from_pretrained(TEST_MODEL_NAME)
    module_with_hook = AutoModelForCausalLM.from_pretrained(TEST_MODEL_NAME)

    # Only module_with_hook should have the NEFTuneHook
    neftune_hook = NEFTuneHook(neftune_noise_alpha=5)
    module_with_hook = neftune_hook.activate_hook(module_with_hook)

    return module_without_hook, module_with_hook


def _forward_pass_and_assert_neftune_hook(module_without_hook, module_with_hook, mode):
    assert module_with_hook.get_input_embeddings()._forward_hooks
    assert not module_without_hook.get_input_embeddings()._forward_hooks

    if mode == "train":
        module_without_hook.train()
        module_with_hook.train()
    elif mode == "eval":
        module_without_hook.eval()
        module_with_hook.eval()

    input_tensor = torch.tensor([[1, 2, 3]])
    output_tensor_with_noise = module_with_hook.get_input_embeddings()(input_tensor)
    output_tensor_without_noise = module_without_hook.get_input_embeddings()(input_tensor)

    if mode == "train":
        assert not torch.equal(output_tensor_with_noise, output_tensor_without_noise)
    elif mode == "eval":
        assert torch.equal(output_tensor_with_noise, output_tensor_without_noise)


def test_neftune_hook_with_noise_alpha_train_mode():
    """Test that the NEFTuneHook is only applied when the module is in training mode."""
    module_without_hook, module_with_hook = _setup_models_for_neftune()
    _forward_pass_and_assert_neftune_hook(module_without_hook, module_with_hook, mode="train")


def test_neftune_hook_with_noise_alpha_eval_mode():
    """Test that the NEFTuneHook is not applied when the module is in eval mode."""
    module_without_hook, module_with_hook = _setup_models_for_neftune()
    _forward_pass_and_assert_neftune_hook(module_without_hook, module_with_hook, mode="eval")