"""
CW3: Gradient Checking for Neural Network Layers.

Run this script to verify your layer backward() implementations:
    cd cw3_mlp
    python -m tests.test_cw3

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

For each layer, we compare the analytic gradient (from backward()) against
a numerical gradient (from finite differences). If the relative error is
< 1e-5, your implementation is likely correct.
"""

import sys
import os
import numpy as np

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


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_linear_layer():
    """Test Linear layer forward and backward."""
    from layers import Linear

    np.random.seed(42)
    layer = Linear(5, 3)
    X = np.random.randn(4, 5).astype(np.float64)
    # Use float64 for gradient checking precision
    layer.W = layer.W.astype(np.float64)
    layer.b = layer.b.astype(np.float64)

    # Forward
    try:
        output = layer.forward(X)
    except NotImplementedError:
        print("  [Linear forward] NOT IMPLEMENTED — skipping")
        return False

    expected = X @ layer.W + layer.b
    assert np.allclose(output, expected, atol=1e-10), "Linear forward incorrect"
    print("  [Linear forward] PASSED ✓")

    # Backward: check grad_input
    try:
        grad_output = np.random.randn(4, 3).astype(np.float64)
        grad_input = layer.backward(grad_output)
    except NotImplementedError:
        print("  [Linear backward] NOT IMPLEMENTED — skipping")
        return False

    # Numerical gradient for grad_input
    eps = 1e-5
    num_grad_input = np.zeros_like(X)
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            X_plus = X.copy()
            X_plus[i, j] += eps
            out_plus = X_plus @ layer.W + layer.b

            X_minus = X.copy()
            X_minus[i, j] -= eps
            out_minus = X_minus @ layer.W + layer.b

            num_grad_input[i, j] = np.sum(grad_output * (out_plus - out_minus)) / (2 * eps)

    rel_err = relative_error(grad_input, num_grad_input)
    passed = rel_err < 1e-5
    status = "PASSED ✓" if passed else "FAILED ✗"
    print(f"  [Linear backward - grad_input] {status}  (rel error: {rel_err:.2e})")

    # Check grad_W
    num_grad_W = np.zeros_like(layer.W)
    for i in range(layer.W.shape[0]):
        for j in range(layer.W.shape[1]):
            W_orig = layer.W[i, j]

            layer.W[i, j] = W_orig + eps
            out_plus = X @ layer.W + layer.b

            layer.W[i, j] = W_orig - eps
            out_minus = X @ layer.W + layer.b

            num_grad_W[i, j] = np.sum(grad_output * (out_plus - out_minus)) / (2 * eps)
            layer.W[i, j] = W_orig

    rel_err_W = relative_error(layer.grad_W, num_grad_W)
    passed_W = rel_err_W < 1e-5
    status_W = "PASSED ✓" if passed_W else "FAILED ✗"
    print(f"  [Linear backward - grad_W]     {status_W}  (rel error: {rel_err_W:.2e})")

    # Check grad_b
    num_grad_b = np.zeros_like(layer.b)
    for j in range(layer.b.shape[0]):
        b_orig = layer.b[j]

        layer.b[j] = b_orig + eps
        out_plus = X @ layer.W + layer.b

        layer.b[j] = b_orig - eps
        out_minus = X @ layer.W + layer.b

        num_grad_b[j] = np.sum(grad_output * (out_plus - out_minus)) / (2 * eps)
        layer.b[j] = b_orig

    rel_err_b = relative_error(layer.grad_b, num_grad_b)
    passed_b = rel_err_b < 1e-5
    status_b = "PASSED ✓" if passed_b else "FAILED ✗"
    print(f"  [Linear backward - grad_b]     {status_b}  (rel error: {rel_err_b:.2e})")

    return passed and passed_W and passed_b


def test_relu_layer():
    """Test ReLU layer forward and backward."""
    from layers import ReLU

    np.random.seed(42)
    layer = ReLU()
    X = np.random.randn(4, 5).astype(np.float64)

    try:
        output = layer.forward(X)
    except NotImplementedError:
        print("  [ReLU forward] NOT IMPLEMENTED — skipping")
        return False

    expected = np.maximum(0, X)
    assert np.allclose(output, expected), "ReLU forward incorrect"
    print("  [ReLU forward] PASSED ✓")

    try:
        grad_output = np.random.randn(4, 5).astype(np.float64)
        grad_input = layer.backward(grad_output)
    except NotImplementedError:
        print("  [ReLU backward] NOT IMPLEMENTED — skipping")
        return False

    expected_grad = grad_output * (X > 0)
    rel_err = relative_error(grad_input, expected_grad)
    passed = rel_err < 1e-5
    status = "PASSED ✓" if passed else "FAILED ✗"
    print(f"  [ReLU backward] {status}  (rel error: {rel_err:.2e})")
    return passed


def test_sigmoid_layer():
    """Test Sigmoid layer forward and backward."""
    from layers import Sigmoid

    np.random.seed(42)
    layer = Sigmoid()
    X = np.random.randn(4, 5).astype(np.float64)

    try:
        output = layer.forward(X)
    except NotImplementedError:
        print("  [Sigmoid forward] NOT IMPLEMENTED — skipping")
        return False

    expected = 1.0 / (1.0 + np.exp(-X))
    assert np.allclose(output, expected, atol=1e-10), "Sigmoid forward incorrect"
    print("  [Sigmoid forward] PASSED ✓")

    try:
        grad_output = np.random.randn(4, 5).astype(np.float64)
        grad_input = layer.backward(grad_output)
    except NotImplementedError:
        print("  [Sigmoid backward] NOT IMPLEMENTED — skipping")
        return False

    expected_grad = grad_output * output * (1 - output)
    rel_err = relative_error(grad_input, expected_grad)
    passed = rel_err < 1e-5
    status = "PASSED ✓" if passed else "FAILED ✗"
    print(f"  [Sigmoid backward] {status}  (rel error: {rel_err:.2e})")
    return passed


def test_cross_entropy_loss():
    """Test CrossEntropyLoss forward and backward."""
    from losses import CrossEntropyLoss

    np.random.seed(42)
    loss_fn = CrossEntropyLoss()

    batch_size, num_classes = 4, 5
    logits = np.random.randn(batch_size, num_classes).astype(np.float64)
    labels = np.random.randint(0, num_classes, size=batch_size)

    try:
        loss = loss_fn.forward(logits, labels)
    except NotImplementedError:
        print("  [CrossEntropyLoss forward] NOT IMPLEMENTED — skipping")
        return False

    assert isinstance(loss, (float, np.floating)), "Loss should be scalar"
    assert loss > 0, "Cross-entropy loss should be positive"
    print(f"  [CrossEntropyLoss forward] PASSED ✓  (loss={loss:.4f})")

    try:
        grad = loss_fn.backward()
    except NotImplementedError:
        print("  [CrossEntropyLoss backward] NOT IMPLEMENTED — skipping")
        return False

    # Numerical gradient
    eps = 1e-5
    num_grad = np.zeros_like(logits)
    for i in range(batch_size):
        for j in range(num_classes):
            logits_plus = logits.copy()
            logits_plus[i, j] += eps
            loss_plus = loss_fn.forward(logits_plus, labels)

            logits_minus = logits.copy()
            logits_minus[i, j] -= eps
            loss_minus = loss_fn.forward(logits_minus, labels)

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

    # Restore original forward for correct cache state
    loss_fn.forward(logits, labels)

    rel_err = relative_error(grad, num_grad)
    passed = rel_err < 1e-5
    status = "PASSED ✓" if passed else "FAILED ✗"
    print(f"  [CrossEntropyLoss backward] {status}  (rel error: {rel_err:.2e})")
    return passed


def test_mlp_forward_backward():
    """Test full MLP forward and backward pass."""
    from model import MLP
    from losses import CrossEntropyLoss

    np.random.seed(42)

    try:
        model = MLP(input_dim=10, hidden_dims=[8, 6], num_classes=3,
                     activation='relu')
    except NotImplementedError:
        print("  [MLP construction] NOT IMPLEMENTED — skipping")
        return False

    X = np.random.randn(4, 10).astype(np.float64)
    labels = np.array([0, 1, 2, 1])

    # Convert all weights to float64 for numerical precision
    for layer in model.layers:
        if hasattr(layer, 'W'):
            layer.W = layer.W.astype(np.float64)
            layer.b = layer.b.astype(np.float64)

    try:
        logits = model.forward(X)
    except NotImplementedError:
        print("  [MLP forward] NOT IMPLEMENTED — skipping")
        return False

    assert logits.shape == (4, 3), f"Expected shape (4, 3), got {logits.shape}"
    print(f"  [MLP forward] PASSED ✓  (output shape: {logits.shape})")

    loss_fn = CrossEntropyLoss()
    try:
        loss = loss_fn.forward(logits, labels)
        grad_logits = loss_fn.backward()
        model.backward(grad_logits)
    except NotImplementedError:
        print("  [MLP backward] NOT IMPLEMENTED — skipping")
        return False

    # Check that all Linear layers have non-None gradients
    all_grads_ok = True
    for i, layer in enumerate(model.layers):
        if hasattr(layer, 'grad_W'):
            if layer.grad_W is None or layer.grad_b is None:
                print(f"  [MLP backward] Layer {i} has None gradients!")
                all_grads_ok = False

    if all_grads_ok:
        print("  [MLP backward] PASSED ✓  (all gradients computed)")
    return all_grads_ok


def main():
    print("\n" + "=" * 50)
    print("  CW3: Gradient Checking")
    print("=" * 50)

    results = []

    print("\n--- Linear Layer ---")
    results.append(test_linear_layer())

    print("\n--- ReLU Layer ---")
    results.append(test_relu_layer())

    print("\n--- Sigmoid Layer ---")
    results.append(test_sigmoid_layer())

    print("\n--- CrossEntropyLoss ---")
    results.append(test_cross_entropy_loss())

    print("\n--- Full MLP Forward/Backward ---")
    results.append(test_mlp_forward_backward())

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


if __name__ == "__main__":
    main()
