import pytest
import torch
import blaze as bl
from conftest import COMPILE_BACKENDS


def test_torch_compile():
    def forward(x):
        x = bl.Linear(10, 20)(x)
        x = bl.ReLU()(x)
        x = bl.Linear(20, 5)(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))

    compiled = torch.compile(model, backend="eager")
    out = compiled(torch.randn(2, 10))
    assert out.shape == (2, 5)


def test_torch_compile_training():
    def forward(x):
        x = bl.Linear(10, 20)(x)
        x = bl.ReLU()(x)
        x = bl.Linear(20, 5)(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))

    compiled = torch.compile(model, backend="eager")
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    x = torch.randn(2, 10)
    for _ in range(3):
        optimizer.zero_grad()
        out = compiled(x)
        loss = out.sum()
        loss.backward()
        optimizer.step()


def test_if_else_compile():
    """If/else under torch.compile."""
    def forward(x):
        x = bl.Linear(10, 20)(x)
        if x.shape[0] > 1:
            x = bl.ReLU()(x)
        else:
            x = bl.GELU()(x)
        x = bl.Linear(20, 5)(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))

    compiled = torch.compile(model, backend="eager")
    out = compiled(torch.randn(4, 10))
    assert out.shape == (4, 5)


def test_loop_compile():
    """Fixed-count loop under torch.compile."""
    def forward(x):
        for _ in range(3):
            x = bl.Linear(10, 10)(x)
            x = bl.ReLU()(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))

    compiled = torch.compile(model, backend="eager")
    out = compiled(torch.randn(2, 10))
    assert out.shape == (2, 10)


def test_loop_with_if_compile():
    """Loop with conditional activations under torch.compile."""
    def forward(x):
        for i in range(4):
            x = bl.Linear(10, 10)(x)
            if i % 2 == 0:
                x = bl.ReLU()(x)
            else:
                x = bl.GELU()(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))

    compiled = torch.compile(model, backend="eager")
    out = compiled(torch.randn(2, 10))
    assert out.shape == (2, 10)


def test_loop_compile_training():
    """Loop under torch.compile with backward pass."""
    def forward(x):
        for _ in range(3):
            x = bl.Linear(10, 10)(x)
            x = bl.ReLU()(x)
        return bl.Linear(10, 5)(x)

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))

    compiled = torch.compile(model, backend="eager")
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for _ in range(3):
        optimizer.zero_grad()
        out = compiled(torch.randn(2, 10))
        loss = out.sum()
        loss.backward()
        optimizer.step()


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_correctness(backend):
    """Compiled output matches eager output numerically."""
    def forward(x):
        x = bl.Linear(10, 20)(x)
        x = bl.ReLU()(x)
        x = bl.Linear(20, 5)(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))
    model.eval()

    x = torch.randn(2, 10)
    eager_out = model(x)

    compiled = torch.compile(model, backend=backend)
    compiled_out = compiled(x)
    torch.testing.assert_close(compiled_out, eager_out)


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_training(backend):
    """Backward pass works under compile."""
    def forward(x):
        x = bl.Linear(10, 20)(x)
        x = bl.ReLU()(x)
        x = bl.Linear(20, 5)(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))

    compiled = torch.compile(model, backend=backend)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    params_before = {n: p.clone() for n, p in model.named_parameters()}

    for _ in range(3):
        optimizer.zero_grad()
        out = compiled(torch.randn(2, 10))
        out.sum().backward()
        optimizer.step()

    # Weights should have changed
    for n, p in model.named_parameters():
        assert not torch.equal(p, params_before[n]), f"{n} unchanged after training"


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_nested_module(backend):
    """Nested bl.Module hierarchies under compile."""
    class Block(bl.Module):
        def __call__(self, x, dim=32):
            x = bl.Linear(x.shape[-1], dim)(x)
            x = bl.ReLU()(x)
            return x

    def forward(x):
        x = Block()(x, dim=64)
        x = Block()(x, dim=10)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 16))
    model.eval()

    x = torch.randn(2, 16)
    eager_out = model(x)

    compiled = torch.compile(model, backend=backend)
    torch.testing.assert_close(compiled(x), eager_out)


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_multiple_inputs_outputs(backend):
    """Multiple inputs and outputs under compile."""
    def forward(x, y):
        x = bl.Linear(10, 20)(x)
        y = bl.Linear(10, 20)(y)
        z = x + y
        z = bl.ReLU()(z)
        w = bl.Linear(20, 5)(z)
        v = bl.Linear(20, 5)(z)
        return w, v

    model = bl.transform(forward)
    model.init(torch.randn(2, 10), torch.randn(2, 10))
    model.eval()

    x, y = torch.randn(2, 10), torch.randn(2, 10)
    eager_w, eager_v = model(x, y)

    compiled = torch.compile(model, backend=backend)
    comp_w, comp_v = compiled(x, y)
    torch.testing.assert_close(comp_w, eager_w)
    torch.testing.assert_close(comp_v, eager_v)


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_batchnorm(backend):
    """BatchNorm stateful behavior under compile (train + eval)."""
    def forward(x):
        x = bl.Linear(10, 20)(x)
        x = bl.BatchNorm1d(20)(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(4, 10))

    compiled = torch.compile(model, backend=backend)

    # Train mode: running stats should update
    model.train()
    bn = model._registry["batch_norm1d"]
    mean_before = bn.running_mean.clone()
    compiled(torch.randn(4, 10) + 5.0)
    assert not torch.equal(bn.running_mean, mean_before)

    # Eval mode: deterministic
    model.eval()
    x = torch.randn(4, 10)
    out1 = compiled(x)
    out2 = compiled(x)
    torch.testing.assert_close(out1, out2)


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_get_parameter(backend):
    """get_parameter under compile."""
    def forward(x):
        scale = bl.get_parameter("scale", (5,), init_fn=torch.ones)
        bias = bl.get_parameter("bias", (5,), init_fn=torch.zeros)
        x = bl.Linear(10, 5)(x)
        return x * scale + bias

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))
    model.eval()

    x = torch.randn(2, 10)
    eager_out = model(x)

    compiled = torch.compile(model, backend=backend)
    torch.testing.assert_close(compiled(x), eager_out)


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_get_state(backend):
    """get_state (non-trainable buffer) under compile."""
    def forward(x):
        offset = bl.get_state("offset", (5,), init_fn=torch.ones)
        x = bl.Linear(10, 5)(x)
        return x + offset

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))
    model.eval()

    x = torch.randn(2, 10)
    eager_out = model(x)

    compiled = torch.compile(model, backend=backend)
    torch.testing.assert_close(compiled(x), eager_out)


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_set_state(backend):
    """Mutable state via set_state under compile."""
    def forward(x):
        count = bl.get_state("count", (1,), init_fn=torch.zeros)
        bl.set_state("count", count + 1)
        return x + count

    model = bl.transform(forward)
    model.init(torch.randn(1, 3))

    compiled = torch.compile(model, backend=backend)

    x = torch.zeros(1, 3)
    out1 = compiled(x)
    out2 = compiled(x)
    torch.testing.assert_close(out2 - out1, torch.ones(1, 3))


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_train_eval_switch(backend):
    """Mode switching (train/eval) with compiled model."""
    def forward(x):
        x = bl.Linear(10, 10)(x)
        x = bl.Dropout(0.5)(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))

    compiled = torch.compile(model, backend=backend)

    # Eval: deterministic
    model.eval()
    x = torch.randn(100, 10)
    out1 = compiled(x)
    out2 = compiled(x)
    torch.testing.assert_close(out1, out2)

    # Train: should run without error
    model.train()
    _ = compiled(x)


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_state_dict_roundtrip(backend):
    """Save and load state_dict after compiling."""
    def forward(x):
        x = bl.Linear(10, 20)(x)
        x = bl.ReLU()(x)
        x = bl.Linear(20, 5)(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))
    model.eval()

    compiled = torch.compile(model, backend=backend)

    # Run once to trigger compilation
    x = torch.randn(2, 10)
    out_before = compiled(x)

    # Save and reload
    sd = model.state_dict()
    model2 = bl.transform(forward)
    model2.init(torch.randn(2, 10))
    model2.load_state_dict(sd)
    model2.eval()

    compiled2 = torch.compile(model2, backend=backend)
    torch.testing.assert_close(compiled2(x), out_before)


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_control_flow(backend):
    """Control flow (if/else + loop) under compile."""
    def forward(x):
        for i in range(3):
            x = bl.Linear(10, 10)(x)
            if i % 2 == 0:
                x = bl.ReLU()(x)
            else:
                x = bl.GELU()(x)
        return bl.Linear(10, 5)(x)

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))
    model.eval()

    x = torch.randn(2, 10)
    eager_out = model(x)

    compiled = torch.compile(model, backend=backend)
    torch.testing.assert_close(compiled(x), eager_out)


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_compile_deeply_nested(backend):
    """Deeply nested modules under compile."""
    class Inner(bl.Module):
        def __call__(self, x):
            return bl.Linear(x.shape[-1], x.shape[-1])(x)

    class Outer(bl.Module):
        def __call__(self, x):
            x = Inner()(x)
            x = Inner()(x)
            return x

    def forward(x):
        return Outer()(x)

    model = bl.transform(forward)
    model.init(torch.randn(2, 8))
    model.eval()

    x = torch.randn(2, 8)
    eager_out = model(x)

    compiled = torch.compile(model, backend=backend)
    torch.testing.assert_close(compiled(x), eager_out)
