"""Tests comparing LPIPS-NNX to PyTorch reference implementation."""

import flax.nnx as nnx
import jax.numpy as jnp
import lpips as lpips_pt
import numpy as np
import torch


def test_lpips_matches_pytorch(random_images_jax, random_images_torch):
    """Test that LPIPS-NNX produces same output as PyTorch LPIPS."""
    from lpips_nnx import LPIPS

    x_jax, y_jax = random_images_jax
    x_pt, y_pt = random_images_torch

    # JAX LPIPS
    lpips_jax = LPIPS(rngs=nnx.Rngs(0))
    dist_jax = lpips_jax(x_jax, y_jax)

    # PyTorch LPIPS
    lpips_torch = lpips_pt.LPIPS(net="vgg")
    with torch.no_grad():
        dist_pt = lpips_torch(x_pt, y_pt)

    # Convert to numpy for comparison
    dist_jax_np = np.array(dist_jax).squeeze()
    dist_pt_np = dist_pt.numpy().squeeze()

    np.testing.assert_allclose(
        dist_jax_np,
        dist_pt_np,
        atol=1e-4,
        rtol=1e-3,
        err_msg="LPIPS outputs do not match between JAX and PyTorch",
    )


def test_lpips_zero_for_identical_images(random_images_jax):
    """Test that LPIPS is zero for identical images."""
    from lpips_nnx import LPIPS

    x_jax, _ = random_images_jax

    lpips_jax = LPIPS(rngs=nnx.Rngs(0))
    dist = lpips_jax(x_jax, x_jax)

    np.testing.assert_allclose(
        np.array(dist).squeeze(),
        0.0,
        atol=1e-6,
        err_msg="LPIPS should be zero for identical images",
    )


def test_lpips_symmetric(random_images_jax):
    """Test that LPIPS(x, y) == LPIPS(y, x)."""
    from lpips_nnx import LPIPS

    x_jax, y_jax = random_images_jax

    lpips_jax = LPIPS(rngs=nnx.Rngs(0))
    dist_xy = lpips_jax(x_jax, y_jax)
    dist_yx = lpips_jax(y_jax, x_jax)

    np.testing.assert_allclose(
        np.array(dist_xy),
        np.array(dist_yx),
        atol=1e-6,
        err_msg="LPIPS should be symmetric",
    )


def test_lpips_normalize_flag(random_images_jax, random_images_torch):
    """Test the normalize flag for [0, 1] input range."""
    from lpips_nnx import LPIPS

    # Convert images to [0, 1] range
    x_jax, y_jax = random_images_jax
    x_01 = (x_jax + 1) / 2
    y_01 = (y_jax + 1) / 2

    x_pt, y_pt = random_images_torch
    x_pt_01 = (x_pt + 1) / 2
    y_pt_01 = (y_pt + 1) / 2

    # JAX LPIPS with normalize=True
    lpips_jax = LPIPS(rngs=nnx.Rngs(0))
    dist_jax = lpips_jax(x_01, y_01, normalize=True)

    # PyTorch LPIPS with normalize=True
    lpips_torch = lpips_pt.LPIPS(net="vgg")
    with torch.no_grad():
        dist_pt = lpips_torch(x_pt_01, y_pt_01, normalize=True)

    np.testing.assert_allclose(
        np.array(dist_jax).squeeze(),
        dist_pt.numpy().squeeze(),
        atol=1e-4,
        rtol=1e-3,
        err_msg="LPIPS with normalize=True should match PyTorch",
    )


def test_lpips_different_resolutions():
    """Test LPIPS works with different image resolutions."""
    from lpips_nnx import LPIPS

    lpips_jax = LPIPS(rngs=nnx.Rngs(0))

    for size in [32, 64, 128, 224]:
        x = jnp.zeros((1, size, size, 3))
        y = jnp.ones((1, size, size, 3)) * 0.5
        dist = lpips_jax(x, y)
        assert dist.shape == (1, 1, 1, 1), f"Unexpected shape for size {size}"
        assert jnp.isfinite(dist).all(), f"Non-finite LPIPS for size {size}"


def test_lpips_batch_processing(random_images_jax, random_images_torch):
    """Test that batch processing matches element-wise processing."""
    from lpips_nnx import LPIPS

    x_jax, y_jax = random_images_jax
    x_pt, y_pt = random_images_torch

    lpips_jax = LPIPS(rngs=nnx.Rngs(0))
    lpips_torch = lpips_pt.LPIPS(net="vgg")

    # Batch processing
    dist_batch_jax = lpips_jax(x_jax, y_jax)
    with torch.no_grad():
        dist_batch_pt = lpips_torch(x_pt, y_pt)

    # Element-wise processing
    for i in range(x_jax.shape[0]):
        dist_i_jax = lpips_jax(x_jax[i : i + 1], y_jax[i : i + 1])
        with torch.no_grad():
            dist_i_pt = lpips_torch(x_pt[i : i + 1], y_pt[i : i + 1])

        np.testing.assert_allclose(
            np.array(dist_batch_jax[i]).squeeze(),
            np.array(dist_i_jax).squeeze(),
            atol=1e-6,
            err_msg=f"Batch vs element-wise mismatch for JAX at index {i}",
        )
        np.testing.assert_allclose(
            np.array(dist_i_jax).squeeze(),
            dist_i_pt.numpy().squeeze(),
            atol=1e-4,
            rtol=1e-3,
            err_msg=f"JAX vs PyTorch element-wise mismatch at index {i}",
        )
