"""Flow Euler sampler tests — schedule, integration, and guidance call counts."""
import sys
from pathlib import Path
import mlx.core as mx
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from pixal3d_mlx.sampler import FlowEulerSampler  # noqa: E402


def test_schedule_matches_upstream():
    s = FlowEulerSampler()
    worst = 0.0
    for r in (1.0, 3.0, 0.5):
        got = s.timesteps(7, r)
        seq = np.linspace(1, 0, 8)
        seq = r * seq / (1 + (r - 1) * seq)
        want = [(seq[i], seq[i + 1]) for i in range(7)]
        worst = max(worst, max(abs(a - c) + abs(b - d) for (a, b), (c, d) in zip(got, want)))
    assert worst < 1e-12, f"schedule err {worst:.3g}"
    return worst


def test_constant_velocity_integrates_exactly():
    """v constant => x_end = x_start - v*(1-0), regardless of step count."""
    s = FlowEulerSampler()
    class Const:
        def __call__(self, x, t, cond):
            return mx.ones_like(x) * 2.0
    for steps in (1, 10, 37):
        out = np.asarray(s.sample(Const(), mx.zeros((1, 4)), cond=None, steps=steps))
        err = abs(out[0, 0] - (-2.0))
        assert err < 1e-5, f"steps={steps} gave {out[0,0]}, want -2.0"
    return 0.0


def test_guidance_interval_skips_negative_pass():
    """Outside the interval only the conditional branch runs — half the model calls."""
    s = FlowEulerSampler()
    calls = []
    class Counting:
        def __call__(self, x, t, cond):
            calls.append(float(t[0]) / 1000.0)
            return mx.zeros_like(x)
    s.sample(Counting(), mx.zeros((1, 4)), cond="p", neg_cond="n", steps=4,
             guidance_strength=3.0, guidance_interval=(0.0, 0.5))
    assert len(calls) == 6, f"expected 6 calls (2 in-window + 1 out), got {len(calls)}"
    from collections import Counter
    c = Counter(round(t, 3) for t in calls)
    assert c[1.0] == 1 and c[0.5] == 2, dict(c)
    return 0.0


def test_guidance_is_a_lerp():
    """upstream uses g*pos + (1-g)*neg, not neg + g*(pos-neg)."""
    s = FlowEulerSampler()
    class Two:
        def __call__(self, x, t, cond):
            return mx.ones_like(x) * (1.0 if cond == "p" else 3.0)
    out = np.asarray(s.sample(Two(), mx.zeros((1, 1)), cond="p", neg_cond="n",
                              steps=1, guidance_strength=2.0))
    # g*1 + (1-g)*3 = 2 - 3 = -1  -> x = 0 - 1*(-1) = +1
    assert abs(out[0, 0] - 1.0) < 1e-5, f"got {out[0,0]}, want +1.0 for a LERP"
    return 0.0


if __name__ == "__main__":
    tests = [("schedule vs upstream", test_schedule_matches_upstream),
             ("constant v integrates", test_constant_velocity_integrates_exactly),
             ("guidance interval", test_guidance_interval_skips_negative_pass),
             ("guidance is a lerp", test_guidance_is_a_lerp)]
    failed = 0
    for n, fn in tests:
        try:
            fn(); print(f"  PASS  {n}")
        except AssertionError as e:
            print(f"  FAIL  {n}: {e}"); failed += 1
    print(f"\n{len(tests)-failed}/{len(tests)} passed")
    sys.exit(1 if failed else 0)