viva_math

-module(viva_math@scheduler).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/scheduler.gleam").
-export([step_decay/4, multi_step_decay/4, exponential/3, inverse/3, inverse_sqrt/3, polynomial/4, linear_warmup/3, cosine_annealing/4, cosine_warm_restarts/5, one_cycle_defaults/2, one_cycle/2, triangle/2]).
-export_type([one_cycle_config/0]).

-if(?OTP_RELEASE >= 27).
-define(MODULEDOC(Str), -moduledoc(Str)).
-define(DOC(Str), -doc(Str)).
-else.
-define(MODULEDOC(Str), -compile([])).
-define(DOC(Str), -compile([])).
-endif.

?MODULEDOC(
    " Learning-rate / coefficient schedulers.\n"
    "\n"
    " Pure functions `(step, config) -> Float`. No internal state. Useful for:\n"
    "\n"
    " - `viva_tensor` learning-rate schedules during training\n"
    " - `viva_emotion` decay of stimulus weights / mood half-lives\n"
    " - any time-varying coefficient\n"
    "\n"
    " All schedulers are total — even out-of-range inputs are clamped to a\n"
    " sensible value rather than erroring. This makes them safe inside hot\n"
    " training loops.\n"
).

-type one_cycle_config() :: {one_cycle_config,
        float(),
        float(),
        float(),
        integer(),
        float()}.

-file("src/viva_math/scheduler.gleam", 268).
-spec int_to_float(integer()) -> float().
int_to_float(N) ->
    erlang:float(N).

-file("src/viva_math/scheduler.gleam", 23).
?DOC(
    " Step decay: lr · γ^(step / step_size).\n"
    "\n"
    " Drops by factor γ every `step_size` steps.\n"
).
-spec step_decay(float(), integer(), integer(), float()) -> float().
step_decay(Base_lr, Step, Step_size, Gamma) ->
    case Step_size =< 0 of
        true ->
            Base_lr;

        false ->
            Drops = case Step_size of
                0 -> 0;
                Gleam@denominator -> Step div Gleam@denominator
            end,
            Base_lr * math:pow(Gamma, int_to_float(Drops))
    end.

-file("src/viva_math/scheduler.gleam", 49).
-spec count_le(list(integer()), integer(), integer()) -> integer().
count_le(Xs, Step, Acc) ->
    case Xs of
        [] ->
            Acc;

        [M | Rest] ->
            case Step >= M of
                true ->
                    count_le(Rest, Step, Acc + 1);

                false ->
                    count_le(Rest, Step, Acc)
            end
    end.

-file("src/viva_math/scheduler.gleam", 39).
?DOC(" Multi-step decay: drops by γ each time step crosses a milestone.\n").
-spec multi_step_decay(float(), integer(), list(integer()), float()) -> float().
multi_step_decay(Base_lr, Step, Milestones, Gamma) ->
    Crossed = count_le(Milestones, Step, 0),
    Base_lr * math:pow(Gamma, int_to_float(Crossed)).

-file("src/viva_math/scheduler.gleam", 65).
?DOC(" Exponential decay: lr · γ^step.\n").
-spec exponential(float(), integer(), float()) -> float().
exponential(Base_lr, Step, Gamma) ->
    Base_lr * math:pow(Gamma, int_to_float(Step)).

-file("src/viva_math/scheduler.gleam", 70).
?DOC(" Inverse decay: lr / (1 + step / τ).\n").
-spec inverse(float(), integer(), float()) -> float().
inverse(Base_lr, Step, Tau) ->
    case Tau =< +0.0 of
        true ->
            Base_lr;

        false ->
            case (1.0 + (case Tau of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> int_to_float(Step) / Gleam@denominator
            end)) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@1 -> Base_lr / Gleam@denominator@1
            end
    end.

-file("src/viva_math/scheduler.gleam", 78).
?DOC(" Inverse square-root decay: lr / √(1 + step / τ). Common in transformers.\n").
-spec inverse_sqrt(float(), integer(), float()) -> float().
inverse_sqrt(Base_lr, Step, Tau) ->
    case Tau =< +0.0 of
        true ->
            Base_lr;

        false ->
            case math:sqrt(1.0 + (case Tau of
                    +0.0 -> +0.0;
                    -0.0 -> -0.0;
                    Gleam@denominator -> int_to_float(Step) / Gleam@denominator
                end)) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@1 -> Base_lr / Gleam@denominator@1
            end
    end.

-file("src/viva_math/scheduler.gleam", 86).
?DOC(" Polynomial decay: lr · (1 - step / total)^power, clamped at zero past `total`.\n").
-spec polynomial(float(), integer(), integer(), float()) -> float().
polynomial(Base_lr, Step, Total_steps, Power) ->
    case Total_steps =< 0 of
        true ->
            Base_lr;

        false ->
            case Step >= Total_steps of
                true ->
                    +0.0;

                false ->
                    Frac = 1.0 - (case int_to_float(Total_steps) of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator -> int_to_float(Step) / Gleam@denominator
                    end),
                    Base_lr * math:pow(Frac, Power)
            end
    end.

-file("src/viva_math/scheduler.gleam", 110).
?DOC(" Linear warmup: ramps from 0 to base_lr over `warmup_steps`.\n").
-spec linear_warmup(float(), integer(), integer()) -> float().
linear_warmup(Base_lr, Step, Warmup_steps) ->
    case Warmup_steps =< 0 of
        true ->
            Base_lr;

        false ->
            case Step >= Warmup_steps of
                true ->
                    Base_lr;

                false ->
                    case int_to_float(Warmup_steps) of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator -> Base_lr * int_to_float(Step) / Gleam@denominator
                    end
            end
    end.

-file("src/viva_math/scheduler.gleam", 124).
?DOC(
    " Cosine annealing: smooth cosine ride from base_lr to min_lr over T_max steps.\n"
    "\n"
    " lr(t) = min_lr + ½(base_lr - min_lr)·(1 + cos(π · t / T_max))\n"
).
-spec cosine_annealing(float(), integer(), integer(), float()) -> float().
cosine_annealing(Base_lr, Step, T_max, Min_lr) ->
    case T_max =< 0 of
        true ->
            Base_lr;

        false ->
            T = case Step > T_max of
                true ->
                    int_to_float(T_max);

                false ->
                    int_to_float(Step)
            end,
            Progress = case int_to_float(T_max) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> T / Gleam@denominator
            end,
            Cos_term = math:cos(3.141592653589793 * Progress),
            Min_lr + ((0.5 * (Base_lr - Min_lr)) * (1.0 + Cos_term))
    end.

-file("src/viva_math/scheduler.gleam", 163).
-spec advance_period(integer(), integer(), integer(), integer()) -> {integer(),
    integer()}.
advance_period(Step, Current_period, T_mult, Budget) ->
    case Step < Budget of
        true ->
            {Current_period, Step - (Budget - Current_period)};

        false ->
            Mult = case T_mult < 1 of
                true ->
                    1;

                false ->
                    T_mult
            end,
            Next_period = Current_period * Mult,
            advance_period(Step, Next_period, T_mult, Budget + Next_period)
    end.

-file("src/viva_math/scheduler.gleam", 146).
?DOC(
    " Cosine warmup-restart: cosine annealing that resets to base_lr every\n"
    " `period` steps. The period doubles each restart when `t_mult = 2`.\n"
).
-spec cosine_warm_restarts(float(), integer(), integer(), integer(), float()) -> float().
cosine_warm_restarts(Base_lr, Step, Period, T_mult, Min_lr) ->
    case Period =< 0 of
        true ->
            Base_lr;

        false ->
            {Current_period, Step_in_period} = advance_period(
                Step,
                Period,
                T_mult,
                Period
            ),
            cosine_annealing(Base_lr, Step_in_period, Current_period, Min_lr)
    end.

-file("src/viva_math/scheduler.gleam", 198).
?DOC(" Default 1-cycle configuration: 30 % warmup, then anneal.\n").
-spec one_cycle_defaults(float(), integer()) -> one_cycle_config().
one_cycle_defaults(Max_lr, Total_steps) ->
    {one_cycle_config, Max_lr, Max_lr / 25.0, Max_lr / 1.0e4, Total_steps, 0.3}.

-file("src/viva_math/scheduler.gleam", 232).
?DOC(" Cosine interpolation between `a` and `b` for `t ∈ [0, 1]`.\n").
-spec cosine_interp(float(), float(), float()) -> float().
cosine_interp(A, B, T) ->
    T_clamped = viva_math@scalar:clamp_unit(T),
    Cos_term = math:cos(3.141592653589793 * T_clamped),
    A + ((0.5 * (B - A)) * (1.0 - Cos_term)).

-file("src/viva_math/scheduler.gleam", 280).
-spec int_max(integer(), integer()) -> integer().
int_max(A, B) ->
    case A > B of
        true ->
            A;

        false ->
            B
    end.

-file("src/viva_math/scheduler.gleam", 209).
?DOC(" One-cycle policy: cosine ramp up to `max_lr`, then cosine anneal down.\n").
-spec one_cycle(one_cycle_config(), integer()) -> float().
one_cycle(Config, Step) ->
    case erlang:element(5, Config) =< 0 of
        true ->
            erlang:element(3, Config);

        false ->
            Warmup_end = erlang:trunc(
                int_to_float(erlang:element(5, Config)) * erlang:element(
                    6,
                    Config
                )
            ),
            Progress = case Step =< Warmup_end of
                true ->
                    P = case int_to_float(int_max(Warmup_end, 1)) of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator -> int_to_float(Step) / Gleam@denominator
                    end,
                    cosine_interp(
                        erlang:element(3, Config),
                        erlang:element(2, Config),
                        P
                    );

                false ->
                    Span = int_max(erlang:element(5, Config) - Warmup_end, 1),
                    P@1 = case int_to_float(Span) of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator@1 -> int_to_float(Step - Warmup_end) / Gleam@denominator@1
                    end,
                    cosine_interp(
                        erlang:element(2, Config),
                        erlang:element(4, Config),
                        P@1
                    )
            end,
            Progress
    end.

-file("src/viva_math/scheduler.gleam", 287).
-spec float_floor_div(float(), float()) -> float().
float_floor_div(A, B) ->
    _pipe = erlang:trunc(case B of
            +0.0 -> +0.0;
            -0.0 -> -0.0;
            Gleam@denominator -> A / Gleam@denominator
        end),
    int_to_float(_pipe).

-file("src/viva_math/scheduler.gleam", 243).
?DOC(" Triangle waveform: 0 → 1 → 0 over `period` steps.\n").
-spec triangle(integer(), integer()) -> float().
triangle(Step, Period) ->
    case Period =< 0 of
        true ->
            +0.0;

        false ->
            Period_f = int_to_float(Period),
            Half = Period_f / 2.0,
            Step_f = int_to_float(Step),
            Pos = Step_f - (Period_f * float_floor_div(Step_f, Period_f)),
            case Pos =< Half of
                true ->
                    case Half of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator -> Pos / Gleam@denominator
                    end;

                false ->
                    2.0 - (case Half of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator@1 -> Pos / Gleam@denominator@1
                    end)
            end
    end.