"""Exercise: find and repair a wrong causal mask.

Run from this lab's directory:

    python exercises/wrong_mask.py
    python -m pytest tests/test_wrong_mask.py -q

The buggy mask is upper-triangular. Every tensor shape still matches the
reference implementation, so shape assertions pass. Only a behavioural check
(does any query attend to a future key?) catches the bug.
"""

from __future__ import annotations

import sys
from pathlib import Path

import torch

LAB_ROOT = Path(__file__).resolve().parents[1]
if str(LAB_ROOT) not in sys.path:
    sys.path.insert(0, str(LAB_ROOT))

from attention import scaled_dot_product_attention  # noqa: E402

T = 4
D_HEAD = 8


def buggy_causal_mask(seq_len: int) -> torch.Tensor:
    """Fixed: keeps the lower triangle, so query t sees keys 0..t."""
    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))


def fixed_causal_mask(seq_len: int) -> torch.Tensor:
    """RIGHT: keeps the lower triangle, so query t sees keys 0..t."""
    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))


def make_qkv(seq_len: int = T, d_head: int = D_HEAD, seed: int = 0):
    generator = torch.Generator().manual_seed(seed)
    q = torch.randn(1, 1, seq_len, d_head, generator=generator)
    k = torch.randn(1, 1, seq_len, d_head, generator=generator)
    v = torch.randn(1, 1, seq_len, d_head, generator=generator)
    return q, k, v


def future_leak_violations(
    attn_weights: torch.Tensor, tol: float = 1e-6
) -> torch.Tensor:
    """Return the weights placed on keys strictly after the query position.

    ``attn_weights`` is ``[..., T, T]``. An empty result means no leakage.
    """
    if attn_weights.shape[-1] != attn_weights.shape[-2]:
        raise ValueError(
            f"expected square attention weights, got {tuple(attn_weights.shape)}"
        )
    seq_len = attn_weights.shape[-1]
    future = torch.triu(
        torch.ones(seq_len, seq_len, dtype=torch.bool), diagonal=1
    ).expand_as(attn_weights)
    future_weights = attn_weights.masked_select(future)
    return future_weights[future_weights > tol]


def assert_no_future_leak(
    attn_weights: torch.Tensor, tol: float = 1e-6, label: str = ""
) -> None:
    violations = future_leak_violations(attn_weights, tol=tol)
    if violations.numel() > 0:
        raise AssertionError(
            f"{label}: {violations.numel()} future-token weights above {tol:g}; "
            f"largest={violations.max().item():.6f}"
        )


def weights_for(mask: torch.Tensor) -> torch.Tensor:
    q, k, v = make_qkv()
    _, weights = scaled_dot_product_attention(
        q, k, v, mask=mask, return_weights=True
    )
    return weights


def main() -> None:
    buggy = buggy_causal_mask(T)
    fixed = fixed_causal_mask(T)

    # Shape checks are identical for the wrong and right mask. This is why a
    # shape-only test suite cannot catch the bug.
    assert buggy.shape == fixed.shape == (T, T)
    print(f"shape check: both masks are {tuple(buggy.shape)} (passes for both)")

    buggy_weights = weights_for(buggy)
    fixed_weights = weights_for(fixed)

    buggy_leaks = future_leak_violations(buggy_weights)
    fixed_leaks = future_leak_violations(fixed_weights)

    if buggy_leaks.numel() > 0:
        print(f"buggy mask: {buggy_leaks.numel()} future weights, "
              f"largest={buggy_leaks.max().item():.6f}")
    else:
        print("buggy mask: 0 future weights (mask is lower-triangular, no leak)")

    if fixed_leaks.numel() > 0:
        print(f"fixed mask: {fixed_leaks.numel()} future weights")
    else:
        print("fixed mask: 0 future weights")

    assert_no_future_leak(fixed_weights, label="fixed mask")
    print("fixed mask passes the future-token isolation check")


if __name__ == "__main__":
    main()