import torch
import blaze as bl


def test_jit_trace():
    """Test that the model can be JIT traced."""

    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))

    jit_model = torch.jit.trace(model, torch.randn(2, 10))
    out = jit_model(torch.randn(2, 10))
    assert out.shape == (2, 5)


def test_jit_trace_training():
    """Test that a JIT-traced model can be trained."""

    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))

    jit_model = torch.jit.trace(model, torch.randn(2, 10))
    optimizer = torch.optim.Adam(jit_model.parameters(), lr=1e-3)

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


def test_jit_trace_multiple_inputs():
    """Test that the model can handle multiple inputs in JIT trace."""

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

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

    jit_model = torch.jit.trace(model, (torch.randn(2, 10), torch.randn(2, 10)))
    out = jit_model(torch.randn(2, 10), torch.randn(2, 10))
    assert out.shape == (2, 5)


def test_jit_trace_multiple_outputs():
    """Test that the model can handle multiple outputs in JIT trace."""

    def forward(x):
        x = bl.Linear(10, 20)(x)
        y = bl.Linear(20, 20)(x)
        z = bl.Linear(20, 30)(x)
        return y, z

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

    jit_model = torch.jit.trace(model, torch.randn(2, 10))
    y, z = jit_model(torch.randn(2, 10))
    assert y.shape == (2, 20)
    assert z.shape == (2, 30)
