"""Tests for giant/train.py helpers."""

from giant.train import _gumbel_tau


def test_gumbel_tau_at_step_zero_is_start():
    assert _gumbel_tau(0, 1000, 1.0, 0.1) == 1.0


def test_gumbel_tau_at_total_steps_is_end():
    assert abs(_gumbel_tau(1000, 1000, 1.0, 0.1) - 0.1) < 1e-9


def test_gumbel_tau_interpolates_linearly_midway():
    assert abs(_gumbel_tau(500, 1000, 1.0, 0.1) - 0.55) < 1e-9


def test_gumbel_tau_clamps_beyond_total_steps():
    assert _gumbel_tau(5000, 1000, 1.0, 0.1) == _gumbel_tau(1000, 1000, 1.0, 0.1)


def test_gumbel_tau_handles_zero_total_steps():
    # total_steps=0 is guarded to 1 internally: step=0 gives zero progress
    # (still tau_start), any step>=1 immediately clamps to full progress.
    assert _gumbel_tau(0, 0, 1.0, 0.1) == 1.0
    assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9