"""
CW4: Shape Checking for CNN Models.

Run this script to verify your CNN architecture is correct:
    cd cw4_cnn
    python -m tests.test_cw4

Or from the project root:
    python -m tests.test_cw4

Checks:
    - Model accepts (B, 1, 28, 28) input
    - Model produces (B, 10) output
    - Gradients flow through all layers (no dead layers)
"""

import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "cw4_cnn"))

import torch
import torch.nn as nn


def test_model(model_class, model_name, num_classes=10):
    """
    Test a CNN model for correct input/output shapes and gradient flow.

    Args:
        model_class: nn.Module subclass
        model_name: Name for display
        num_classes: Expected number of output classes
    """
    print(f"\n--- Testing {model_name} ---")

    # Create model
    try:
        model = model_class(num_classes=num_classes)
    except NotImplementedError:
        print(f"  [{model_name}] NOT IMPLEMENTED — skipping")
        return False

    # Test forward pass shape
    batch_size = 4
    x = torch.randn(batch_size, 1, 28, 28)

    try:
        output = model(x)
    except Exception as e:
        print(f"  [{model_name} forward] FAILED ✗ — {e}")
        return False

    if output.shape != (batch_size, num_classes):
        print(f"  [{model_name} forward] FAILED ✗ — "
              f"Expected output shape ({batch_size}, {num_classes}), "
              f"got {tuple(output.shape)}")
        return False

    print(f"  [{model_name} forward] PASSED ✓  (output shape: {tuple(output.shape)})")

    # Test gradient flow
    loss = output.sum()
    loss.backward()

    has_grad = True
    for name, param in model.named_parameters():
        if param.grad is None:
            print(f"  [{model_name} grad] WARNING: {name} has no gradient!")
            has_grad = False
        elif torch.all(param.grad == 0):
            print(f"  [{model_name} grad] WARNING: {name} has zero gradient!")
            has_grad = False

    if has_grad:
        print(f"  [{model_name} gradient flow] PASSED ✓")

    # Count parameters
    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"  [{model_name} params] {num_params:,} trainable parameters")

    # Test with different batch sizes
    for bs in [1, 16]:
        x = torch.randn(bs, 1, 28, 28)
        out = model(x)
        assert out.shape == (bs, num_classes), \
            f"Failed for batch_size={bs}: got shape {out.shape}"
    print(f"  [{model_name} batch sizes] PASSED ✓  (tested bs=1, 16)")

    return True


def main():
    from model import SimpleCNN, DeepCNN

    print("\n" + "=" * 50)
    print("  CW4: CNN Shape Checking")
    print("=" * 50)

    results = []
    results.append(test_model(SimpleCNN, "SimpleCNN"))
    results.append(test_model(DeepCNN, "DeepCNN"))

    # Summary
    print("\n" + "=" * 50)
    passed = sum(r for r in results if r)
    total = len(results)
    print(f"  Results: {passed}/{total} models passed")
    if passed == total:
        print("  All models have correct shapes and gradient flow!")
    elif passed > 0:
        print("  SimpleCNN passed. DeepCNN is bonus — implement it for extra credit.")
    else:
        print("  Please check your CNN implementations.")
    print("=" * 50)


if __name__ == "__main__":
    main()
