viva_glyph

-module(viva_glyph@codebook).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_glyph/codebook.gleam").
-export([empty/2, from_vectors/1, get/2, quantize/2, dequantize/2, update_centroid/3, init_deterministic/3]).
-export_type([codebook/0, quantize_result/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(
    " Codebook - Learned vocabulary for vector quantization\n"
    "\n"
    " A codebook is a set of prototype vectors (centroids).\n"
    " Each vector in latent space is mapped to its nearest centroid.\n"
    "\n"
    " ## Theory\n"
    "\n"
    " Based on Vector Quantization (VQ) from signal processing:\n"
    " - K centroids partition the space into Voronoi cells\n"
    " - Encoding: find nearest centroid index\n"
    " - Decoding: lookup centroid by index\n"
    "\n"
    " ## References\n"
    "\n"
    " - Gray (1984). Vector Quantization. IEEE ASSP Magazine.\n"
    " - van den Oord et al. (2017). VQ-VAE. NeurIPS.\n"
).

-type codebook() :: {codebook, list(list(float())), integer(), integer()}.

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

-file("src/viva_glyph/codebook.gleam", 45).
?DOC(" Create empty codebook\n").
-spec empty(integer(), integer()) -> codebook().
empty(Dimension, Size) ->
    Centroids = begin
        _pipe = gleam@list:range(0, Size - 1),
        gleam@list:map(_pipe, fun(_) -> viva_glyph@vector:zeros(Dimension) end)
    end,
    {codebook, Centroids, Dimension, Size}.

-file("src/viva_glyph/codebook.gleam", 54).
?DOC(" Create codebook from list of vectors\n").
-spec from_vectors(list(list(float()))) -> gleam@option:option(codebook()).
from_vectors(Vectors) ->
    case Vectors of
        [] ->
            none;

        [First | _] ->
            Dim = viva_glyph@vector:dimension(First),
            Size = erlang:length(Vectors),
            {some, {codebook, Vectors, Dim, Size}}
    end.

-file("src/viva_glyph/codebook.gleam", 66).
?DOC(" Get centroid by index\n").
-spec get(codebook(), integer()) -> gleam@option:option(list(float())).
get(Codebook, Index) ->
    case (Index >= 0) andalso (Index < erlang:element(4, Codebook)) of
        true ->
            _pipe = erlang:element(2, Codebook),
            _pipe@1 = gleam@list:drop(_pipe, Index),
            _pipe@2 = gleam@list:first(_pipe@1),
            gleam@option:from_result(_pipe@2);

        false ->
            none
    end.

-file("src/viva_glyph/codebook.gleam", 80).
?DOC(
    " Find nearest centroid to input vector\n"
    " Returns index and quantization error\n"
).
-spec quantize(codebook(), list(float())) -> quantize_result().
quantize(Codebook, Input) ->
    Indexed = begin
        _pipe = erlang:element(2, Codebook),
        gleam@list:index_map(_pipe, fun(Centroid, Idx) -> {Idx, Centroid} end)
    end,
    Initial = {quantize_result, 0, 999999.0},
    gleam@list:fold(
        Indexed,
        Initial,
        fun(Best, Pair) ->
            {Idx@1, Centroid@1} = Pair,
            Dist = viva_glyph@vector:distance_squared(Input, Centroid@1),
            case Dist < erlang:element(3, Best) of
                true ->
                    {quantize_result, Idx@1, Dist};

                false ->
                    Best
            end
        end
    ).

-file("src/viva_glyph/codebook.gleam", 98).
?DOC(" Dequantize: get centroid vector from index\n").
-spec dequantize(codebook(), integer()) -> list(float()).
dequantize(Codebook, Index) ->
    case get(Codebook, Index) of
        {some, V} ->
            V;

        none ->
            viva_glyph@vector:zeros(erlang:element(3, Codebook))
    end.

-file("src/viva_glyph/codebook.gleam", 123).
?DOC(" Update a single centroid (for training)\n").
-spec update_centroid(codebook(), integer(), list(float())) -> codebook().
update_centroid(Codebook, Index, New_centroid) ->
    New_centroids = begin
        _pipe = erlang:element(2, Codebook),
        gleam@list:index_map(_pipe, fun(C, I) -> case I =:= Index of
                    true ->
                        New_centroid;

                    false ->
                        C
                end end)
    end,
    {codebook,
        New_centroids,
        erlang:element(3, Codebook),
        erlang:element(4, Codebook)}.

-file("src/viva_glyph/codebook.gleam", 148).
-spec int_to_float_positive(integer(), float()) -> float().
int_to_float_positive(I, Acc) ->
    case I of
        0 ->
            Acc;

        _ ->
            int_to_float_positive(I - 1, Acc + 1.0)
    end.

-file("src/viva_glyph/codebook.gleam", 141).
-spec int_to_float(integer()) -> float().
int_to_float(I) ->
    case I >= 0 of
        true ->
            int_to_float_positive(I, +0.0);

        false ->
            +0.0 - int_to_float_positive(- I, +0.0)
    end.

-file("src/viva_glyph/codebook.gleam", 107).
?DOC(
    " Initialize codebook with random-ish values based on seed\n"
    " (Deterministic pseudo-random for reproducibility)\n"
).
-spec init_deterministic(integer(), integer(), integer()) -> codebook().
init_deterministic(Dimension, Size, Seed) ->
    Centroids = begin
        _pipe = gleam@list:range(0, Size - 1),
        gleam@list:map(
            _pipe,
            fun(I) -> _pipe@1 = gleam@list:range(0, Dimension - 1),
                gleam@list:map(
                    _pipe@1,
                    fun(D) ->
                        Hash = (((I * 31) + (D * 17)) + Seed) rem 1000,
                        (int_to_float(Hash) - 500.0) / 500.0
                    end
                ) end
        )
    end,
    {codebook, Centroids, Dimension, Size}.