viva_math

-module(viva_math@entropy).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/entropy.gleam").
-export([shannon/1, kl_divergence/2, symmetric_kl/2, jensen_shannon/2, cross_entropy/2, binary_cross_entropy/2, mutual_information/3, conditional_entropy/2, shannon_normalized/1, relative_entropy_rate/2]).

-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(
    " Entropy and information theory functions.\n"
    "\n"
    " Based on Shannon (1948) and Kullback-Leibler (1951).\n"
    " Used for memory consolidation scoring and uncertainty quantification.\n"
    "\n"
    " References:\n"
    " - Shannon (1948) \"A Mathematical Theory of Communication\"\n"
    " - Cover & Thomas (2006) \"Elements of Information Theory\"\n"
).

-file("src/viva_math/entropy.gleam", 26).
?DOC(
    " Shannon entropy: H(X) = -Σ p(x) log₂ p(x)\n"
    "\n"
    " Measures uncertainty/information content of a distribution.\n"
    " Higher entropy = more uncertainty.\n"
    "\n"
    " ## Examples\n"
    "\n"
    " ```gleam\n"
    " shannon([0.5, 0.5])     // -> 1.0 (maximum for 2 outcomes)\n"
    " shannon([1.0, 0.0])     // -> 0.0 (no uncertainty)\n"
    " shannon([0.25, 0.25, 0.25, 0.25]) // -> 2.0\n"
    " ```\n"
).
-spec shannon(list(float())) -> float().
shannon(Probabilities) ->
    gleam@list:fold(Probabilities, +0.0, fun(Acc, P) -> case P =< +0.0 of
                true ->
                    Acc;

                false ->
                    case gleam_community@maths:logarithm_2(P) of
                        {ok, Log_p} ->
                            Acc - (P * Log_p);

                        {error, _} ->
                            Acc
                    end
            end end).

-file("src/viva_math/entropy.gleam", 67).
?DOC(
    " KL Divergence: D_KL(P || Q) = Σ p(x) log(p(x) / q(x))\n"
    "\n"
    " Measures how much P diverges from Q (not symmetric!).\n"
    " P is the \"true\" distribution, Q is the approximation.\n"
    "\n"
    " Returns Error if distributions have different lengths or Q has zeros where P is non-zero.\n"
).
-spec kl_divergence(list(float()), list(float())) -> {ok, float()} |
    {error, nil}.
kl_divergence(P, Q) ->
    case erlang:length(P) =:= erlang:length(Q) of
        false ->
            {error, nil};

        true ->
            Pairs = gleam@list:zip(P, Q),
            Result = gleam@list:fold(
                Pairs,
                {ok, +0.0},
                fun(Acc, Pair) -> case Acc of
                        {error, nil} ->
                            {error, nil};

                        {ok, Sum} ->
                            {Pi, Qi} = Pair,
                            case Pi =< +0.0 of
                                true ->
                                    {ok, Sum};

                                false ->
                                    case Qi =< +0.0 of
                                        true ->
                                            {error, nil};

                                        false ->
                                            case gleam_community@maths:natural_logarithm(
                                                case Qi of
                                                    +0.0 -> +0.0;
                                                    -0.0 -> -0.0;
                                                    Gleam@denominator -> Pi / Gleam@denominator
                                                end
                                            ) of
                                                {ok, Log_ratio} ->
                                                    {ok, Sum + (Pi * Log_ratio)};

                                                {error, _} ->
                                                    {error, nil}
                                            end
                                    end
                            end
                    end end
            ),
            Result
    end.

-file("src/viva_math/entropy.gleam", 103).
?DOC(
    " Symmetric KL Divergence (Jensen-Shannon divergence without the 1/2).\n"
    "\n"
    " D_sym(P, Q) = D_KL(P || Q) + D_KL(Q || P)\n"
).
-spec symmetric_kl(list(float()), list(float())) -> {ok, float()} | {error, nil}.
symmetric_kl(P, Q) ->
    case {kl_divergence(P, Q), kl_divergence(Q, P)} of
        {{ok, D1}, {ok, D2}} ->
            {ok, D1 + D2};

        {_, _} ->
            {error, nil}
    end.

-file("src/viva_math/entropy.gleam", 114).
?DOC(
    " Jensen-Shannon Divergence: JS(P, Q) = (D_KL(P || M) + D_KL(Q || M)) / 2\n"
    " where M = (P + Q) / 2\n"
    "\n"
    " This is symmetric and bounded [0, 1] when using log₂.\n"
).
-spec jensen_shannon(list(float()), list(float())) -> {ok, float()} |
    {error, nil}.
jensen_shannon(P, Q) ->
    case erlang:length(P) =:= erlang:length(Q) of
        false ->
            {error, nil};

        true ->
            M = begin
                _pipe = gleam@list:zip(P, Q),
                gleam@list:map(
                    _pipe,
                    fun(Pair) ->
                        {Pi, Qi} = Pair,
                        (Pi + Qi) / 2.0
                    end
                )
            end,
            case {kl_divergence(P, M), kl_divergence(Q, M)} of
                {{ok, D_pm}, {ok, D_qm}} ->
                    {ok, (D_pm + D_qm) / 2.0};

                {_, _} ->
                    {error, nil}
            end
    end.

-file("src/viva_math/entropy.gleam", 138).
?DOC(
    " Cross-entropy: H(P, Q) = -Σ p(x) log q(x)\n"
    "\n"
    " Used in machine learning loss functions.\n"
    " H(P, Q) = H(P) + D_KL(P || Q)\n"
).
-spec cross_entropy(list(float()), list(float())) -> {ok, float()} |
    {error, nil}.
cross_entropy(P, Q) ->
    case erlang:length(P) =:= erlang:length(Q) of
        false ->
            {error, nil};

        true ->
            Pairs = gleam@list:zip(P, Q),
            Result = gleam@list:fold(
                Pairs,
                {ok, +0.0},
                fun(Acc, Pair) -> case Acc of
                        {error, nil} ->
                            {error, nil};

                        {ok, Sum} ->
                            {Pi, Qi} = Pair,
                            case Pi =< +0.0 of
                                true ->
                                    {ok, Sum};

                                false ->
                                    case Qi =< +0.0 of
                                        true ->
                                            {error, nil};

                                        false ->
                                            case gleam_community@maths:natural_logarithm(
                                                Qi
                                            ) of
                                                {ok, Log_q} ->
                                                    {ok, Sum - (Pi * Log_q)};

                                                {error, _} ->
                                                    {error, nil}
                                            end
                                    end
                            end
                    end end
            ),
            Result
    end.

-file("src/viva_math/entropy.gleam", 172).
?DOC(
    " Binary cross-entropy for single probability.\n"
    "\n"
    " H(p, q) = -[p log(q) + (1-p) log(1-q)]\n"
).
-spec binary_cross_entropy(float(), float()) -> {ok, float()} | {error, nil}.
binary_cross_entropy(P, Q) ->
    Q_clamped = gleam@float:max(gleam@float:min(Q, 0.999999), 0.000001),
    Q_inv = 1.0 - Q_clamped,
    P_inv = 1.0 - P,
    case {gleam_community@maths:natural_logarithm(Q_clamped),
        gleam_community@maths:natural_logarithm(Q_inv)} of
        {{ok, Log_q}, {ok, Log_q_inv}} ->
            Result = (P * Log_q) + (P_inv * Log_q_inv),
            {ok, +0.0 - Result};

        {_, _} ->
            {error, nil}
    end.

-file("src/viva_math/entropy.gleam", 191).
?DOC(
    " Mutual Information: I(X; Y) = H(X) + H(Y) - H(X, Y)\n"
    "\n"
    " Measures shared information between two variables.\n"
    " Takes marginal distributions and joint distribution as input.\n"
).
-spec mutual_information(list(float()), list(float()), list(list(float()))) -> float().
mutual_information(Px, Py, Pxy) ->
    Hx = shannon(Px),
    Hy = shannon(Py),
    Pxy_flat = lists:append(Pxy),
    Hxy = shannon(Pxy_flat),
    (Hx + Hy) - Hxy.

-file("src/viva_math/entropy.gleam", 209).
?DOC(
    " Conditional entropy: H(X|Y) = H(X, Y) - H(Y)\n"
    "\n"
    " Uncertainty in X given knowledge of Y.\n"
).
-spec conditional_entropy(list(float()), list(list(float()))) -> float().
conditional_entropy(Px, Pxy) ->
    Hx = shannon(Px),
    Pxy_flat = lists:append(Pxy),
    Hxy = shannon(Pxy_flat),
    Hxy - Hx.

-file("src/viva_math/entropy.gleam", 240).
-spec int_to_float(integer()) -> float().
int_to_float(N) ->
    case N of
        0 ->
            +0.0;

        1 ->
            1.0;

        2 ->
            2.0;

        _ ->
            Half = N div 2,
            Remainder = N - (Half * 2),
            (int_to_float(Half) * 2.0) + int_to_float(Remainder)
    end.

-file("src/viva_math/entropy.gleam", 43).
?DOC(
    " Normalized Shannon entropy (0 to 1 range).\n"
    "\n"
    " Divides by log₂(n) where n is number of outcomes.\n"
).
-spec shannon_normalized(list(float())) -> float().
shannon_normalized(Probabilities) ->
    N = erlang:length(Probabilities),
    case N =< 1 of
        true ->
            +0.0;

        false ->
            H = shannon(Probabilities),
            case gleam_community@maths:logarithm_2(int_to_float(N)) of
                {ok, Max_h} ->
                    case Max_h =:= +0.0 of
                        true ->
                            +0.0;

                        false ->
                            case Max_h of
                                +0.0 -> +0.0;
                                -0.0 -> -0.0;
                                Gleam@denominator -> H / Gleam@denominator
                            end
                    end;

                {error, _} ->
                    +0.0
            end
    end.

-file("src/viva_math/entropy.gleam", 219).
?DOC(
    " Relative entropy rate for sequences.\n"
    "\n"
    " Used for measuring \"surprise\" in temporal data.\n"
).
-spec relative_entropy_rate(list(float()), list(float())) -> {ok, float()} |
    {error, nil}.
relative_entropy_rate(Observed, Expected) ->
    case erlang:length(Observed) =:= erlang:length(Expected) of
        false ->
            {error, nil};

        true ->
            N = erlang:length(Observed),
            case N =:= 0 of
                true ->
                    {ok, +0.0};

                false ->
                    case kl_divergence(Observed, Expected) of
                        {ok, Kl} ->
                            {ok, case int_to_float(N) of
                                    +0.0 -> +0.0;
                                    -0.0 -> -0.0;
                                    Gleam@denominator -> Kl / Gleam@denominator
                                end};

                        {error, nil} ->
                            {error, nil}
                    end
            end
    end.