import torch

from inteprelens.core.attribution import CausalCircuitAttribution
from inteprelens.core.chg import CHG


def test_gradient_flows_without_mask(shared_model, llama_tokenizer):
    text = "Hello, this is a test."
    encoding = llama_tokenizer(text, return_tensors="pt")
    input_ids = encoding["input_ids"].to("cuda")

    embed_layer = shared_model.get_input_embeddings()
    embeds = embed_layer(input_ids)
    embeds = embeds.clone().detach().requires_grad_(True)

    outputs = shared_model(inputs_embeds=embeds)
    loss = outputs.logits[0, -1, :].sum()
    loss.backward()

    assert embeds.grad is not None
    assert embeds.grad.abs().sum() > 0


def test_gradient_flows_with_all_ones_mask(shared_model, llama_tokenizer, llama_chg):
    text = "Hello, this is a test."
    encoding = llama_tokenizer(text, return_tensors="pt")
    input_ids = encoding["input_ids"].to("cuda")

    mask = torch.ones(1, llama_chg.num_layers, llama_chg.num_heads, device="cuda")
    embed_layer = shared_model.get_input_embeddings()
    embeds = embed_layer(input_ids)
    embeds = embeds.clone().detach().requires_grad_(True)

    with llama_chg.register_mask_hooks(mask):
        outputs = shared_model(inputs_embeds=embeds)

    loss = outputs.logits[0, -1, :].sum()
    loss.backward()

    assert embeds.grad is not None
    assert embeds.grad.abs().sum() > 0


def test_gradient_magnitude_changes_with_mask(shared_model, llama_tokenizer, llama_chg):
    text = "The quick brown fox jumps over the lazy dog."
    encoding = llama_tokenizer(text, return_tensors="pt")
    input_ids = encoding["input_ids"].to("cuda")

    embed_layer = shared_model.get_input_embeddings()

    embeds_full = embed_layer(input_ids).clone().detach().requires_grad_(True)
    mask_full = torch.ones(1, llama_chg.num_layers, llama_chg.num_heads, device="cuda")
    with llama_chg.register_mask_hooks(mask_full):
        outputs = shared_model(inputs_embeds=embeds_full)
        loss = outputs.logits[0, -1, :].sum()
    loss.backward()
    grad_full = embeds_full.grad.abs().sum().item()

    embeds_half = embed_layer(input_ids).clone().detach().requires_grad_(True)
    mask_half = torch.zeros(1, llama_chg.num_layers, llama_chg.num_heads, device="cuda")
    mask_half[:, :, ::2] = 1.0
    with llama_chg.register_mask_hooks(mask_half):
        outputs = shared_model(inputs_embeds=embeds_half)
        loss = outputs.logits[0, -1, :].sum()
    loss.backward()
    grad_half = embeds_half.grad.abs().sum().item()

    assert grad_full > 0
    assert grad_half > 0
    assert abs(grad_full - grad_half) / max(grad_full, grad_half, 1e-8) > 0.01


def test_compute_token_gradients(shared_model, llama_tokenizer):
    attribution = CausalCircuitAttribution(
        shared_model,
        llama_tokenizer,
        enable_gradient_checkpointing=False,
    )
    document = "The quick brown fox jumps over the lazy dog."
    summary = "Fox jumps dog."

    doc_enc = llama_tokenizer(document, return_tensors="pt")
    sum_enc = llama_tokenizer(summary, add_special_tokens=False, return_tensors="pt")

    input_ids = doc_enc["input_ids"].to("cuda")
    target_ids = sum_enc["input_ids"].to("cuda")

    chg = CHG(shared_model)
    mask = torch.ones(chg.num_layers, chg.num_heads)

    grad_norms = attribution.compute_token_gradients(input_ids, target_ids, mask)

    assert grad_norms.shape[0] == input_ids.shape[1]
    assert grad_norms.min() >= 0
    assert grad_norms.sum() > 0


def test_compute_sentence_importance(shared_model, llama_tokenizer):
    attribution = CausalCircuitAttribution(
        shared_model,
        llama_tokenizer,
        enable_gradient_checkpointing=False,
    )
    document = "The quick brown fox. It jumps over the lazy dog. The end."
    sentences = ["The quick brown fox.", "It jumps over the lazy dog.", "The end."]
    summary = "Fox jumps over dog."

    chg = CHG(shared_model)
    mask = torch.ones(chg.num_layers, chg.num_heads)

    scores = attribution.compute_sentence_importance(
        document=document,
        sentences=sentences,
        summary=summary,
        facilitating_mask=mask,
        use_full_target=False,
    )

    assert len(scores) == len(sentences)
    assert scores.min() >= 0
    assert scores.max() <= 1.0 + 1e-6
