import mlx.core as mx
import numpy as np
import pytest

from h3_mlx.ops import modulate, rms_norm, rotate
from h3_mlx.tuning import (
    TunedFC2,
    _gemm,
    norm_modulate_exact,
    qk_norm_rope_exact,
    supported_device,
)
from h3_mlx.weights import Linear


@pytest.mark.parametrize("magnitude", [0.001, 1.0, 10000.0])
@pytest.mark.parametrize("rope_dtype", [mx.bfloat16, mx.float32])
@pytest.mark.skipif(
    not supported_device(), reason="Pinned M5 Pro/MLX kernel validation"
)
def test_qk_fusion_preserves_native_reduction_and_rope_bits(magnitude, rope_dtype):
    mx.random.seed(8192)
    qkv = (mx.random.normal((37, 3, 3, 128)) * magnitude).astype(mx.bfloat16)
    qw = mx.random.normal((128,)).astype(mx.bfloat16)
    kw = mx.random.normal((128,)).astype(mx.bfloat16)
    angles = mx.random.normal((1, 37, 1, 48))
    angles = mx.concatenate((angles, angles), axis=-1)
    cosine = mx.cos(angles).astype(rope_dtype)
    sine = mx.sin(angles).astype(rope_dtype)
    reference = tuple(
        rotate(rms_norm(qkv[None, :, i], w), cosine, sine)
        for i, w in ((0, qw), (1, kw))
    )
    result = qk_norm_rope_exact(qkv, qw, kw, cosine, sine)
    for expected, actual in zip(reference, result, strict=True):
        np.testing.assert_array_equal(
            np.asarray(actual.view(mx.uint16)), np.asarray(expected.view(mx.uint16))
        )


@pytest.mark.skipif(not supported_device(), reason="Metal 4 M5 kernel validation")
def test_mpp_fp32_bias_epilogue_and_partial_m_tile():
    mx.random.seed(4352)
    x = (mx.random.normal((71, 512)) * 32).astype(mx.bfloat16)
    weight = (mx.random.normal((256, 512)) * 0.02).astype(mx.bfloat16)
    bias = mx.random.normal((256,)).astype(mx.bfloat16)
    expected = mx.addmm(bias, x, weight.T)
    actual = _gemm(x, weight, bias, 64, 128, 128, 4, 8)
    np.testing.assert_array_equal(
        np.asarray(actual.view(mx.uint16)), np.asarray(expected.view(mx.uint16))
    )


def test_tuned_fc2_keeps_native_fallback_for_unsupported_shapes():
    mx.random.seed(71)
    layer = Linear(mx.random.normal((32, 64)).astype(mx.bfloat16))
    x = mx.random.normal((17, 64)).astype(mx.bfloat16)
    tuned = TunedFC2(layer)
    np.testing.assert_array_equal(
        np.asarray(tuned(x).view(mx.uint16)), np.asarray(layer(x).view(mx.uint16))
    )


@pytest.mark.skipif(not supported_device(), reason="Pinned M5 Pro reduction layout")
def test_full_width_norm_fusion_matches_native_loop_and_bf16_boundaries():
    mx.random.seed(5376)
    x = (mx.random.normal((13, 5376)) * 256).astype(mx.bfloat16)
    weight = mx.random.normal((5376,)).astype(mx.bfloat16)
    shift = mx.random.normal((6, 5376)).astype(mx.bfloat16)
    scale = (mx.random.normal((6, 5376)) * 0.1).astype(mx.bfloat16)
    indices = mx.arange(13, dtype=mx.int32) % 6
    expected = modulate(x, weight, shift, scale, indices)
    actual = norm_modulate_exact(x, weight, shift, scale, indices)
    np.testing.assert_array_equal(
        np.asarray(actual.view(mx.uint16)), np.asarray(expected.view(mx.uint16))
    )