"""Shared test fixtures for LPIPS-NNX tests."""

import jax.numpy as jnp
import numpy as np
import pytest
import torch


@pytest.fixture
def random_images_np():
    """Generate random image pairs as numpy arrays."""
    np.random.seed(42)
    x = np.random.rand(2, 64, 64, 3).astype(np.float32) * 2 - 1  # [-1, 1]
    y = np.random.rand(2, 64, 64, 3).astype(np.float32) * 2 - 1
    return x, y


@pytest.fixture
def random_images_jax(random_images_np):
    """Convert numpy images to JAX arrays (NHWC)."""
    x_np, y_np = random_images_np
    return jnp.array(x_np), jnp.array(y_np)


@pytest.fixture
def random_images_torch(random_images_np):
    """Convert numpy images to PyTorch tensors (NCHW)."""
    x_np, y_np = random_images_np
    # Convert NHWC to NCHW for PyTorch
    x_pt = torch.from_numpy(np.transpose(x_np, (0, 3, 1, 2)))
    y_pt = torch.from_numpy(np.transpose(y_np, (0, 3, 1, 2)))
    return x_pt, y_pt
