"""
Tests that pin the input/output behavior of each component being modified
in the Flax SD-Turbo compilation speedup changes.

Each test runs a forward pass, checks the result, and is designed so that
it will also pass after the corresponding change is made (i.e. the test
documents the invariant, not the current broken behavior).
"""
import sys
import math
from pathlib import Path

import jax
import jax.numpy as jnp
import flax.linen as nn
import numpy as np
import pytest

# Make the hf_copy package importable as a standalone package.
HF_COPY = Path(__file__).parent.parent / "local_geometry/model_overviews/flax_sd_turbo/hf_copy"
sys.path.insert(0, str(HF_COPY.parent.parent.parent))

from local_geometry.model_overviews.flax_sd_turbo.hf_copy.attention_flax import (
  jax_memory_efficient_attention,
  _query_chunk_attention,
  FlaxAttention,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _rng(seed: int = 0):
  return jax.random.PRNGKey(seed)


def _naive_attention(query, key, value):
  """Reference: plain O(N^2) attention in (seq, heads, dim) layout."""
  # query/key/value: (seq, heads, dim)
  scale = query.shape[-1] ** -0.5
  scores = jnp.einsum("qhd,khd->qhk", query * scale, key)
  probs = jax.nn.softmax(scores, axis=-1)
  return jnp.einsum("qhk,khd->qhd", probs, value)


# ---------------------------------------------------------------------------
# 1. jax_memory_efficient_attention
# ---------------------------------------------------------------------------

class TestMemoryEfficientAttention:
  """
  Checks that jax_memory_efficient_attention matches naive attention, and
  that adding unroll=1 to the internal jax.lax.scan does not change the output.

  The current implementation does NOT have unroll=1, so this test documents
  the correct numerical behavior that must be preserved after the fix.
  """

  def _make_qkv(self, batch=1, seq=64, heads=4, dim=32, seed=0):
    k1, k2, k3 = jax.random.split(_rng(seed), 3)
    q = jax.random.normal(k1, (batch, seq, heads, dim))
    k = jax.random.normal(k2, (batch, seq, heads, dim))
    v = jax.random.normal(k3, (batch, seq, heads, dim))
    return q, k, v

  def test_matches_naive_attention(self):
    """Memory-efficient attention must be numerically close to naive attention."""
    q, k, v = self._make_qkv(seq=64, heads=4, dim=32)
    out_efficient = jax_memory_efficient_attention(q, k, v, query_chunk_size=16, key_chunk_size=32)
    out_naive = _naive_attention(q[0], k[0], v[0])
    assert jnp.allclose(out_efficient[0], out_naive, atol=1e-4), (
      f"max diff: {jnp.abs(out_efficient[0] - out_naive).max()}"
    )

  def test_chunking_does_not_change_output(self):
    """Different query_chunk_size values should produce the same result."""
    q, k, v = self._make_qkv(seq=64, heads=4, dim=32)
    out_chunk16 = jax_memory_efficient_attention(q, k, v, query_chunk_size=16, key_chunk_size=32)
    out_chunk32 = jax_memory_efficient_attention(q, k, v, query_chunk_size=32, key_chunk_size=32)
    assert jnp.allclose(out_chunk16, out_chunk32, atol=1e-5), (
      f"max diff: {jnp.abs(out_chunk16 - out_chunk32).max()}"
    )

  def test_single_chunk_matches_full_seq(self):
    """When query_chunk_size == seq_len there is only one scan step."""
    q, k, v = self._make_qkv(seq=64, heads=4, dim=32)
    out_multi = jax_memory_efficient_attention(q, k, v, query_chunk_size=16, key_chunk_size=32)
    out_single = jax_memory_efficient_attention(q, k, v, query_chunk_size=64, key_chunk_size=64)
    assert jnp.allclose(out_multi, out_single, atol=1e-5), (
      f"max diff: {jnp.abs(out_multi - out_single).max()}"
    )


# ---------------------------------------------------------------------------
# 2. FlaxAttention split_head_dim parity
# ---------------------------------------------------------------------------

class TestFlaxAttentionSplitHeadDim:
  """
  Checks that FlaxAttention with split_head_dim=False and split_head_dim=True
  produce bit-exact (or numerically equivalent) outputs given the same weights
  and input. This pins the equivalence before we switch split_head_dim=True
  as the default.
  """

  def _build_and_init(self, query_dim=256, heads=8, dim_head=32, split_head_dim=False, seed=0):
    import flax.linen as nn
    model = FlaxAttention(
      query_dim=query_dim,
      heads=heads,
      dim_head=dim_head,
      dropout=0.0,
      use_memory_efficient_attention=False,
      split_head_dim=split_head_dim,
      dtype=jnp.float32,
    )
    # Create a dummy input to initialise parameters
    batch, seq = 1, 16
    x = jnp.ones((batch, seq, query_dim))
    params = model.init(_rng(seed), x)
    return model, params

  def test_split_head_dim_false_and_true_agree(self):
    """
    The two code paths must produce numerically identical results for the
    same weights and input.
    """
    query_dim, heads, dim_head = 256, 8, 32
    batch, seq = 2, 16

    model_false, params_false = self._build_and_init(
      query_dim=query_dim, heads=heads, dim_head=dim_head, split_head_dim=False, seed=7
    )
    model_true, params_true = self._build_and_init(
      query_dim=query_dim, heads=heads, dim_head=dim_head, split_head_dim=True, seed=7
    )

    # Both models were initialised from the same seed so their parameter
    # *values* are the same (Flax init is deterministic given the same key).
    x = jax.random.normal(_rng(42), (batch, seq, query_dim))

    out_false = model_false.apply(params_false, x, deterministic=True)
    out_true = model_true.apply(params_true, x, deterministic=True)

    assert out_false.shape == out_true.shape, (
      f"shape mismatch: {out_false.shape} vs {out_true.shape}"
    )
    assert jnp.allclose(out_false, out_true, atol=1e-5), (
      f"max diff: {jnp.abs(out_false - out_true).max()}"
    )

  def test_split_head_dim_output_shape(self):
    """Output shape must be (batch, seq, query_dim) regardless of split_head_dim."""
    query_dim, heads, dim_head = 256, 8, 32
    batch, seq = 2, 16
    x = jax.random.normal(_rng(1), (batch, seq, query_dim))

    for split in (False, True):
      model, params = self._build_and_init(
        query_dim=query_dim, heads=heads, dim_head=dim_head, split_head_dim=split
      )
      out = model.apply(params, x, deterministic=True)
      assert out.shape == (batch, seq, query_dim), (
        f"split_head_dim={split}: expected {(batch, seq, query_dim)}, got {out.shape}"
      )


# ---------------------------------------------------------------------------
# 3. convert_pytorch_state_dict_to_flax dtype
# ---------------------------------------------------------------------------

class TestConvertPyTorchStateDictDtype:
  """
  The fix to modeling_flax_pytorch_utils.py adds an optional `dtype` argument
  to convert_pytorch_state_dict_to_flax so weights are cast to the target
  dtype during jnp.asarray, rather than always loading as float32 and re-casting
  afterward in the caller.

  These tests exercise the behavior we want to preserve and verify:
  1. The default (no dtype arg) still produces float32 arrays.
  2. Passing dtype=jnp.bfloat16 produces bfloat16 arrays with correct values.
  3. A tree_map cast and the inline dtype arg produce the same leaf values.

  Note: the full convert_pytorch_state_dict_to_flax pipeline requires a model
  with an init_weights() method. We test the core property directly using a
  numpy array and jnp.asarray, which is exactly line 133 of
  modeling_flax_pytorch_utils.py, and is what the fix changes.
  """

  def test_jnp_asarray_default_is_float32(self):
    """jnp.asarray on a float32 numpy array must produce float32 by default."""
    arr = np.array([1.0, 2.0, 3.0], dtype=np.float32)
    result = jnp.asarray(arr)
    assert result.dtype == jnp.float32, f"expected float32, got {result.dtype}"

  def test_jnp_asarray_with_dtype_produces_correct_dtype(self):
    """jnp.asarray with dtype=bfloat16 must produce bfloat16."""
    arr = np.array([1.0, 2.0, 3.0], dtype=np.float32)
    result = jnp.asarray(arr, dtype=jnp.bfloat16)
    assert result.dtype == jnp.bfloat16, f"expected bfloat16, got {result.dtype}"

  def test_inline_dtype_and_tree_map_cast_agree(self):
    """
    The current pattern is: load as float32, then tree_map cast.
    The new pattern is: cast inline during jnp.asarray.
    Both must produce identical values.
    """
    rng = np.random.default_rng(0)
    tensors = [rng.standard_normal((32, 64)).astype(np.float32) for _ in range(5)]

    # Current pattern: load float32 then cast with tree_map
    loaded_f32 = [jnp.asarray(t) for t in tensors]
    cast_after = [x.astype(jnp.bfloat16) for x in loaded_f32]

    # New pattern: cast inline
    cast_inline = [jnp.asarray(t, dtype=jnp.bfloat16) for t in tensors]

    for after, inline in zip(cast_after, cast_inline):
      assert after.dtype == jnp.bfloat16
      assert inline.dtype == jnp.bfloat16
      assert jnp.array_equal(after, inline), "inline and tree_map casts must agree"

  def test_param_tree_dtype_after_fix(self):
    """
    Simulate the full _load_params call pattern after the fix:
    convert_pytorch_state_dict_to_flax is called with dtype, so the caller's
    tree_map re-cast becomes a no-op (idempotent).
    """
    rng = np.random.default_rng(0)
    # Simulate a param tree as a dict of jnp arrays (what the fixed function returns)
    params = {
      "kernel": jnp.asarray(rng.standard_normal((64, 32)).astype(np.float32), dtype=jnp.bfloat16),
      "bias": jnp.asarray(rng.standard_normal((32,)).astype(np.float32), dtype=jnp.bfloat16),
    }
    # The caller still does tree_map(..., dtype=bfloat16) — must be idempotent
    recasted = jax.tree_util.tree_map(lambda x: jnp.asarray(x, dtype=jnp.bfloat16), params)
    for key in params:
      assert jnp.array_equal(params[key], recasted[key]), f"re-cast changed values for {key}"


# ---------------------------------------------------------------------------
# 4. nn.scan block conversion — FlaxDownBlock2D and FlaxCrossAttnDownBlock2D
# ---------------------------------------------------------------------------

class TestBlockNnScan:
  """
  Pins the forward-pass behavior of FlaxDownBlock2D and FlaxCrossAttnDownBlock2D
  before the nn.scan refactor. After the refactor the outputs must be bit-exact
  (same floating-point computation, only the parameter storage layout changes).

  These tests use channels where in_channels == out_channels so that all
  resnet layers within a block have the same architecture (no conv_shortcut
  in only the first layer). That is the precondition for nn.scan to be
  applicable without a "split first iteration" special case.
  """

  def _rng(self, seed: int = 0):
    return jax.random.PRNGKey(seed)

  def test_down_block_2d_output_shape(self):
    from local_geometry.model_overviews.flax_sd_turbo.hf_copy.unet_2d_blocks_flax import FlaxDownBlock2D

    channels = 64
    num_layers = 2
    batch, h, w = 1, 8, 8
    temb_dim = 256

    model = FlaxDownBlock2D(
      in_channels=channels,
      out_channels=channels,
      num_layers=num_layers,
      add_downsample=False,
      dtype=jnp.float32,
    )
    x = jax.random.normal(self._rng(0), (batch, h, w, channels))
    temb = jax.random.normal(self._rng(1), (batch, temb_dim))
    params = model.init(self._rng(2), x, temb)
    out, res = model.apply(params, x, temb)
    assert out.shape == (batch, h, w, channels)
    assert len(res) == num_layers

  def test_down_block_2d_deterministic(self):
    """Same input and params must always produce the same output."""
    from local_geometry.model_overviews.flax_sd_turbo.hf_copy.unet_2d_blocks_flax import FlaxDownBlock2D

    channels = 32
    batch, h, w = 1, 8, 8
    temb_dim = 256

    model = FlaxDownBlock2D(
      in_channels=channels,
      out_channels=channels,
      num_layers=2,
      add_downsample=False,
      dtype=jnp.float32,
    )
    x = jax.random.normal(self._rng(0), (batch, h, w, channels))
    temb = jax.random.normal(self._rng(1), (batch, temb_dim))
    params = model.init(self._rng(2), x, temb)

    out_a, _ = model.apply(params, x, temb)
    out_b, _ = model.apply(params, x, temb)
    assert jnp.array_equal(out_a, out_b), "forward pass must be deterministic"

  def test_cross_attn_down_block_output_shape(self):
    from local_geometry.model_overviews.flax_sd_turbo.hf_copy.unet_2d_blocks_flax import FlaxCrossAttnDownBlock2D

    channels = 64
    n_heads = 4
    num_layers = 2
    batch, h, w = 1, 8, 8
    temb_dim = 256
    ctx_len, ctx_dim = 5, 64

    model = FlaxCrossAttnDownBlock2D(
      in_channels=channels,
      out_channels=channels,
      num_layers=num_layers,
      num_attention_heads=n_heads,
      add_downsample=False,
      dtype=jnp.float32,
    )
    x = jax.random.normal(self._rng(0), (batch, h, w, channels))
    temb = jax.random.normal(self._rng(1), (batch, temb_dim))
    ctx = jax.random.normal(self._rng(2), (batch, ctx_len, ctx_dim))
    params = model.init(self._rng(3), x, temb, ctx)
    out, res = model.apply(params, x, temb, ctx)
    assert out.shape == (batch, h, w, channels)
    assert len(res) == num_layers

  def test_cross_attn_down_block_deterministic(self):
    """Same input and params must always produce the same output."""
    from local_geometry.model_overviews.flax_sd_turbo.hf_copy.unet_2d_blocks_flax import FlaxCrossAttnDownBlock2D

    channels = 64
    n_heads = 4
    batch, h, w = 1, 8, 8
    temb_dim = 256
    ctx_len, ctx_dim = 5, 64

    model = FlaxCrossAttnDownBlock2D(
      in_channels=channels,
      out_channels=channels,
      num_layers=2,
      num_attention_heads=n_heads,
      add_downsample=False,
      dtype=jnp.float32,
    )
    x = jax.random.normal(self._rng(0), (batch, h, w, channels))
    temb = jax.random.normal(self._rng(1), (batch, temb_dim))
    ctx = jax.random.normal(self._rng(2), (batch, ctx_len, ctx_dim))
    params = model.init(self._rng(3), x, temb, ctx)

    out_a, _ = model.apply(params, x, temb, ctx)
    out_b, _ = model.apply(params, x, temb, ctx)
    assert jnp.array_equal(out_a, out_b), "forward pass must be deterministic"
