viva_math

-module(viva_math@free_energy).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/free_energy.gleam").
-export([default_thresholds/0, prediction_error/2, precision_weighted_prediction_error/3, gaussian_kl_divergence/3, gaussian_kl_divergence_full/4, complexity/3, complexity_weighted/3, free_energy/5, classify_feeling_normalized/2, compute_state/6, classify_feeling/1, compute_state_simple/4, update_thresholds/3, surprise/3, active_inference_delta/3, precision_weighted_error_vec/3, belief_update/4, generalized_free_energy/3, variational_bound/2, estimate_precision/1]).
-export_type([free_energy_state/0, feeling/0, feeling_thresholds/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(
    " Free Energy Principle (FEP) calculations.\n"
    "\n"
    " Based on Karl Friston's work (2010, 2019).\n"
    " Free Energy bounds surprise (negative log evidence) and can be decomposed as:\n"
    "\n"
    " F = Π · (μ - o)² + D_KL(q || p)\n"
    "     ↑              ↑\n"
    "     Accuracy       Complexity\n"
    "     (weighted      (deviation\n"
    "     prediction     from priors)\n"
    "     error)\n"
    "\n"
    " In VIVA, this is used for interoception - sensing internal state\n"
    " and minimizing \"surprise\" through prediction.\n"
    "\n"
    " References:\n"
    " - Friston (2010) \"The free-energy principle: a unified brain theory?\"\n"
    " - Parr & Friston (2019) \"Generalised free energy and active inference\"\n"
    " - Validated by DeepSeek R1 671B (2025)\n"
).

-type free_energy_state() :: {free_energy_state,
        float(),
        float(),
        float(),
        float(),
        feeling()}.

-type feeling() :: homeostatic | surprised | alarmed | overwhelmed.

-type feeling_thresholds() :: {feeling_thresholds, float(), float()}.

-file("src/viva_math/free_energy.gleam", 67).
?DOC(
    " Default thresholds calibrated for PAD space.\n"
    " Mean and std_dev derived from typical emotional dynamics.\n"
).
-spec default_thresholds() -> feeling_thresholds().
default_thresholds() ->
    {feeling_thresholds, 0.5, 0.3}.

-file("src/viva_math/free_energy.gleam", 73).
?DOC(
    " Compute raw prediction error between expected and actual state.\n"
    " Uses squared Euclidean distance (L2 loss).\n"
).
-spec prediction_error(viva_math@vector:vec3(), viva_math@vector:vec3()) -> float().
prediction_error(Expected, Actual) ->
    viva_math@vector:distance_squared(Expected, Actual).

-file("src/viva_math/free_energy.gleam", 83).
?DOC(
    " Compute precision-weighted prediction error.\n"
    "\n"
    " F_accuracy = Π · (expected - actual)²\n"
    "\n"
    " Precision (Π) = 1/variance. Higher precision = more weight on prediction errors.\n"
    " This is critical for biological systems where uncertainty should attenuate errors.\n"
).
-spec precision_weighted_prediction_error(
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    float()
) -> float().
precision_weighted_prediction_error(Expected, Actual, Precision) ->
    Pe = prediction_error(Expected, Actual),
    Precision * Pe.

-file("src/viva_math/free_energy.gleam", 100).
?DOC(
    " Compute KL divergence between Gaussian distributions (closed form).\n"
    "\n"
    " CORRECTED per DeepSeek R1 validation - Full KL for Gaussians:\n"
    " D_KL(N(μ₁,σ₁²) || N(μ₂,σ₂²)) = (μ₁ - μ₂)²/(2σ₂²) + (σ₁² - σ₂²)/(2σ₂²) - 1/2\n"
    "\n"
    " When variances are equal (σ₁ = σ₂), reduces to: (μ₁ - μ₂)²/(2σ²)\n"
    "\n"
    " This measures how much the posterior (current belief) diverges from prior.\n"
).
-spec gaussian_kl_divergence(
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    float()
) -> float().
gaussian_kl_divergence(Posterior_mean, Prior_mean, Variance) ->
    Diff_squared = viva_math@vector:distance_squared(Posterior_mean, Prior_mean),
    case Variance =< +0.0 of
        true ->
            +0.0;

        false ->
            case (2.0 * Variance) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> Diff_squared / Gleam@denominator
            end
    end.

-file("src/viva_math/free_energy.gleam", 117).
?DOC(
    " Full KL divergence between Gaussians with different variances.\n"
    "\n"
    " D_KL(N(μ₁,σ₁²) || N(μ₂,σ₂²)) = log(σ₂/σ₁) + (σ₁² + (μ₁-μ₂)²)/(2σ₂²) - 1/2\n"
    "\n"
    " This is the complete formula from DeepSeek R1 validation.\n"
).
-spec gaussian_kl_divergence_full(
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    float(),
    float()
) -> float().
gaussian_kl_divergence_full(
    Posterior_mean,
    Prior_mean,
    Posterior_variance,
    Prior_variance
) ->
    case (Posterior_variance =< +0.0) orelse (Prior_variance =< +0.0) of
        true ->
            +0.0;

        false ->
            Diff_squared = viva_math@vector:distance_squared(
                Posterior_mean,
                Prior_mean
            ),
            Log_ratio = case {gleam_community@maths:natural_logarithm(
                    Prior_variance
                ),
                gleam_community@maths:natural_logarithm(Posterior_variance)} of
                {{ok, Log_prior}, {ok, Log_posterior}} ->
                    (Log_prior - Log_posterior) / 2.0;

                {_, _} ->
                    +0.0
            end,
            Ratio_term = case (2.0 * Prior_variance) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> (Posterior_variance + Diff_squared) / Gleam@denominator
            end,
            (Log_ratio + Ratio_term) - 0.5
    end.

-file("src/viva_math/free_energy.gleam", 150).
?DOC(
    " Compute complexity term using KL divergence.\n"
    "\n"
    " Complexity = D_KL(q(θ) || p(θ))\n"
    "\n"
    " Where q is posterior belief and p is prior belief (homeostatic setpoint).\n"
    " Weight controls the regularization strength.\n"
).
-spec complexity(viva_math@vector:vec3(), viva_math@vector:vec3(), float()) -> float().
complexity(Current, Baseline, Prior_variance) ->
    gaussian_kl_divergence(Current, Baseline, Prior_variance).

-file("src/viva_math/free_energy.gleam", 159).
?DOC(" Legacy complexity function for backwards compatibility.\n").
-spec complexity_weighted(
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    float()
) -> float().
complexity_weighted(Current, Baseline, Weight) ->
    Weight * viva_math@vector:distance_squared(Current, Baseline).

-file("src/viva_math/free_energy.gleam", 175).
?DOC(
    " Compute full Free Energy: F = Π·(μ-o)² + D_KL(q||p)\n"
    "\n"
    " ## Parameters\n"
    " - expected: predicted/expected state (μ)\n"
    " - actual: observed/actual state (o)\n"
    " - baseline: prior baseline state (p) - e.g., personality/homeostatic setpoint\n"
    " - precision: inverse variance of predictions (Π)\n"
    " - prior_variance: variance of prior beliefs (for KL term)\n"
).
-spec free_energy(
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    float(),
    float()
) -> float().
free_energy(Expected, Actual, Baseline, Precision, Prior_variance) ->
    Accuracy = precision_weighted_prediction_error(Expected, Actual, Precision),
    Cx = complexity(Actual, Baseline, Prior_variance),
    Accuracy + Cx.

-file("src/viva_math/free_energy.gleam", 237).
?DOC(
    " Classify feeling using normalized thresholds.\n"
    "\n"
    " - Homeostatic: F < μ - σ (better than expected)\n"
    " - Surprised: μ - σ ≤ F < μ (slightly worse)\n"
    " - Alarmed: μ ≤ F < μ + σ (worse than average)\n"
    " - Overwhelmed: F ≥ μ + σ (much worse)\n"
).
-spec classify_feeling_normalized(float(), feeling_thresholds()) -> feeling().
classify_feeling_normalized(Free_energy, Thresholds) ->
    Lower = erlang:element(2, Thresholds) - erlang:element(3, Thresholds),
    Upper = erlang:element(2, Thresholds) + erlang:element(3, Thresholds),
    case Free_energy of
        Fe when Fe < Lower ->
            homeostatic;

        Fe@1 when Fe@1 < erlang:element(2, Thresholds) ->
            surprised;

        Fe@2 when Fe@2 < Upper ->
            alarmed;

        _ ->
            overwhelmed
    end.

-file("src/viva_math/free_energy.gleam", 189).
?DOC(
    " Compute free energy and return full state with feeling.\n"
    " Uses normalized thresholds for feeling classification.\n"
).
-spec compute_state(
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    float(),
    float(),
    feeling_thresholds()
) -> free_energy_state().
compute_state(Expected, Actual, Baseline, Precision, Prior_variance, Thresholds) ->
    Accuracy = precision_weighted_prediction_error(Expected, Actual, Precision),
    Cx = complexity(Actual, Baseline, Prior_variance),
    Fe = Accuracy + Cx,
    {free_energy_state,
        Fe,
        Accuracy,
        Cx,
        Precision,
        classify_feeling_normalized(Fe, Thresholds)}.

-file("src/viva_math/free_energy.gleam", 254).
?DOC(
    " Legacy classify_feeling with fixed thresholds.\n"
    " Calibrated for PAD space (max distance ~3.46).\n"
).
-spec classify_feeling(float()) -> feeling().
classify_feeling(Free_energy) ->
    case Free_energy of
        Fe when Fe < 0.1 ->
            homeostatic;

        Fe@1 when Fe@1 < 0.5 ->
            surprised;

        Fe@2 when Fe@2 < 1.5 ->
            alarmed;

        _ ->
            overwhelmed
    end.

-file("src/viva_math/free_energy.gleam", 212).
?DOC(
    " Simplified compute_state with default thresholds and legacy interface.\n"
    " For backwards compatibility.\n"
).
-spec compute_state_simple(
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    float()
) -> free_energy_state().
compute_state_simple(Expected, Actual, Baseline, Complexity_weight) ->
    Pe = prediction_error(Expected, Actual),
    Cx = complexity_weighted(Actual, Baseline, Complexity_weight),
    Fe = Pe + Cx,
    {free_energy_state, Fe, Pe, Cx, 1.0, classify_feeling(Fe)}.

-file("src/viva_math/free_energy.gleam", 265).
?DOC(
    " Update thresholds based on observed free energy history.\n"
    " Uses exponential moving average for online learning.\n"
).
-spec update_thresholds(feeling_thresholds(), float(), float()) -> feeling_thresholds().
update_thresholds(Current, Observed_fe, Alpha) ->
    New_mean = (Alpha * Observed_fe) + ((1.0 - Alpha) * erlang:element(
        2,
        Current
    )),
    Diff = Observed_fe - erlang:element(2, Current),
    New_var = (Alpha * (Diff * Diff)) + ((1.0 - Alpha) * (erlang:element(
        3,
        Current
    )
    * erlang:element(3, Current))),
    New_std = case gleam_community@maths:nth_root(New_var, 2) of
        {ok, S} ->
            S;

        {error, _} ->
            erlang:element(3, Current)
    end,
    {feeling_thresholds, New_mean, gleam@float:max(New_std, 0.01)}.

-file("src/viva_math/free_energy.gleam", 290).
?DOC(
    " Compute surprise for a single dimension.\n"
    "\n"
    " Surprise = -log(p(observation | model))\n"
    " Using Gaussian approximation: surprise ∝ (x - μ)² / (2σ²)\n"
).
-spec surprise(float(), float(), float()) -> float().
surprise(Expected, Observed, Sigma) ->
    Diff = Observed - Expected,
    Sigma_sq = Sigma * Sigma,
    case Sigma_sq =< +0.0 of
        true ->
            +0.0;

        false ->
            case (2.0 * Sigma_sq) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> (Diff * Diff) / Gleam@denominator
            end
    end.

-file("src/viva_math/free_energy.gleam", 303).
?DOC(
    " Active Inference: compute action that minimizes expected free energy.\n"
    "\n"
    " This returns the delta to apply to current state to move toward target.\n"
    " Rate controls how quickly to move (0 = no movement, 1 = instant).\n"
).
-spec active_inference_delta(
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    float()
) -> viva_math@vector:vec3().
active_inference_delta(Current, Target, Rate) ->
    Diff = viva_math@vector:sub(Target, Current),
    viva_math@vector:scale(Diff, Rate).

-file("src/viva_math/free_energy.gleam", 316).
?DOC(
    " Precision-weighted prediction error for Vec3.\n"
    "\n"
    " Each dimension can have different precision.\n"
    " Returns weighted sum of squared errors.\n"
).
-spec precision_weighted_error_vec(
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    viva_math@vector:vec3()
) -> float().
precision_weighted_error_vec(Expected, Actual, Precisions) ->
    Diff = viva_math@vector:sub(Expected, Actual),
    Diff_sq = viva_math@vector:multiply(Diff, Diff),
    Weighted = viva_math@vector:multiply(Diff_sq, Precisions),
    viva_math@vector:sum(Weighted).

-file("src/viva_math/free_energy.gleam", 359).
?DOC(
    " Bayesian belief update: combine prior with likelihood.\n"
    "\n"
    " posterior ∝ likelihood × prior\n"
    " Using precision-weighted combination:\n"
    " new_belief = (Π_prior × prior + Π_likelihood × observation) /\n"
    "              (Π_prior + Π_likelihood)\n"
).
-spec belief_update(float(), float(), float(), float()) -> float().
belief_update(Prior, Observation, Precision_prior, Precision_likelihood) ->
    Total_precision = Precision_prior + Precision_likelihood,
    case Total_precision =< +0.0 of
        true ->
            Prior;

        false ->
            case Total_precision of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> ((Precision_prior * Prior) + (Precision_likelihood
                * Observation))
                / Gleam@denominator
            end
    end.

-file("src/viva_math/free_energy.gleam", 381).
?DOC(
    " Generalized Free Energy (expected free energy for planning).\n"
    "\n"
    " G = ambiguity + risk\n"
    " - ambiguity: expected surprise under model (epistemic value)\n"
    " - risk: KL divergence from preferred outcomes (pragmatic value)\n"
    "\n"
    " Used for action selection in active inference.\n"
).
-spec generalized_free_energy(
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    float()
) -> float().
generalized_free_energy(Expected_state, Preferred_state, Uncertainty) ->
    Ambiguity = Uncertainty,
    Risk = viva_math@vector:distance_squared(Expected_state, Preferred_state),
    Ambiguity + Risk.

-file("src/viva_math/free_energy.gleam", 396).
?DOC(
    " Variational Free Energy bound.\n"
    "\n"
    " F ≤ -log p(o) + D_KL(q||p)\n"
    "\n"
    " The free energy bounds the negative log evidence (surprise).\n"
).
-spec variational_bound(float(), float()) -> float().
variational_bound(Observation_likelihood, Kl_divergence) ->
    Neg_log_likelihood = case Observation_likelihood =< +0.0 of
        true ->
            100.0;

        false ->
            case gleam_community@maths:natural_logarithm(Observation_likelihood) of
                {ok, Log_l} ->
                    +0.0 - Log_l;

                {error, _} ->
                    100.0
            end
    end,
    Neg_log_likelihood + Kl_divergence.

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

        1 ->
            1.0;

        2 ->
            2.0;

        3 ->
            3.0;

        4 ->
            4.0;

        5 ->
            5.0;

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

-file("src/viva_math/free_energy.gleam", 331).
?DOC(
    " Estimate precision from recent prediction errors.\n"
    "\n"
    " Precision = 1 / variance of errors\n"
    " Higher precision means more reliable predictions.\n"
).
-spec estimate_precision(list(float())) -> float().
estimate_precision(Errors) ->
    case erlang:length(Errors) of
        0 ->
            1.0;

        1 ->
            1.0;

        N ->
            N_float = int_to_float(N),
            Mean = case N_float of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> gleam@list:fold(
                    Errors,
                    +0.0,
                    fun(Acc, E) -> Acc + E end
                )
                / Gleam@denominator
            end,
            Variance = case N_float of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@1 -> gleam@list:fold(
                    Errors,
                    +0.0,
                    fun(Acc@1, E@1) ->
                        Diff = E@1 - Mean,
                        Acc@1 + (Diff * Diff)
                    end
                )
                / Gleam@denominator@1
            end,
            case Variance < 0.001 of
                true ->
                    100.0;

                false ->
                    case Variance of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator@2 -> 1.0 / Gleam@denominator@2
                    end
            end
    end.