#!/usr/bin/env python3
"""Test functional Qwen implementation against HuggingFace model.

Run locally (requires GPU):
    python -m tools.functional_extractor.test_qwen

Run on remote GPU:
    python -m tools.functional_extractor.test_qwen --remote
    python -m tools.functional_extractor.test_qwen --remote --gpu-id <id>
"""

from __future__ import annotations

import argparse
import sys


def test_on_gpu() -> None:
    """Run full verification test on GPU."""
    import os
    import sys

    # Add parent dirs to path for imports to work both locally and remotely
    # Script is at: <repo>/tools/functional_extractor/test_qwen.py
    # We need <repo> in sys.path so "from tools.functional_extractor import ..." works
    script_dir = os.path.dirname(os.path.abspath(__file__))  # functional_extractor/
    tools_dir = os.path.dirname(script_dir)  # tools/
    _repo_dir = os.path.dirname(tools_dir)  # <repo> (rollouts/)

    # Add the functional_extractor dir itself to path for direct import
    if script_dir not in sys.path:
        sys.path.insert(0, script_dir)

    import torch

    # Import functional implementation (direct import from same directory)
    from qwen_functional import qwen_forward
    from transformers import AutoModelForCausalLM

    print("=" * 60)
    print("Qwen2.5-0.5B Functional Implementation Test")
    print("=" * 60)

    print("\nLoading model...")
    model = AutoModelForCausalLM.from_pretrained(
        "Qwen/Qwen2.5-0.5B",
        torch_dtype=torch.bfloat16,
        device_map="cuda:0",
    )
    model.eval()

    weights = {k: v for k, v in model.state_dict().items()}

    # ========== Test Suite ==========
    all_passed = True
    test_count = 0

    def run_test(name: str, input_ids: torch.Tensor) -> bool:
        """Run a single test and return True if passed."""
        nonlocal test_count
        test_count += 1

        with torch.no_grad():
            original_logits = model(input_ids).logits
            functional_logits = qwen_forward(input_ids, weights)

        matches = torch.allclose(original_logits, functional_logits, rtol=1e-5, atol=1e-5)
        max_diff = (original_logits - functional_logits).abs().max().item()

        batch, seq_len, vocab = original_logits.shape
        status = "PASS" if matches else "FAIL"
        print(f"  {name}: batch={batch}, seq={seq_len}, max_diff={max_diff:.2e} [{status}]")

        if not matches:
            diff = (original_logits - functional_logits).abs()
            max_idx = diff.argmax()
            print(f"    Max diff at flat index: {max_idx.item()}")

        return matches

    # --- Sequence Length Tests ---
    print("\n### Sequence Length Tests ###")
    seq_lengths = [1, 4, 16, 32, 64, 128]
    for seq_len in seq_lengths:
        input_ids = torch.randint(1, 1000, (1, seq_len), device="cuda:0")
        if not run_test(f"seq_len={seq_len:3d}", input_ids):
            all_passed = False

    # --- Batch Size Tests ---
    print("\n### Batch Size Tests ###")
    batch_sizes = [1, 2, 4, 8]
    for batch_size in batch_sizes:
        input_ids = torch.randint(1, 1000, (batch_size, 16), device="cuda:0")
        if not run_test(f"batch={batch_size}, seq=16", input_ids):
            all_passed = False

    # --- Edge Cases ---
    print("\n### Edge Cases ###")

    # Single token
    input_ids = torch.tensor([[42]], device="cuda:0")
    if not run_test("single_token", input_ids):
        all_passed = False

    # Repeated tokens
    input_ids = torch.full((1, 32), 100, device="cuda:0")
    if not run_test("repeated_token", input_ids):
        all_passed = False

    # Sequential tokens
    input_ids = torch.arange(1, 65, device="cuda:0").unsqueeze(0)
    if not run_test("sequential_tokens", input_ids):
        all_passed = False

    # High token IDs (near vocab boundary)
    input_ids = torch.randint(150000, 151000, (1, 16), device="cuda:0")
    if not run_test("high_token_ids", input_ids):
        all_passed = False

    # --- Dtype Tests ---
    print("\n### Dtype Tests ###")

    def run_dtype_test(name: str, dtype: torch.dtype) -> bool:
        """Run test with specific dtype by reloading model."""
        nonlocal test_count
        test_count += 1

        # Load model with specified dtype
        model_typed = AutoModelForCausalLM.from_pretrained(
            "Qwen/Qwen2.5-0.5B",
            torch_dtype=dtype,
            device_map="cuda:0",
        )
        model_typed.eval()
        weights_typed = {k: v for k, v in model_typed.state_dict().items()}

        input_ids = torch.randint(1, 1000, (1, 16), device="cuda:0")

        with torch.no_grad():
            original_logits = model_typed(input_ids).logits
            functional_logits = qwen_forward(input_ids, weights_typed)

        # Use looser tolerance for fp16/fp32 which have different precision characteristics
        # fp32: RMSNorm computes in fp32 so there's some accumulation difference at 1e-5
        # fp16: lower precision overall
        if dtype == torch.float16:
            rtol, atol = 1e-3, 1e-3
        elif dtype == torch.float32:
            rtol, atol = 1e-4, 1e-4
        else:  # bf16
            rtol, atol = 1e-5, 1e-5
        matches = torch.allclose(original_logits, functional_logits, rtol=rtol, atol=atol)
        max_diff = (original_logits - functional_logits).abs().max().item()

        status = "PASS" if matches else "FAIL"
        print(f"  {name}: max_diff={max_diff:.2e} [{status}]")

        # Cleanup to free GPU memory
        del model_typed, weights_typed

        return matches

    for dtype_name, dtype in [
        ("bf16", torch.bfloat16),
        ("fp16", torch.float16),
        ("fp32", torch.float32),
    ]:
        if not run_dtype_test(dtype_name, dtype):
            all_passed = False

    # --- Attention Mask Tests ---
    # With sequential position IDs (matching HF), both implementations should match exactly.
    print("\n### Attention Mask Tests ###")

    def run_mask_test(name: str, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> bool:
        """Run test with attention mask."""
        nonlocal test_count
        test_count += 1

        with torch.no_grad():
            original_logits = model(input_ids, attention_mask=attention_mask).logits
            functional_logits = qwen_forward(input_ids, weights, attention_mask=attention_mask)

        batch_size, seq_len, vocab_size = original_logits.shape

        # Calculate diff only at non-padded positions
        mask_expanded = attention_mask.unsqueeze(-1).expand_as(original_logits)
        diff = (original_logits - functional_logits).abs()
        diff_masked = torch.where(mask_expanded == 1, diff, torch.zeros_like(diff))
        max_diff_real = diff_masked.max().item()

        # Use same tolerance as other tests
        matches = max_diff_real < 1e-4

        status = "PASS" if matches else "FAIL"
        print(
            f"  {name}: batch={batch_size}, seq={seq_len}, max_diff={max_diff_real:.2e} [{status}]"
        )

        if not matches:
            # Show where the largest diff is
            max_idx = diff_masked.argmax()
            print(f"    Max diff at flat index: {max_idx.item()}")

        return matches

    # All ones (no padding) - should be equivalent to no mask
    input_ids = torch.randint(1, 1000, (1, 16), device="cuda:0")
    attention_mask = torch.ones_like(input_ids)
    if not run_mask_test("all_ones", input_ids, attention_mask):
        all_passed = False

    # Left padding (common for batched inference)
    input_ids = torch.randint(1, 1000, (2, 16), device="cuda:0")
    attention_mask = torch.ones_like(input_ids)
    attention_mask[0, :4] = 0  # First 4 tokens are padding for batch item 0
    attention_mask[1, :2] = 0  # First 2 tokens are padding for batch item 1
    if not run_mask_test("left_padding", input_ids, attention_mask):
        all_passed = False

    # Variable length sequences
    input_ids = torch.randint(1, 1000, (4, 32), device="cuda:0")
    attention_mask = torch.ones_like(input_ids)
    attention_mask[0, :8] = 0  # 24 real tokens
    attention_mask[1, :16] = 0  # 16 real tokens
    attention_mask[2, :24] = 0  # 8 real tokens
    attention_mask[3, :4] = 0  # 28 real tokens
    if not run_mask_test("variable_length", input_ids, attention_mask):
        all_passed = False

    # --- Summary ---
    print("\n" + "=" * 60)
    if all_passed:
        print(f"ALL {test_count} TESTS PASSED!")
        print("Functional implementation is numerically identical to HF model.")
    else:
        print(f"SOME TESTS FAILED (out of {test_count})")
        sys.exit(1)


def main() -> None:
    parser = argparse.ArgumentParser(description="Test functional Qwen implementation")
    parser.add_argument("--remote", action="store_true", help="Run on remote GPU")
    parser.add_argument("--gpu-id", type=str, help="Reuse existing GPU instance")
    parser.add_argument("--keep-alive", action="store_true", help="Keep GPU alive after test")
    args = parser.parse_args()

    if args.remote:
        from tools.functional_extractor.config import DeploymentConfig
        from tools.functional_extractor.verify import run_on_gpu

        run_on_gpu(
            script_path=__file__,
            deployment=DeploymentConfig(vram_gb=16),
            gpu_id=args.gpu_id,
            keep_alive=args.keep_alive or bool(args.gpu_id),
        )
    else:
        # Check if we have GPU
        import torch

        if not torch.cuda.is_available():
            print("No GPU available locally. Use --remote to run on remote GPU.")
            sys.exit(1)

        test_on_gpu()


if __name__ == "__main__":
    main()
