import json

import mlx.core as mx
import numpy as np
import pytest
import torch
import torch.nn.functional as F

from h3_mlx.config import canvas_size, load_sigmas
from h3_mlx.layout import image_positions, patchify, unpatchify
from h3_mlx.ops import attention, layer_norm, modulate, rms_norm, rotate, silu
from h3_mlx.weights import Linear, WeightStore


def array(t):
    return mx.array(t.float().numpy()).astype(
        mx.bfloat16 if t.dtype == torch.bfloat16 else mx.float32
    )


def numpy(x):
    return np.asarray(x.astype(mx.float32))


def assert_bf16_near(actual, expected, *, equality=0.9999):
    a, b = numpy(actual), expected.float().numpy()
    # CPU/Metal reduction orders may cross a BF16 rounding midpoint. Require
    # one ULP at most, while detecting the much larger wrong-arithmetic case.
    aa = np.asarray(actual.view(mx.uint16)).astype(np.int32)
    bb = expected.view(torch.int16).numpy().view(np.uint16).astype(np.int32)
    assert np.max(abs(aa - bb)) <= 1
    assert np.mean(a == b) >= equality


@pytest.mark.parametrize("width", [128, 5376])
def test_distinct_aten_and_qwen_rms_rounding(width):
    gen = torch.Generator().manual_seed(9)
    x = torch.randn(17, width, generator=gen).bfloat16()
    w = torch.randn(width, generator=gen).bfloat16()
    aten = F.rms_norm(x, (width,), w, 1e-5)
    qwen = (
        x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-5)
    ).bfloat16() * w
    assert_bf16_near(rms_norm(array(x), array(w)), aten)
    assert_bf16_near(rms_norm(array(x), array(w), rounded=True), qwen)
    assert not torch.equal(aten, qwen)


def test_bf16_bias_epilogue_and_silu():
    torch.manual_seed(21)
    x, w, b = (
        torch.randn(13, 128).bfloat16(),
        torch.randn(256, 128).bfloat16(),
        torch.randn(256).bfloat16(),
    )
    layer = Linear(array(w), array(b))
    assert_bf16_near(layer(array(x)), F.linear(x, w, b), equality=0.998)
    np.testing.assert_array_equal(numpy(silu(array(x))), F.silu(x).float().numpy())


def test_attention_dense_gqa_and_causal_chunks():
    torch.manual_seed(8)
    q, k, v = (
        torch.randn(1, 37, 8, 128),
        torch.randn(1, 37, 2, 128),
        torch.randn(1, 37, 2, 128),
    )
    ref = F.scaled_dot_product_attention(
        q.transpose(1, 2),
        k.transpose(1, 2),
        v.transpose(1, 2),
        is_causal=True,
        enable_gqa=True,
    ).transpose(1, 2)
    for chunk in (0, 11):
        out = attention(array(q), array(k), array(v), causal=True, query_chunk=chunk)
        np.testing.assert_allclose(numpy(out), ref.numpy(), atol=2e-6, rtol=2e-5)


def test_rotary_keeps_passthrough_and_bf16_products():
    torch.manual_seed(41)
    x = torch.randn(1, 17, 3, 128).bfloat16()
    angles = torch.randn(1, 17, 1, 48)
    angles = torch.cat((angles, angles), dim=-1)
    cos, sin = angles.cos().bfloat16(), angles.sin().bfloat16()
    a = x[..., :96]
    ref = torch.cat(
        (a * cos + torch.cat((-a[..., 48:], a[..., :48]), -1) * sin, x[..., 96:]), -1
    )
    np.testing.assert_array_equal(
        numpy(rotate(array(x), array(cos), array(sin))), ref.float().numpy()
    )


def test_layer_norm_affine_does_not_round_normalized_input():
    torch.manual_seed(98)
    x = torch.randn(17, 1024).bfloat16()
    w, b = torch.randn(1024).bfloat16(), torch.randn(1024).bfloat16()
    ref = F.layer_norm(x, (1024,), w, b, 1e-5)
    assert_bf16_near(layer_norm(array(x), array(w), array(b)), ref, equality=0.999)


def test_safetensors_bf16_bit_loading_and_grouped_qkv(tmp_path):
    from safetensors.torch import save_file

    tensor = torch.arange(2 * 3 * 4 * 8, dtype=torch.float32).reshape(24, 8).bfloat16()
    save_file({"qkv.weight": tensor}, str(tmp_path / "test.safetensors"))
    store = WeightStore(tmp_path)
    np.testing.assert_array_equal(
        numpy(store.get("qkv.weight")), tensor.float().numpy()
    )
    layer = Linear.load(store, "qkv", grouped_heads=(2, 4))
    ref = tensor.reshape(2, 3, 4, 8).permute(1, 0, 2, 3).reshape(24, 8)
    np.testing.assert_array_equal(numpy(layer.weight), ref.float().numpy())
    store.clear()


def test_patch_order_matches_native_einsum():
    x = torch.arange(24 * 8 * 10, dtype=torch.float32).reshape(1, 24, 8, 10)
    ref = torch.einsum("nchpwq->nhwcpq", x.reshape(1, 24, 4, 2, 5, 2)).reshape(-1, 96)
    rows = patchify(array(x))
    np.testing.assert_array_equal(numpy(rows), ref.numpy())
    np.testing.assert_array_equal(numpy(unpatchify(rows, 8, 10)), x.numpy())


def test_positions_are_native_ref2va_grid():
    out = numpy(image_positions(7, 8, 12, [(4, 6)]))
    assert out.shape == (7 + 6 + 24, 3)
    np.testing.assert_array_equal(out[:7, 0], np.arange(7))
    np.testing.assert_array_equal(out[7:13, 0], 7)
    np.testing.assert_array_equal(out[13:, 0], 8)
    np.testing.assert_allclose(out[13:, 1].min(), (1 - 8 / np.sqrt(96)) * 16, rtol=1e-6)


def test_native_schedule_and_canvas_validation(tmp_path):
    assert canvas_size(1024, 1024) == (2048, 2048)
    assert canvas_size(512, 768, 1024) == (512, 768)
    for width, height in [(0, 2048), (512, 0), (-1, 32)]:
        with pytest.raises(ValueError):
            canvas_size(width, height)
    (tmp_path / "pdd_config.json").write_text(json.dumps({"sigmas": [1, 0.5, 0]}))
    with pytest.raises(ValueError, match="six-step"):
        load_sigmas(tmp_path)


@pytest.mark.parametrize("threads", [128, 256, 512])
def test_metal_modulation_with_odd_row_count(threads):
    from h3_mlx.kernels import norm_modulate

    mx.random.seed(71)
    x = mx.random.normal((65, 5376)).astype(mx.bfloat16)
    w = mx.random.normal((5376,)).astype(mx.bfloat16)
    s = (mx.random.normal((3, 5376)) * 0.1).astype(mx.bfloat16)
    ids = mx.arange(65) % 3
    ref = modulate(x, w, s, s, ids)
    out = norm_modulate(x, w, s, s, ids, threads=threads)
    a, b = numpy(ref), numpy(out)
    # Different parallel reduction orders are allowed only within one BF16 ULP.
    assert np.sqrt(np.mean((a - b) ** 2) / (np.mean(a * a) + 1e-20)) < 1e-4
    assert np.mean(a == b) > 0.9999


def test_fused_qk_norm_rope_layout_and_rounding():
    from h3_mlx.kernels import qk_norm_rope

    mx.random.seed(33)
    qkv = mx.random.normal((67, 3, 7, 128)).astype(mx.bfloat16)
    w = mx.random.normal((128,)).astype(mx.bfloat16)
    angles = mx.random.normal((1, 67, 1, 48))
    angles = mx.concatenate((angles, angles), axis=-1)
    cos, sin = mx.cos(angles).astype(mx.bfloat16), mx.sin(angles).astype(mx.bfloat16)
    actual = qk_norm_rope(qkv, w, w, cos, sin)
    for i, out in enumerate(actual):
        ref = rotate(rms_norm(qkv[None, :, i], w), cos, sin)
        a, b = numpy(ref), numpy(out)
        assert np.sqrt(np.mean((a - b) ** 2) / (np.mean(a * a) + 1e-20)) < 1e-4
        assert np.mean(a == b) > 0.9999