"""Property tests for rtc.masking (see its docstring).

Run: ``python -m pytest tests/`` or ``python tests/test_masking.py``.
"""
import numpy as np

from rtc.masking import _SCHEDULES, get_prefix_weights_np as W

# The reference implementation's documented example, exactly.
got = W(start=2, end=6, total=10, schedule="linear")
want = np.array([1, 1, 4/5, 3/5, 2/5, 1/5, 0, 0, 0, 0])
assert np.allclose(got, want), f"linear mismatch:\n got={got}\nwant={want}"
print(f"linear(2,6,10) = {got}  == reference example  OK")

for sched in _SCHEDULES:
    for start, end, total in [(0, 8, 16), (2, 6, 10), (4, 10, 16), (0, 0, 8),
                              (8, 8, 16), (6, 4, 16), (3, 16, 16)]:
        w = W(start, end, total, sched)
        s, e = min(start, end), min(max(end, 0), total)
        assert w.shape == (total,)
        assert np.all((w >= 0) & (w <= 1)), f"{sched} {start},{end}: out of [0,1]"
        # 1) frozen prefix is a HARD constraint
        if sched != "ones":
            assert np.all(w[:s] == 1.0), f"{sched} {start},{end},{total}: prefix not 1"
        # 2) monotone non-increasing
        assert np.all(np.diff(w) <= 1e-12), f"{sched} {start},{end}: not monotone"
        # 3) nothing beyond the overlap
        if sched != "ones" or e < total:
            assert np.all(w[e:] == 0.0), f"{sched} {start},{end},{total}: tail not 0"
print(f"properties 1-3 hold for all of {_SCHEDULES} across 7 (start,end,total) cases  OK")

# zeros == the naive hard-masking baseline: prefix only, no soft overlap
z = W(2, 6, 10, "zeros")
assert np.array_equal(z, np.array([1., 1., 0, 0, 0, 0, 0, 0, 0, 0])), z
print(f"zeros(2,6,10)  = {z}  (naive hard-mask baseline)  OK")

e = W(2, 6, 10, "exp")
print(f"exp(2,6,10)    = {np.round(e, 4)}")
assert e[0] == e[1] == 1.0 and e[6] == 0.0
assert np.all(e[2:6] <= W(2, 6, 10, "linear")[2:6] + 1e-12), "exp should not exceed linear"
print("exp endpoints + convexity vs linear  OK")