#!/usr/bin/env python
"""
Quick model testing script to verify model creation and inference.
"""

import os
import sys
import time

import torch

# Add project root to the path to enable imports
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from models.detection import get_maskrcnn_model
from utils.data_utils import PennFudanDataset, get_transform


def test_model_creation():
    """Test that we can create the model."""
    print("Testing model creation...")
    model = get_maskrcnn_model(
        num_classes=2, pretrained=False, pretrained_backbone=False
    )
    print("✓ Model created successfully")
    return model


def test_model_forward(model, device):
    """Test model forward pass with random inputs."""
    print("\nTesting model forward pass...")

    # Create a random batch
    image = torch.rand(3, 300, 400, device=device)  # Random image

    # Create a random target
    target = {
        "boxes": torch.tensor(
            [[100, 100, 200, 200]], dtype=torch.float32, device=device
        ),
        "labels": torch.tensor([1], dtype=torch.int64, device=device),
        "masks": torch.randint(0, 2, (1, 300, 400), dtype=torch.uint8, device=device),
        "image_id": torch.tensor([0], device=device),
        "area": torch.tensor([10000.0], dtype=torch.float32, device=device),
        "iscrowd": torch.tensor([0], dtype=torch.uint8, device=device),
    }

    # Test inference mode (no targets)
    model.eval()
    with torch.no_grad():
        start_time = time.time()
        output_inference = model([image])
        inference_time = time.time() - start_time

    # Verify inference output
    print(f"✓ Inference mode output: {type(output_inference)}")
    print(f"✓ Inference time: {inference_time:.3f}s")
    print(f"✓ Detection boxes shape: {output_inference[0]['boxes'].shape}")
    print(f"✓ Detection scores shape: {output_inference[0]['scores'].shape}")

    # Test training mode (with targets)
    model.train()
    start_time = time.time()
    output_train = model([image], [target])
    train_time = time.time() - start_time

    # Verify training output
    print(f"✓ Training mode output: {type(output_train)}")
    print(f"✓ Training time: {train_time:.3f}s")

    # Print loss values
    for loss_name, loss_value in output_train.items():
        print(f"✓ {loss_name}: {loss_value.item():.4f}")

    return output_train


def test_model_backward(model, loss_dict, device):
    """Test model backward pass."""
    print("\nTesting model backward pass...")

    # Calculate total loss
    total_loss = sum(loss for loss in loss_dict.values())

    # Create optimizer
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

    # Backward pass
    start_time = time.time()
    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()
    backward_time = time.time() - start_time

    print("✓ Backward pass and optimization completed")
    print(f"✓ Backward time: {backward_time:.3f}s")

    # Check that gradients were calculated
    has_gradients = any(
        param.grad is not None for param in model.parameters() if param.requires_grad
    )
    print(f"✓ Model has gradients: {has_gradients}")


def test_dataset():
    """Test that we can load the dataset."""
    print("\nTesting dataset loading...")

    data_root = "data/PennFudanPed"
    if not os.path.exists(data_root):
        print("✗ Dataset not found at", data_root)
        return None

    # Create dataset
    dataset = PennFudanDataset(root=data_root, transforms=get_transform(train=True))
    print(f"✓ Dataset loaded with {len(dataset)} samples")

    # Test loading a sample
    start_time = time.time()
    img, target = dataset[0]
    load_time = time.time() - start_time

    print(f"✓ Sample loaded in {load_time:.3f}s")
    print(f"✓ Image shape: {img.shape}")
    print(f"✓ Target boxes shape: {target['boxes'].shape}")

    return dataset


def main():
    """Run all tests."""
    print("=== Quick Model Testing Script ===")

    # Set device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Run tests
    model = test_model_creation()
    model.to(device)

    loss_dict = test_model_forward(model, device)

    test_model_backward(model, loss_dict, device)

    test_dataset()

    print("\n=== All tests completed successfully ===")


if __name__ == "__main__":
    main()