"""Tests for core/sampler.py -- flow-matching ODE/SDE sampler."""

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

from mlx_music.core.sampler import (
    _TURBO_SCHEDULES,
    VALID_SHIFTS,
    VALID_TIMESTEPS,
    apply_cfg,
    euler_step,
    get_x0_from_noise,
    prepare_noise,
    sample_loop,
    sde_step,
    time_schedule,
)


class TestTimeSchedule:
    """Verify timestep schedules match the reference implementation."""

    def test_shift3_steps8(self):
        schedule = time_schedule(8, shift=3.0)
        expected = mx.array(_TURBO_SCHEDULES[3.0])
        assert mx.allclose(schedule, expected, atol=1e-6).item()

    def test_shift1_steps8(self):
        schedule = time_schedule(8, shift=1.0)
        expected = mx.array(_TURBO_SCHEDULES[1.0])
        assert mx.allclose(schedule, expected, atol=1e-6).item()

    def test_shift2_steps8(self):
        schedule = time_schedule(8, shift=2.0)
        expected = mx.array(_TURBO_SCHEDULES[2.0])
        assert mx.allclose(schedule, expected, atol=1e-6).item()

    def test_length_matches_num_steps(self):
        for n in [4, 8, 16, 32]:
            schedule = time_schedule(n, shift=3.0)
            assert schedule.shape == (n,), f"Expected ({n},), got {schedule.shape}"

    def test_monotonically_decreasing(self):
        """Schedule should go from high t (noisy) to low t (clean)."""
        for shift in [1.0, 2.0, 3.0]:
            schedule = time_schedule(8, shift=shift)
            vals = schedule.tolist()
            for i in range(len(vals) - 1):
                assert vals[i] > vals[i + 1], (
                    f"shift={shift}: schedule not decreasing at index {i}: "
                    f"{vals[i]} vs {vals[i+1]}"
                )

    def test_starts_at_one(self):
        for shift in [1.0, 2.0, 3.0]:
            schedule = time_schedule(8, shift=shift)
            assert abs(schedule[0].item() - 1.0) < 1e-6

    def test_all_positive(self):
        """All timesteps should be > 0 (we never include the final 0)."""
        for shift in [1.0, 2.0, 3.0]:
            schedule = time_schedule(8, shift=shift)
            assert mx.all(schedule > 0).item()

    def test_shift1_is_linear(self):
        """With shift=1, the schedule is a simple linear ramp."""
        schedule = time_schedule(8, shift=1.0)
        expected = mx.linspace(1.0, 0.0, 9)[:-1]
        assert mx.allclose(schedule, expected, atol=1e-6).item()

    def test_arbitrary_steps_shift3(self):
        """Non-standard step count should still use the general formula."""
        schedule = time_schedule(4, shift=3.0)
        assert schedule.shape == (4,)
        assert abs(schedule[0].item() - 1.0) < 1e-6
        # Should be decreasing
        vals = schedule.tolist()
        for i in range(len(vals) - 1):
            assert vals[i] > vals[i + 1]

    def test_higher_shift_more_frontloaded(self):
        """Higher shift should keep timesteps higher (more time at high noise)."""
        s1 = time_schedule(8, shift=1.0)
        s3 = time_schedule(8, shift=3.0)
        # At the midpoint (index 4), shift=3 should have a higher value
        assert s3[4].item() > s1[4].item()


class TestPrepareNoise:
    def test_shape(self):
        noise = prepare_noise((2, 100, 64))
        assert noise.shape == (2, 100, 64)

    def test_dtype_default(self):
        noise = prepare_noise((1, 10, 8))
        assert noise.dtype == mx.float32

    def test_dtype_float16(self):
        noise = prepare_noise((1, 10, 8), dtype=mx.float16)
        assert noise.dtype == mx.float16

    def test_seeded_reproducibility(self):
        """Same seed should produce the same noise."""
        a = prepare_noise((2, 50, 64), seed=42)
        b = prepare_noise((2, 50, 64), seed=42)
        assert mx.array_equal(a, b).item()

    def test_different_seeds_differ(self):
        """Different seeds should produce different noise."""
        a = prepare_noise((2, 50, 64), seed=1)
        b = prepare_noise((2, 50, 64), seed=2)
        assert not mx.array_equal(a, b).item()

    def test_approximately_standard_normal(self):
        """Large sample should have ~0 mean and ~1 std."""
        noise = prepare_noise((1000, 100, 64), seed=123)
        mean = mx.mean(noise).item()
        std = mx.sqrt(mx.mean(noise * noise)).item()  # RMS approx std when mean~0
        assert abs(mean) < 0.05, f"Mean too far from 0: {mean}"
        assert abs(std - 1.0) < 0.05, f"Std too far from 1: {std}"


class TestEulerStep:
    def test_identity_when_velocity_zero(self):
        """Zero velocity should leave xt unchanged."""
        xt = mx.ones((1, 10, 4))
        vt = mx.zeros((1, 10, 4))
        result = euler_step(xt, vt, 1.0, 0.5)
        assert mx.array_equal(result, xt).item()

    def test_basic_step(self):
        """xt_next = xt - vt * dt."""
        xt = mx.array([[[2.0, 2.0]]])
        vt = mx.array([[[1.0, 1.0]]])
        result = euler_step(xt, vt, t_curr=1.0, t_next=0.5)
        expected = mx.array([[[1.5, 1.5]]])  # 2 - 1 * 0.5
        assert mx.allclose(result, expected, atol=1e-6).item()

    def test_full_step_recovers_x0(self):
        """If v = noise - data and we step from t=1 to t=0, we get data."""
        data = mx.array([[[3.0, 4.0]]])
        noise = mx.array([[[1.0, 2.0]]])
        # At t=1, xt = 1*noise + 0*data = noise
        xt = noise
        vt = noise - data  # velocity = noise - x0
        result = euler_step(xt, vt, t_curr=1.0, t_next=0.0)
        # xt - vt * 1.0 = noise - (noise - data) = data
        assert mx.allclose(result, data, atol=1e-6).item()


class TestGetX0FromNoise:
    def test_at_t1_recovers_x0(self):
        """At t=1, x0 = zt - vt * 1 = zt - vt."""
        zt = mx.array([[[5.0, 6.0]]])
        vt = mx.array([[[2.0, 3.0]]])
        t = mx.array([1.0])
        x0 = get_x0_from_noise(zt, vt, t)
        expected = mx.array([[[3.0, 3.0]]])
        assert mx.allclose(x0, expected, atol=1e-6).item()

    def test_at_t0_returns_zt(self):
        """At t=0, x0 = zt - vt * 0 = zt (already clean)."""
        zt = mx.array([[[5.0, 6.0]]])
        vt = mx.array([[[99.0, 99.0]]])  # should be irrelevant
        t = mx.array([0.0])
        x0 = get_x0_from_noise(zt, vt, t)
        assert mx.allclose(x0, zt, atol=1e-6).item()

    def test_batch_dimension(self):
        """Should handle batch dim correctly."""
        zt = mx.ones((3, 10, 4))
        vt = mx.ones((3, 10, 4)) * 0.5
        t = mx.array([0.8, 0.5, 0.2])
        x0 = get_x0_from_noise(zt, vt, t)
        assert x0.shape == (3, 10, 4)
        # For batch 0: 1.0 - 0.5 * 0.8 = 0.6
        assert abs(x0[0, 0, 0].item() - 0.6) < 1e-5
        # For batch 1: 1.0 - 0.5 * 0.5 = 0.75
        assert abs(x0[1, 0, 0].item() - 0.75) < 1e-5

    def test_flow_matching_identity(self):
        """Verify: if xt = t*noise + (1-t)*x0 and v = noise - x0,
        then get_x0_from_noise(xt, v, t) == x0."""
        x0_true = mx.array([[[1.0, 2.0, 3.0]]])
        noise = mx.array([[[10.0, 20.0, 30.0]]])
        t_val = 0.7
        # Construct xt at this timestep
        xt = t_val * noise + (1.0 - t_val) * x0_true
        v = noise - x0_true
        t = mx.array([t_val])
        x0_recovered = get_x0_from_noise(xt, v, t)
        assert mx.allclose(x0_recovered, x0_true, atol=1e-5).item()


class TestSdeStep:
    def test_output_shape(self):
        xt = mx.ones((2, 10, 4))
        vt = mx.zeros((2, 10, 4))
        result = sde_step(xt, vt, t_curr=1.0, t_next=0.5)
        assert result.shape == (2, 10, 4)

    def test_no_nan(self):
        """SDE step should not produce NaN values."""
        mx.random.seed(42)
        xt = mx.random.normal((2, 50, 64))
        vt = mx.random.normal((2, 50, 64))
        result = sde_step(xt, vt, t_curr=0.8, t_next=0.5)
        mx.eval(result)
        assert not mx.any(mx.isnan(result)).item()

    def test_t_next_zero_returns_x0(self):
        """When t_next=0, re-noising adds 0*noise + 1*x0 = x0."""
        xt = mx.array([[[5.0, 6.0]]])
        vt = mx.array([[[2.0, 3.0]]])
        result = sde_step(xt, vt, t_curr=1.0, t_next=0.0)
        # x0 = xt - vt * 1.0 = [[3, 3]]
        # x_next = 0*noise + 1*x0 = x0
        expected = mx.array([[[3.0, 3.0]]])
        assert mx.allclose(result, expected, atol=1e-6).item()


class TestApplyCfg:
    def test_scale_one_returns_cond(self):
        """scale=1 should return pred_cond exactly."""
        cond = mx.array([1.0, 2.0, 3.0])
        uncond = mx.array([0.0, 0.0, 0.0])
        result = apply_cfg(cond, uncond, scale=1.0)
        assert mx.allclose(result, cond, atol=1e-6).item()

    def test_scale_zero_returns_uncond(self):
        """scale=0 should return pred_uncond exactly."""
        cond = mx.array([10.0, 20.0])
        uncond = mx.array([1.0, 2.0])
        result = apply_cfg(cond, uncond, scale=0.0)
        assert mx.allclose(result, uncond, atol=1e-6).item()

    def test_amplification(self):
        """scale=2 should amplify the difference."""
        cond = mx.array([2.0])
        uncond = mx.array([1.0])
        # 1 + 2*(2-1) = 3
        result = apply_cfg(cond, uncond, scale=2.0)
        assert abs(result.item() - 3.0) < 1e-6

    def test_negative_scale(self):
        """Negative scale should invert the guidance direction."""
        cond = mx.array([2.0])
        uncond = mx.array([1.0])
        # 1 + (-1)*(2-1) = 0
        result = apply_cfg(cond, uncond, scale=-1.0)
        assert abs(result.item() - 0.0) < 1e-6


class TestSampleLoop:
    """Integration tests for the main denoising loop."""

    @staticmethod
    def _constant_denoise_fn(xt: mx.array, t: float) -> mx.array:
        """Dummy denoise function: always predicts zero velocity."""
        return mx.zeros_like(xt)

    @staticmethod
    def _identity_denoise_fn(xt: mx.array, t: float) -> mx.array:
        """Denoise function that predicts v = xt (moves toward origin)."""
        return xt

    def test_ode_constant_velocity_preserves_input(self):
        """Zero velocity should leave noise unchanged (except final x0 step)."""
        mx.random.seed(99)
        noise = mx.random.normal((1, 10, 4))
        schedule = time_schedule(4, shift=1.0)
        result = sample_loop(self._constant_denoise_fn, noise, schedule, method="ode")
        # With zero velocity: euler steps do nothing, final step is x0 = xt - 0*t = xt
        assert mx.allclose(result, noise, atol=1e-5).item()

    def test_sde_produces_output(self):
        """SDE method should produce valid (non-NaN) output."""
        mx.random.seed(42)
        noise = mx.random.normal((1, 10, 4))
        schedule = time_schedule(4, shift=1.0)
        result = sample_loop(self._constant_denoise_fn, noise, schedule, method="sde")
        mx.eval(result)
        assert not mx.any(mx.isnan(result)).item()
        assert result.shape == (1, 10, 4)

    def test_invalid_method_raises(self):
        noise = mx.zeros((1, 4, 2))
        schedule = time_schedule(4)
        with pytest.raises(ValueError, match="Unknown sampling method"):
            sample_loop(self._constant_denoise_fn, noise, schedule, method="rk4")

    def test_cfg_requires_uncond_fn(self):
        noise = mx.zeros((1, 4, 2))
        schedule = time_schedule(4)
        with pytest.raises(ValueError, match="cfg_scale requires"):
            sample_loop(self._constant_denoise_fn, noise, schedule, cfg_scale=7.5)

    def test_cfg_blending(self):
        """CFG should blend conditional and unconditional predictions."""
        mx.random.seed(10)
        noise = mx.ones((1, 4, 2))
        schedule = time_schedule(2, shift=1.0)

        def cond_fn(xt, t):
            return mx.ones_like(xt)

        def uncond_fn(xt, t):
            return mx.zeros_like(xt)

        result = sample_loop(
            cond_fn,
            noise,
            schedule,
            method="ode",
            cfg_scale=2.0,
            denoise_fn_uncond=uncond_fn,
        )
        assert result.shape == (1, 4, 2)
        assert not mx.any(mx.isnan(result)).item()

    def test_callback_called(self):
        """Callback should be called once per step."""
        noise = mx.zeros((1, 4, 2))
        schedule = time_schedule(4, shift=1.0)
        call_log = []

        def cb(step, total, xt):
            call_log.append((step, total))

        sample_loop(
            self._constant_denoise_fn,
            noise,
            schedule,
            method="ode",
            callback=cb,
        )
        assert len(call_log) == 4
        assert call_log[0] == (0, 4)
        assert call_log[-1] == (3, 4)

    def test_output_shape_matches_input(self):
        """Output shape should match the input noise shape."""
        shapes = [(1, 10, 64), (2, 50, 32), (4, 100, 8)]
        schedule = time_schedule(4, shift=1.0)
        for shape in shapes:
            noise = mx.zeros(shape)
            result = sample_loop(
                self._constant_denoise_fn, noise, schedule, method="ode"
            )
            assert result.shape == shape, f"Expected {shape}, got {result.shape}"

    def test_ode_convergence_with_known_target(self):
        """With a perfect velocity predictor, ODE should recover x0.

        Set up: x0 = 0 (zeros), noise is random, v = noise - x0 = noise.
        The model always returns v = noise (the true velocity for x0=0).
        After denoising, we should get close to zero.
        """
        mx.random.seed(77)
        noise = mx.random.normal((1, 20, 8))
        schedule = time_schedule(8, shift=1.0)

        # The model "knows" the true velocity is the noise itself (since x0=0)
        captured_noise = noise

        def perfect_velocity_fn(xt, t):
            # v = noise - x0, with x0 = 0, so v = noise
            # But actually: at any time t, if data = 0, then xt = t * noise
            # and v should be noise (constant). So v = noise works for all t.
            return captured_noise

        result = sample_loop(perfect_velocity_fn, noise, schedule, method="ode")
        # Should be close to zero (the target x0)
        max_abs = mx.max(mx.abs(result)).item()
        assert max_abs < 1e-4, f"Expected near-zero, got max_abs={max_abs}"
