"""Unit tests for the differentiable PyTorch VLM pipeline.

Slow tests (marked ``slow``) download model weights from HuggingFace on first
run.  Deselect them with ``-m "not slow"``.

Generative model used in slow tests: ``stabilityai/sd-turbo``
  - 1-step latent diffusion, 512x512 output
  - latents shape: (B, 4, 64, 64), latent_dim = 16384

Vision model used in slow tests: ``openai/clip-vit-base-patch32``
  - lightweight CLIP vision encoder, pooled embedding output
"""

import pytest
import torch

from diffusers import DDIMScheduler, StableDiffusionPipeline

from local_geometry.experiments.vlm import (
  DiffusionConfig,
  VLMConfig,
  GenerativeModel,
  VisionModel,
  load_generative_model,
  load_vlm,
  preprocess_image,
  ddim_sample,
  vlm_features,
  latent_to_features,
)

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

DEVICE = torch.device("cpu")

DIFF_CFG = DiffusionConfig(model_name="stabilityai/sd-turbo", num_inference_steps=1)

# Expected latent dimensions for sd-turbo.
LATENT_CHANNELS = 4
LATENT_SIZE = 64
PIXEL_SIZE = 512


@pytest.fixture(scope="module")
def gen_model() -> GenerativeModel:
  return load_generative_model(DIFF_CFG, DEVICE)


@pytest.fixture(scope="module")
def vision_model() -> VisionModel:
  return load_vlm(VLMConfig(), DEVICE)


# ---------------------------------------------------------------------------
# preprocess_image -- no model weights needed
# ---------------------------------------------------------------------------


def test_preprocess_image_output_shape():
  """Output should always be (B, 3, SMOLVLM_PATCH_SIZE, SMOLVLM_PATCH_SIZE)."""
  from local_geometry.experiments.vlm import SMOLVLM_PATCH_SIZE
  images = torch.rand(3, 3, 32, 32)
  out = preprocess_image(images)
  assert out.shape == (3, 3, SMOLVLM_PATCH_SIZE, SMOLVLM_PATCH_SIZE)


def test_preprocess_image_normalization():
  """A uniform image equal to the mean value (0.5) should produce zero output."""
  images = torch.full((1, 3, 32, 32), 0.5)
  out = preprocess_image(images)
  assert out.abs().max().item() < 1e-5, "Normalized mean image should be ~zero."


def test_preprocess_image_wrong_shape():
  """Non (B, 3, H, W) input should raise ValueError."""
  images = torch.rand(4, 1, 32, 32)
  with pytest.raises(ValueError):
    preprocess_image(images)


def test_preprocess_image_preserves_grad():
  """Gradients should flow through both resize steps."""
  images = torch.rand(2, 3, 32, 32, requires_grad=True)
  out = preprocess_image(images)
  out.sum().backward()
  assert images.grad is not None
  assert images.grad.abs().max().item() > 0.0


def test_preprocess_image_output_shape_various_inputs():
  """Output is always (B, 3, SMOLVLM_PATCH_SIZE, SMOLVLM_PATCH_SIZE) regardless of input size."""
  from local_geometry.experiments.vlm import SMOLVLM_PATCH_SIZE
  for h, w in [(32, 32), (100, 200), (500, 300), (1500, 1500)]:
    images = torch.rand(2, 3, h, w)
    out = preprocess_image(images)
    assert out.shape == (2, 3, SMOLVLM_PATCH_SIZE, SMOLVLM_PATCH_SIZE), (
      f"Wrong shape for input ({h}, {w}): got {out.shape}"
    )


def test_preprocess_image_longest_edge_clamp():
  """Images whose longest edge exceeds SMOLVLM_MAX_SIDE are downscaled first."""
  from local_geometry.experiments.vlm import SMOLVLM_PATCH_SIZE, SMOLVLM_MAX_SIDE
  # Use an image much larger than SMOLVLM_MAX_SIDE to confirm the clamp step runs.
  h, w = SMOLVLM_MAX_SIDE + 500, SMOLVLM_MAX_SIDE + 200
  images = torch.rand(1, 3, h, w)
  out = preprocess_image(images)
  assert out.shape == (1, 3, SMOLVLM_PATCH_SIZE, SMOLVLM_PATCH_SIZE)


def test_preprocess_image_matches_processor():
  """Numerically compare preprocess_image output to SmolVLMImageProcessor.

  The processor uses PIL + Lanczos resampling while we use bilinear
  interpolation, so pixel-exact agreement is not expected.  We verify:
  - same output shape
  - output value range matches (both should be roughly in [-1, 1])
  - channel-wise mean is close (within 0.05) for a natural-statistics image
  """
  import numpy as np
  from transformers.models.smolvlm.image_processing_smolvlm import SmolVLMImageProcessor
  from local_geometry.experiments.vlm import SMOLVLM_PATCH_SIZE, SMOLVLM_MAX_SIDE

  processor = SmolVLMImageProcessor(
    do_image_splitting=False,
    max_image_size={"longest_edge": SMOLVLM_PATCH_SIZE},
    size={"longest_edge": SMOLVLM_MAX_SIDE},
  )

  torch.manual_seed(42)
  # Use a moderately sized image to exercise the resize path.
  images_torch = torch.rand(1, 3, 256, 256)

  # Run our differentiable preprocessor.
  our_out = preprocess_image(images_torch)

  # Run the official processor: it expects uint8 HWC numpy images.
  img_np = (images_torch[0].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
  proc_out = processor(images=[img_np], return_tensors="pt", do_image_splitting=False)
  # processor output pixel_values shape: (1, 1, C, H, W) due to the patch dim.
  proc_tensor = proc_out["pixel_values"][0, 0].float()  # (C, H, W)

  assert our_out.shape == (1, 3, SMOLVLM_PATCH_SIZE, SMOLVLM_PATCH_SIZE)
  assert proc_tensor.shape == (3, SMOLVLM_PATCH_SIZE, SMOLVLM_PATCH_SIZE)

  # Both should be in [-1, 1].
  assert our_out.min().item() >= -1.1
  assert our_out.max().item() <= 1.1

  # Channel-wise mean should agree within 0.05 (bilinear vs Lanczos difference).
  our_mean = our_out[0].mean(dim=(-2, -1))
  proc_mean = proc_tensor.mean(dim=(-2, -1))
  channel_diff = (our_mean - proc_mean).abs()
  assert channel_diff.max().item() < 0.05, (
    f"Channel-wise mean disagreement between preprocess_image and SmolVLMImageProcessor: "
    f"our={our_mean.tolist()}, proc={proc_mean.tolist()}"
  )


# ---------------------------------------------------------------------------
# GenerativeModel properties (no model weights needed)
# ---------------------------------------------------------------------------


@pytest.mark.slow
def test_generative_model_latent_dim(gen_model: GenerativeModel):
  """latent_dim should equal latent_channels * latent_size^2."""
  assert gen_model.latent_channels == LATENT_CHANNELS
  assert gen_model.latent_size == LATENT_SIZE
  assert gen_model.latent_dim == LATENT_CHANNELS * LATENT_SIZE * LATENT_SIZE
  assert gen_model.pixel_size == PIXEL_SIZE


# ---------------------------------------------------------------------------
# ddim_sample
# ---------------------------------------------------------------------------


@pytest.mark.slow
def test_ddim_sample_output_shape(gen_model: GenerativeModel):
  """Generated images should have the expected pixel shape."""
  batch = 2
  z = torch.randn(batch, gen_model.latent_dim)
  images = ddim_sample(gen_model, z)
  assert images.shape == (batch, 3, PIXEL_SIZE, PIXEL_SIZE)


@pytest.mark.slow
def test_ddim_sample_output_range(gen_model: GenerativeModel):
  """Generated images should be clipped to [0, 1]."""
  z = torch.randn(1, gen_model.latent_dim)
  images = ddim_sample(gen_model, z)
  assert images.min().item() >= 0.0
  assert images.max().item() <= 1.0


@pytest.mark.slow
def test_ddim_sample_latent_shape(gen_model: GenerativeModel):
  """The reshaped input latent should have the expected spatial shape."""
  batch = 2
  z = torch.randn(batch, gen_model.latent_dim)
  latents = z.reshape(batch, gen_model.latent_channels, gen_model.latent_size, gen_model.latent_size)
  assert latents.shape == (batch, LATENT_CHANNELS, LATENT_SIZE, LATENT_SIZE)


@pytest.mark.slow
def test_ddim_sample_deterministic(gen_model: GenerativeModel):
  """Same latent should produce the same image (deterministic DDIM)."""
  z = torch.randn(1, gen_model.latent_dim)
  img1 = ddim_sample(gen_model, z)
  img2 = ddim_sample(gen_model, z)
  assert torch.allclose(img1, img2), "DDIM sample is not deterministic."


# ---------------------------------------------------------------------------
# vlm_features
# ---------------------------------------------------------------------------


@pytest.mark.slow
def test_vlm_features_shape(vision_model: VisionModel):
  """VLM logit output should have shape (B, vocab_size)."""
  batch = 3
  images = torch.rand(batch, 3, 32, 32)
  features = vlm_features(vision_model, images)
  assert features.ndim == 2
  assert features.shape[0] == batch


@pytest.mark.slow
def test_vlm_features_gradient(vision_model: VisionModel):
  """Gradients should flow from VLM features back to the input images."""
  images = torch.rand(2, 3, 32, 32, requires_grad=True)
  features = vlm_features(vision_model, images)
  features.sum().backward()
  assert images.grad is not None
  assert images.grad.abs().max().item() > 0.0


# ---------------------------------------------------------------------------
# Full pipeline
# ---------------------------------------------------------------------------


@pytest.mark.slow
def test_full_pipeline_forward(gen_model: GenerativeModel, vision_model: VisionModel):
  """End-to-end forward pass should complete and return sane shapes."""
  batch = 1
  z = torch.randn(batch, gen_model.latent_dim)
  features, images = latent_to_features(z, gen_model, vision_model)

  assert images.shape == (batch, 3, PIXEL_SIZE, PIXEL_SIZE)
  assert features.shape[0] == batch
  assert features.ndim == 2


@pytest.mark.slow
def test_full_pipeline_gradient_flows(gen_model: GenerativeModel, vision_model: VisionModel):
  """Gradients should flow from the feature loss all the way back to z."""
  batch = 1
  z = torch.randn(batch, gen_model.latent_dim, requires_grad=True)

  features, _ = latent_to_features(z, gen_model, vision_model)
  loss = features.pow(2).mean()
  loss.backward()

  assert z.grad is not None, "z.grad is None -- gradients did not flow."
  assert z.grad.abs().max().item() > 0.0, "z.grad is all zeros -- gradients are vanishing."


# ---------------------------------------------------------------------------
# Backprop correctness: scheduler config and clamp gradient behavior
# ---------------------------------------------------------------------------


@pytest.mark.slow
def test_scheduler_clip_sample_disabled():
  """DDIMScheduler built from sd-turbo config must have clip_sample=False.

  clip_sample=True inserts a clamp inside the denoising loop, silently
  killing gradients at saturated values.  This test fails loudly if the
  model config ever changes or a different model with clip_sample=True
  is substituted.
  """
  pipe = StableDiffusionPipeline.from_pretrained(DIFF_CFG.model_name, torch_dtype=torch.float32)
  scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
  assert not scheduler.config.clip_sample, (
    f"DDIMScheduler has clip_sample=True for {DIFF_CFG.model_name}. "
    "This will kill gradients inside the denoising loop."
  )


def test_ddim_clamp_gradient_behavior():
  """Verify that .clamp(0,1) passes gradients only for unsaturated pixels.

  The final line of ddim_sample applies ((x+1)/2).clamp(0,1).  The VAE
  outputs values nominally in [-1,1], mapping to [0,1] after the shift.
  Saturation occurs when the VAE output falls outside [-1,1], i.e. raw
  values below -1 map to pre-clamp values below 0, and raw values above 1
  map to pre-clamp values above 1.  This test makes the behavior explicit.
  """
  # raw = -3.0 -> pre-clamp = (-3+1)/2 = -1.0 -> saturated (below 0)
  # raw =  0.0 -> pre-clamp = ( 0+1)/2 =  0.5 -> unsaturated
  # raw =  3.0 -> pre-clamp = ( 3+1)/2 =  2.0 -> saturated (above 1)
  raw = torch.tensor(
    [[[[-3.0, 0.0, 3.0]]]],
    dtype=torch.float32,
    requires_grad=True,
  )

  # Replicate to shape (1, 3, 1, 3) for the pipeline's (B, C, H, W) convention.
  images = raw.expand(1, 3, 1, 3)
  clamped = ((images + 1.0) / 2.0).clamp(0.0, 1.0)
  clamped.sum().backward()

  grad = raw.grad  # shape (1, 1, 1, 3) -- values for below / inside / above

  assert grad is not None
  # Pixel whose pre-clamp value is below 0: saturated, gradient must be 0.
  assert grad[0, 0, 0, 0].item() == 0.0, "Expected zero gradient for saturated (below) pixel."
  # Pixel whose pre-clamp value is inside (0,1): gradient must be nonzero.
  assert grad[0, 0, 0, 1].item() != 0.0, "Expected nonzero gradient for unsaturated pixel."
  # Pixel whose pre-clamp value is above 1: saturated, gradient must be 0.
  assert grad[0, 0, 0, 2].item() == 0.0, "Expected zero gradient for saturated (above) pixel."


@pytest.mark.slow
def test_full_pipeline_gradient_finite_difference(
  gen_model: GenerativeModel, vision_model: VisionModel
):
  """Validate the pipeline gradient numerically via central finite differences.

  This test checks that the analytic gradient from backward() agrees with a
  central-difference estimate on a single scalar perturbation of z.  This
  catches cases where clamp or other operations kill most of the gradient
  signal and the aggregate-norm test (test_full_pipeline_gradient_flows)
  would still pass.
  """
  batch = 1
  torch.manual_seed(0)
  z = torch.randn(batch, gen_model.latent_dim, requires_grad=True)

  # Analytic gradient.
  features, _ = latent_to_features(z, gen_model, vision_model)
  loss = features.pow(2).mean()
  loss.backward()
  analytic_grad = z.grad.clone()

  # Central-difference estimate on the first coordinate only.
  eps = 1e-3
  z0 = z.detach().clone()

  z_plus = z0.clone()
  z_plus[0, 0] += eps
  with torch.no_grad():
    feat_plus, _ = latent_to_features(z_plus, gen_model, vision_model)
  loss_plus = feat_plus.pow(2).mean().item()

  z_minus = z0.clone()
  z_minus[0, 0] -= eps
  with torch.no_grad():
    feat_minus, _ = latent_to_features(z_minus, gen_model, vision_model)
  loss_minus = feat_minus.pow(2).mean().item()

  fd_grad_0 = (loss_plus - loss_minus) / (2.0 * eps)
  analytic_grad_0 = analytic_grad[0, 0].item()

  # The clamp in ddim_sample introduces kinks that cause finite-difference
  # estimates to disagree with the analytic gradient by up to ~20% when the
  # probe lands near a saturation boundary.  We use a 50% tolerance, which
  # is loose enough to survive that non-smoothness but tight enough to catch
  # a catastrophically broken gradient (e.g. zero or wrong sign).
  rel_err = abs(analytic_grad_0 - fd_grad_0) / (abs(fd_grad_0) + 1e-8)
  assert rel_err < 0.50, (
    f"Analytic gradient ({analytic_grad_0:.6f}) disagrees with finite-difference "
    f"estimate ({fd_grad_0:.6f}) by {rel_err*100:.1f}%. "
    "This suggests gradient flow through the pipeline is broken."
  )
