#!/usr/bin/env python
"""
Forward Pass Tests for TinyTorch
=================================
Tests that all architectures can do forward passes correctly.
This validates the "plumbing" - data flows through without errors.
"""

import sys
import os
import numpy as np
rng = np.random.default_rng(7)

# Add project root to path
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
sys.path.insert(0, project_root)

from tinytorch.core.tensor import Tensor
from tinytorch.core.layers import Linear
from tinytorch.core.activations import ReLU, Sigmoid, Tanh, Softmax
from tinytorch.core.spatial import Conv2d
from tinytorch.core.transformers import TransformerBlock, LayerNorm
from tinytorch.core.embeddings import Embedding, PositionalEncoding

class Sequential:
    """Simple sequential container for testing."""
    def __init__(self, layers):
        self.layers = layers
    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x
    def parameters(self):
        params = []
        for layer in self.layers:
            if hasattr(layer, 'parameters'):
                params.extend(layer.parameters())
        return params

class F:
    """Functional interface for testing."""
    @staticmethod
    def relu(x):
        from tinytorch.core.activations import ReLU
        return ReLU()(x)
    @staticmethod
    def sigmoid(x):
        from tinytorch.core.activations import Sigmoid
        return Sigmoid()(x)
    @staticmethod
    def tanh(x):
        from tinytorch.core.activations import Tanh
        return Tanh()(x)
    @staticmethod
    def softmax(x, dim=-1):
        from tinytorch.core.activations import Softmax
        return Softmax()(x)
    @staticmethod
    def max_pool2d(x, kernel_size):
        from tinytorch.core.spatial import MaxPool2d
        return MaxPool2d(kernel_size)(x)
    @staticmethod
    def avg_pool2d(x, kernel_size):
        from tinytorch.core.spatial import AvgPool2d
        return AvgPool2d(kernel_size)(x)
    @staticmethod
    def flatten(x, start_dim=1):
        import numpy as np
        shape = x.shape
        new_shape = shape[:start_dim] + (np.prod(shape[start_dim:]),)
        return x.reshape(*new_shape)


class ForwardPassTester:
    """Test forward passes for various architectures."""

    def __init__(self):
        self.passed = []
        self.failed = []

    def test(self, name, func):
        """Run a test and track results."""
        try:
            func()
            self.passed.append(name)
            print(f"✅ {name}")
            return True
        except Exception as e:
            self.failed.append((name, str(e)))
            print(f"❌ {name}: {e}")
            return False

    def summary(self):
        """Print test summary."""
        total = len(self.passed) + len(self.failed)
        print(f"\n{'='*60}")
        print(f"FORWARD PASS TESTS: {len(self.passed)}/{total} passed")
        if self.failed:
            print("\nFailed tests:")
            for name, error in self.failed:
                print(f"  - {name}: {error}")
        return len(self.failed) == 0


# Test different layer types
def test_linear_forward():
    """Test Linear layer forward pass."""
    layer = Linear(10, 5)
    x = Tensor(rng.standard_normal((3, 10)))
    y = layer(x)
    assert y.shape == (3, 5)


def test_conv2d_forward():
    """Test Conv2d forward pass."""
    layer = Conv2d(3, 16, kernel_size=3)
    x = Tensor(rng.standard_normal((2, 3, 32, 32)))
    y = layer(x)
    assert y.shape == (2, 16, 30, 30)


def test_conv2d_with_padding():
    """Test Conv2d with padding."""
    layer = Conv2d(3, 16, kernel_size=3, padding=1)
    x = Tensor(rng.standard_normal((2, 3, 32, 32)))
    y = layer(x)
    assert y.shape == (2, 16, 32, 32)  # Same size with padding=1


def test_conv2d_with_stride():
    """Test Conv2d with stride."""
    layer = Conv2d(3, 16, kernel_size=3, stride=2)
    x = Tensor(rng.standard_normal((2, 3, 32, 32)))
    y = layer(x)
    assert y.shape == (2, 16, 15, 15)  # (32-3)/2 + 1 = 15


# Test activation functions
def test_relu_forward():
    """Test ReLU activation."""
    x = Tensor(np.array([[-1, 0, 1], [2, -3, 4]]))
    y = F.relu(x)
    assert y.shape == x.shape


def test_sigmoid_forward():
    """Test Sigmoid activation."""
    x = Tensor(rng.standard_normal((2, 3)))
    y = F.sigmoid(x)
    assert y.shape == x.shape
    # Check sigmoid bounds
    assert np.all(y.data >= 0) and np.all(y.data <= 1)


def test_tanh_forward():
    """Test Tanh activation."""
    x = Tensor(rng.standard_normal((2, 3)))
    y = F.tanh(x)
    assert y.shape == x.shape
    # Check tanh bounds
    assert np.all(y.data >= -1) and np.all(y.data <= 1)


def test_softmax_forward():
    """Test Softmax activation."""
    x = Tensor(rng.standard_normal((2, 10)))
    y = F.softmax(x, dim=-1)
    assert y.shape == x.shape
    # Check softmax sums to 1
    sums = np.sum(y.data, axis=-1)
    assert np.allclose(sums, 1.0)


# Test pooling operations
def test_maxpool2d_forward():
    """Test MaxPool2d."""
    x = Tensor(rng.standard_normal((2, 16, 32, 32)))
    y = F.max_pool2d(x, kernel_size=2)
    assert y.shape == (2, 16, 16, 16)


def test_avgpool2d_forward():
    """Test AvgPool2d."""
    x = Tensor(rng.standard_normal((2, 16, 32, 32)))
    y = F.avg_pool2d(x, kernel_size=2)
    assert y.shape == (2, 16, 16, 16)


# Test reshape operations
def test_flatten_forward():
    """Test flatten operation."""
    x = Tensor(rng.standard_normal((2, 3, 4, 5)))
    y = F.flatten(x, start_dim=1)
    assert y.shape == (2, 60)  # 3*4*5 = 60


def test_reshape_forward():
    """Test reshape operation."""
    x = Tensor(rng.standard_normal((2, 3, 4)))
    y = x.reshape(6, 4)
    assert y.shape == (6, 4)


# Test normalization layers
def test_layernorm_forward():
    """Test LayerNorm."""
    layer = LayerNorm(128)
    x = Tensor(rng.standard_normal((2, 10, 128)))
    y = layer(x)
    assert y.shape == x.shape


def test_batchnorm_forward():
    """Test BatchNorm (if implemented)."""
    # Skip if not implemented
    try:
        from tinytorch.nn import BatchNorm1d
        layer = BatchNorm1d(128)
        x = Tensor(rng.standard_normal((32, 128)))
        y = layer(x)
        assert y.shape == x.shape
    except ImportError:
        pass  # BatchNorm not implemented yet


# Test complex architectures
def test_sequential_forward():
    """Test Sequential container."""
    model = Sequential([
        Linear(10, 20),
        ReLU(),
        Linear(20, 30),
        ReLU(),
        Linear(30, 5)
    ])
    x = Tensor(rng.standard_normal((4, 10)))
    y = model(x)
    assert y.shape == (4, 5)


def test_mlp_forward():
    """Test Multi-Layer Perceptron."""
    class MLP:
        def __init__(self):
            self.fc1 = Linear(784, 256)
            self.fc2 = Linear(256, 128)
            self.fc3 = Linear(128, 10)

        def forward(self, x):
            x = F.relu(self.fc1(x))
            x = F.relu(self.fc2(x))
            return self.fc3(x)

    model = MLP()
    x = Tensor(rng.standard_normal((32, 784)))  # MNIST batch
    y = model.forward(x)
    assert y.shape == (32, 10)


def test_cnn_forward():
    """Test Convolutional Neural Network."""
    class CNN:
        def __init__(self):
            self.conv1 = Conv2d(1, 32, 3)
            self.conv2 = Conv2d(32, 64, 3)
            self.fc1 = Linear(64 * 5 * 5, 128)
            self.fc2 = Linear(128, 10)

        def forward(self, x):
            x = F.relu(self.conv1(x))
            x = F.max_pool2d(x, 2)
            x = F.relu(self.conv2(x))
            x = F.max_pool2d(x, 2)
            x = F.flatten(x, start_dim=1)
            x = F.relu(self.fc1(x))
            return self.fc2(x)

    model = CNN()
    x = Tensor(rng.standard_normal((16, 1, 28, 28)))  # MNIST batch
    y = model.forward(x)
    assert y.shape == (16, 10)


def test_transformer_forward():
    """Test Transformer architecture."""
    class SimpleTransformer:
        def __init__(self):
            self.embed = Embedding(1000, 128)  # vocab_size=1000, embed_dim=128
            self.pos_enc = PositionalEncoding(100, 128)  # max_seq_len=100, embed_dim=128
            self.transformer = TransformerBlock(128, 8)
            self.ln = LayerNorm(128)
            self.output = Linear(128, 1000)

        def forward(self, x):
            x = self.embed(x)
            x = self.pos_enc(x)
            x = self.transformer(x)
            x = self.ln(x)
            # Reshape for output
            batch, seq, embed = x.shape
            x = x.reshape(batch * seq, embed)
            x = self.output(x)
            return x.reshape(batch, seq, 1000)

    model = SimpleTransformer()
    x = Tensor(rng.integers(0, 1000, (4, 20)))  # Token batch
    y = model.forward(x)
    assert y.shape == (4, 20, 1000)


def test_residual_block_forward():
    """Test Residual Block (ResNet-style)."""
    class ResidualBlock:
        def __init__(self, channels):
            self.conv1 = Conv2d(channels, channels, 3, padding=1)
            self.conv2 = Conv2d(channels, channels, 3, padding=1)

        def forward(self, x):
            identity = x
            out = F.relu(self.conv1(x))
            out = self.conv2(out)
            out = out + identity  # Residual connection
            return F.relu(out)

    block = ResidualBlock(64)
    x = Tensor(rng.standard_normal((2, 64, 16, 16)))
    y = block.forward(x)
    assert y.shape == x.shape


def run_all_forward_tests():
    """Run comprehensive forward pass tests."""
    print("="*60)
    print("FORWARD PASS TEST SUITE")
    print("Testing data flow through all layer types")
    print("="*60)

    tester = ForwardPassTester()

    # Basic layers
    print("\n📦 Basic Layers:")
    tester.test("Linear layer", test_linear_forward)
    tester.test("Conv2d layer", test_conv2d_forward)
    tester.test("Conv2d with padding", test_conv2d_with_padding)
    tester.test("Conv2d with stride", test_conv2d_with_stride)

    # Activations
    print("\n⚡ Activation Functions:")
    tester.test("ReLU", test_relu_forward)
    tester.test("Sigmoid", test_sigmoid_forward)
    tester.test("Tanh", test_tanh_forward)
    tester.test("Softmax", test_softmax_forward)

    # Pooling
    print("\n🏊 Pooling Operations:")
    tester.test("MaxPool2d", test_maxpool2d_forward)
    tester.test("AvgPool2d", test_avgpool2d_forward)

    # Reshaping
    print("\n🔄 Reshape Operations:")
    tester.test("Flatten", test_flatten_forward)
    tester.test("Reshape", test_reshape_forward)

    # Normalization
    print("\n📊 Normalization:")
    tester.test("LayerNorm", test_layernorm_forward)
    tester.test("BatchNorm", test_batchnorm_forward)

    # Full architectures
    print("\n🏗️ Complete Architectures:")
    tester.test("Sequential container", test_sequential_forward)
    tester.test("MLP (MNIST)", test_mlp_forward)
    tester.test("CNN (Images)", test_cnn_forward)
    tester.test("Transformer (NLP)", test_transformer_forward)
    tester.test("Residual Block", test_residual_block_forward)

    return tester.summary()


if __name__ == "__main__":
    success = run_all_forward_tests()
    sys.exit(0 if success else 1)