"""Unit tests for the _waterfill helper.

Tests cover:
1. The known repro case where all sleeves clamp and residual mass is silently discarded.
2. A normal case where the algorithm converges without all-clamped iteration.
3. A seeded stress test over ~2000 random feasible inputs.
"""
import math
import random

from fxhnt.application.paper_book import _waterfill


def _check_result(result: dict, w_floor: float, cap: float) -> None:
    """Assert sum-to-1 and bounds for a _waterfill result."""
    total = sum(result.values())
    assert math.isclose(total, 1.0, abs_tol=1e-9), (
        f"sum={total} != 1.0 (off by {abs(total - 1.0):.2e}); result={result}"
    )
    for k, v in result.items():
        assert v >= w_floor - 1e-9, f"weight[{k}]={v} < w_floor={w_floor}"
        assert v <= cap + 1e-9, f"weight[{k}]={v} > cap={cap}"


def test_repro_all_clamped_sums_to_one():
    """The bug: when every sleeve hits a bound, residual 0.0813 was silently dropped."""
    result = _waterfill(
        {'s0': 0.198, 's1': 0.359, 's2': 0.108, 's3': 0.334},
        w_floor=0.1372,
        cap=0.2605,
    )
    _check_result(result, w_floor=0.1372, cap=0.2605)


def test_normal_case_sums_to_one_and_respects_bounds():
    """Normal case: one sleeve dominates, floor rescues the other."""
    result = _waterfill({"a": 0.9, "b": 0.1}, 0.25, 1.0)
    _check_result(result, w_floor=0.25, cap=1.0)
    # "a" should be capped at 0.75 (1 - 0.25) and "b" gets the rest
    assert result["b"] >= 0.25 - 1e-9


def test_stress_random_feasible_inputs():
    """~2000 random feasible inputs: sum must be 1.0 and weights in [w_floor, cap]."""
    rng = random.Random(0)
    n_cases = 0
    for _ in range(2000):
        n = rng.randint(2, 6)
        w_floor = 0.5 / n
        # cap must satisfy cap >= 1/n (feasibility) and cap*n >= 1
        cap = rng.uniform(1.0 / n, 1.0)
        # random positive weights normalised to 1
        raw = [rng.random() for _ in range(n)]
        total = sum(raw)
        init = {f"s{i}": v / total for i, v in enumerate(raw)}
        result = _waterfill(init, w_floor, cap)
        _check_result(result, w_floor=w_floor, cap=cap)
        n_cases += 1
    assert n_cases == 2000
