"""
CW2: Gradient Checking for Loss Functions.

Run this script to verify your loss function gradients are correct:
    cd cw2_logistic_regression
    python -m tests.test_cw2

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

This uses numerical gradient checking (finite differences) to compare
your analytic gradients against numerical approximations.
If the relative error is < 1e-5, your gradient is likely correct.
"""

import sys
import os
import numpy as np

# Add paths
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "cw2_logistic_regression"))


def numerical_gradient(loss_fn, logits, labels, num_classes, eps=1e-5):
    """
    Compute numerical gradient of loss w.r.t. logits using finite differences.

    Args:
        loss_fn: Loss function that returns (loss, grad)
        logits: (batch_size, num_classes) input logits
        labels: (batch_size,) integer labels
        num_classes: Number of classes
        eps: Perturbation size

    Returns:
        num_grad: (batch_size, num_classes) numerical gradient
    """
    num_grad = np.zeros_like(logits)

    for i in range(logits.shape[0]):
        for j in range(logits.shape[1]):
            # f(x + eps)
            logits_plus = logits.copy()
            logits_plus[i, j] += eps
            loss_plus, _ = loss_fn(logits_plus, labels, num_classes)

            # f(x - eps)
            logits_minus = logits.copy()
            logits_minus[i, j] -= eps
            loss_minus, _ = loss_fn(logits_minus, labels, num_classes)

            # Central difference
            num_grad[i, j] = (loss_plus - loss_minus) / (2 * eps)

    return num_grad


def relative_error(a, b):
    """Compute relative error between two arrays."""
    return np.max(np.abs(a - b) / (np.maximum(np.abs(a) + np.abs(b), 1e-8)))


def test_loss_gradient(loss_fn, name, batch_size=4, num_classes=5):
    """
    Test a single loss function's gradient against numerical gradient.

    Args:
        loss_fn: Loss function that returns (loss, grad)
        name: Name of the loss function (for printing)
        batch_size: Number of samples in test batch
        num_classes: Number of classes
    """
    np.random.seed(42)

    # Create random logits and labels
    logits = np.random.randn(batch_size, num_classes).astype(np.float64)
    labels = np.random.randint(0, num_classes, size=batch_size)

    # Get analytic gradient
    try:
        loss, analytic_grad = loss_fn(logits, labels, num_classes)
    except NotImplementedError:
        print(f"  [{name}] NOT IMPLEMENTED — skipping")
        return False

    # Compute numerical gradient
    num_grad = numerical_gradient(loss_fn, logits, labels, num_classes)

    # Compare
    rel_err = relative_error(analytic_grad, num_grad)
    passed = rel_err < 1e-5

    status = "PASSED ✓" if passed else "FAILED ✗"
    print(f"  [{name}] {status}  (relative error: {rel_err:.2e})")

    if not passed:
        print(f"    Analytic gradient (sample):\n      {analytic_grad[0, :3]}")
        print(f"    Numerical gradient (sample):\n      {num_grad[0, :3]}")

    return passed


def test_softmax():
    """Test softmax implementation."""
    from losses import softmax

    np.random.seed(42)
    logits = np.random.randn(4, 5).astype(np.float64)

    try:
        probs = softmax(logits)
    except NotImplementedError:
        print("  [Softmax] NOT IMPLEMENTED — skipping")
        return False

    # Check: all positive
    assert np.all(probs > 0), "Softmax outputs should be positive"

    # Check: rows sum to 1
    row_sums = probs.sum(axis=1)
    assert np.allclose(row_sums, 1.0, atol=1e-6), \
        f"Softmax rows should sum to 1, got {row_sums}"

    # Check: numerical stability (large logits)
    large_logits = np.array([[1000, 1001, 1002]], dtype=np.float64)
    probs_large = softmax(large_logits)
    assert not np.any(np.isnan(probs_large)), "Softmax should handle large logits"
    assert not np.any(np.isinf(probs_large)), "Softmax should handle large logits"

    print("  [Softmax] PASSED ✓")
    return True


def main():
    from losses import cross_entropy_loss, hinge_loss, exponential_loss, squared_loss

    print("\n" + "=" * 50)
    print("  CW2: Gradient Checking")
    print("=" * 50)

    results = []

    # Test softmax first
    print("\n--- Softmax Test ---")
    results.append(test_softmax())

    # Test each loss function
    print("\n--- Loss Function Gradient Tests ---")

    loss_functions = [
        (cross_entropy_loss, "Cross-Entropy"),
        (hinge_loss, "Hinge"),
        (exponential_loss, "Exponential"),
        (squared_loss, "Squared"),
    ]

    for loss_fn, name in loss_functions:
        results.append(test_loss_gradient(loss_fn, name))

    # Summary
    print("\n" + "=" * 50)
    passed = sum(r for r in results if r)
    total = len(results)
    print(f"  Results: {passed}/{total} tests passed")
    if passed == total:
        print("  All tests passed! Your gradients are correct.")
    else:
        print("  Some tests failed. Check your implementations.")
    print("=" * 50)


if __name__ == "__main__":
    main()
