-module(viva_tensor@nn@embedding).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/nn/embedding.gleam").
-export([embedding_init/2, embedding_init_uniform/2, embedding_forward/2, sinusoidal_encoding/2, learned_positional_init/2, learned_positional_forward/2, rope/2]).
-export_type([embedding/0, learned_positional_encoding/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(false).

-type embedding() :: {embedding,
        integer(),
        integer(),
        viva_tensor@tensor:tensor()}.

-type learned_positional_encoding() :: {learned_positional_encoding,
        integer(),
        integer(),
        viva_tensor@tensor:tensor()}.

-file("src/viva_tensor/nn/embedding.gleam", 51).
?DOC(false).
-spec embedding_init(integer(), integer()) -> embedding().
embedding_init(Num_embeddings, Embedding_dim) ->
    Weight = viva_tensor@tensor:zeros([Num_embeddings, Embedding_dim]),
    {embedding, Num_embeddings, Embedding_dim, Weight}.

-file("src/viva_tensor/nn/embedding.gleam", 72).
?DOC(false).
-spec embedding_init_uniform(integer(), integer()) -> embedding().
embedding_init_uniform(Num_embeddings, Embedding_dim) ->
    Dim_f = erlang:float(Embedding_dim),
    Limit = case Dim_f > +0.0 of
        true ->
            case viva_tensor@core@ffi:sqrt(Dim_f) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> 1.0 / Gleam@denominator
            end;

        false ->
            +0.0
    end,
    Size = Num_embeddings * Embedding_dim,
    Data = case Size =< 0 of
        true ->
            [];

        false ->
            _pipe = gleam@list:range(1, Size),
            gleam@list:map(
                _pipe,
                fun(_) ->
                    R = viva_tensor@core@ffi:random_uniform(),
                    ((R * 2.0) * Limit) - Limit
                end
            )
    end,
    Weight = {tensor, Data, [Num_embeddings, Embedding_dim]},
    {embedding, Num_embeddings, Embedding_dim, Weight}.

-file("src/viva_tensor/nn/embedding.gleam", 389).
?DOC(false).
-spec take_row(list(float()), integer(), integer()) -> list(float()).
take_row(Data, Row, Dim) ->
    _pipe = Data,
    _pipe@1 = gleam@list:drop(_pipe, Row * Dim),
    gleam@list:take(_pipe@1, Dim).

-file("src/viva_tensor/nn/embedding.gleam", 369).
?DOC(false).
-spec gather_rows(list(float()), list(integer()), integer()) -> {ok,
        list(float())} |
    {error, viva_tensor@core@error:tensor_error()}.
gather_rows(Data, Indices, Dim) ->
    _pipe = gleam@list:try_fold(
        Indices,
        [],
        fun(Acc, Idx) ->
            Row = take_row(Data, Idx, Dim),
            case erlang:length(Row) =:= Dim of
                true ->
                    {ok, [Row | Acc]};

                false ->
                    {error,
                        {index_out_of_bounds, Idx, case gleam@int:max(Dim, 1) of
                                0 -> 0;
                                Gleam@denominator -> erlang:length(Data) div Gleam@denominator
                            end}}
            end
        end
    ),
    gleam@result:map(_pipe, fun(Rows) -> _pipe@1 = Rows,
            _pipe@2 = lists:reverse(_pipe@1),
            lists:append(_pipe@2) end).

-file("src/viva_tensor/nn/embedding.gleam", 350).
?DOC(false).
-spec normalize_indices(list(float()), integer()) -> {ok, list(integer())} |
    {error, viva_tensor@core@error:tensor_error()}.
normalize_indices(Raw, Num_embeddings) ->
    _pipe = gleam@list:try_fold(
        Raw,
        [],
        fun(Acc, Value) ->
            Idx = erlang:trunc(Value),
            Wrapped = case Idx < 0 of
                true ->
                    Idx + Num_embeddings;

                false ->
                    Idx
            end,
            case (Wrapped < 0) orelse (Wrapped >= Num_embeddings) of
                true ->
                    {error, {index_out_of_bounds, Idx, Num_embeddings}};

                false ->
                    {ok, [Wrapped | Acc]}
            end
        end
    ),
    gleam@result:map(_pipe, fun lists:reverse/1).

-file("src/viva_tensor/nn/embedding.gleam", 115).
?DOC(false).
-spec embedding_forward(embedding(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
embedding_forward(Layer, Indices) ->
    case viva_tensor@tensor:rank(Indices) =:= 1 of
        false ->
            {error,
                {dimension_error,
                    <<"embedding_forward: indices must be 1D, got rank "/utf8,
                        (erlang:integer_to_binary(
                            viva_tensor@tensor:rank(Indices)
                        ))/binary>>}};

        true ->
            gleam@result:'try'(
                viva_tensor@tensor:try_to_list(Indices),
                fun(Idx_floats) ->
                    gleam@result:'try'(
                        normalize_indices(Idx_floats, erlang:element(2, Layer)),
                        fun(Idx_ints) ->
                            gleam@result:'try'(
                                viva_tensor@tensor:try_to_list(
                                    erlang:element(4, Layer)
                                ),
                                fun(Weight_data) ->
                                    gleam@result:'try'(
                                        gather_rows(
                                            Weight_data,
                                            Idx_ints,
                                            erlang:element(3, Layer)
                                        ),
                                        fun(Rows) ->
                                            {ok,
                                                {tensor,
                                                    Rows,
                                                    [erlang:length(Idx_ints),
                                                        erlang:element(3, Layer)]}}
                                        end
                                    )
                                end
                            )
                        end
                    )
                end
            )
    end.

-file("src/viva_tensor/nn/embedding.gleam", 439).
?DOC(false).
-spec pow_safe(float(), float()) -> float().
pow_safe(Base, Exponent) ->
    case gleam@float:power(Base, Exponent) of
        {ok, V} ->
            V;

        {error, _} ->
            1.0
    end.

-file("src/viva_tensor/nn/embedding.gleam", 165).
?DOC(false).
-spec sinusoidal_encoding(integer(), integer()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
sinusoidal_encoding(Max_len, Embedding_dim) ->
    case Max_len < 0 of
        true ->
            {error,
                {invalid_shape,
                    <<"sinusoidal_encoding: max_len must be non-negative"/utf8>>}};

        false ->
            case Embedding_dim =< 0 of
                true ->
                    {error,
                        {invalid_shape,
                            <<"sinusoidal_encoding: embedding_dim must be > 0"/utf8>>}};

                false ->
                    case (Embedding_dim rem 2) =:= 0 of
                        false ->
                            {error,
                                {invalid_shape,
                                    <<"sinusoidal_encoding: embedding_dim must be even"/utf8>>}};

                        true ->
                            Dim_f = erlang:float(Embedding_dim),
                            Positions = case Max_len =< 0 of
                                true ->
                                    [];

                                false ->
                                    gleam@list:range(0, Max_len - 1)
                            end,
                            Pair_indices = case (Embedding_dim div 2) =< 0 of
                                true ->
                                    [];

                                false ->
                                    gleam@list:range(
                                        0,
                                        (Embedding_dim div 2) - 1
                                    )
                            end,
                            Data = begin
                                _pipe = Positions,
                                gleam@list:flat_map(
                                    _pipe,
                                    fun(Pos) ->
                                        Pos_f = erlang:float(Pos),
                                        _pipe@1 = Pair_indices,
                                        gleam@list:flat_map(
                                            _pipe@1,
                                            fun(I) ->
                                                Two_i = erlang:float(2 * I),
                                                Exponent = case Dim_f of
                                                    +0.0 -> +0.0;
                                                    -0.0 -> -0.0;
                                                    Gleam@denominator -> Two_i / Gleam@denominator
                                                end,
                                                Denom = pow_safe(
                                                    10000.0,
                                                    Exponent
                                                ),
                                                Angle = case Denom of
                                                    +0.0 -> +0.0;
                                                    -0.0 -> -0.0;
                                                    Gleam@denominator@1 -> Pos_f
                                                    / Gleam@denominator@1
                                                end,
                                                [gleam_community@maths:sin(
                                                        Angle
                                                    ),
                                                    gleam_community@maths:cos(
                                                        Angle
                                                    )]
                                            end
                                        )
                                    end
                                )
                            end,
                            {ok, {tensor, Data, [Max_len, Embedding_dim]}}
                    end
            end
    end.

-file("src/viva_tensor/nn/embedding.gleam", 233).
?DOC(false).
-spec learned_positional_init(integer(), integer()) -> learned_positional_encoding().
learned_positional_init(Max_len, Embedding_dim) ->
    Base = embedding_init_uniform(Max_len, Embedding_dim),
    {learned_positional_encoding,
        Max_len,
        Embedding_dim,
        erlang:element(4, Base)}.

-file("src/viva_tensor/nn/embedding.gleam", 252).
?DOC(false).
-spec learned_positional_forward(learned_positional_encoding(), integer()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
learned_positional_forward(Layer, Len) ->
    case Len < 0 of
        true ->
            {error, {index_out_of_bounds, Len, erlang:element(2, Layer)}};

        false ->
            case Len > erlang:element(2, Layer) of
                true ->
                    {error,
                        {index_out_of_bounds, Len, erlang:element(2, Layer)}};

                false ->
                    gleam@result:'try'(
                        viva_tensor@tensor:try_to_list(erlang:element(4, Layer)),
                        fun(Weight_data) ->
                            Indices = case Len of
                                0 ->
                                    [];

                                _ ->
                                    gleam@list:range(0, Len - 1)
                            end,
                            gleam@result:'try'(
                                gather_rows(
                                    Weight_data,
                                    Indices,
                                    erlang:element(3, Layer)
                                ),
                                fun(Rows) ->
                                    {ok,
                                        {tensor,
                                            Rows,
                                            [Len, erlang:element(3, Layer)]}}
                                end
                            )
                        end
                    )
            end
    end.

-file("src/viva_tensor/nn/embedding.gleam", 407).
?DOC(false).
-spec rotate_row_loop(
    list(float()),
    float(),
    float(),
    integer(),
    integer(),
    float(),
    list(float())
) -> list(float()).
rotate_row_loop(Remaining, Pos_f, Dim_f, Pair_idx, Num_pairs, Base, Acc) ->
    case {Pair_idx >= Num_pairs, Remaining} of
        {true, _} ->
            Acc;

        {_, []} ->
            Acc;

        {_, [_]} ->
            Acc;

        {_, [X, Y | Rest]} ->
            Two_i = erlang:float(2 * Pair_idx),
            Theta = case pow_safe(Base, case Dim_f of
                    +0.0 -> +0.0;
                    -0.0 -> -0.0;
                    Gleam@denominator -> Two_i / Gleam@denominator
                end) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@1 -> 1.0 / Gleam@denominator@1
            end,
            Angle = Pos_f * Theta,
            C = gleam_community@maths:cos(Angle),
            S = gleam_community@maths:sin(Angle),
            X_new = (X * C) - (Y * S),
            Y_new = (X * S) + (Y * C),
            rotate_row_loop(
                Rest,
                Pos_f,
                Dim_f,
                Pair_idx + 1,
                Num_pairs,
                Base,
                [Y_new, X_new | Acc]
            )
    end.

-file("src/viva_tensor/nn/embedding.gleam", 396).
?DOC(false).
-spec rotate_row(list(float()), float(), float(), integer(), float()) -> list(float()).
rotate_row(Row, Pos_f, Dim_f, Num_pairs, Base) ->
    _pipe = rotate_row_loop(Row, Pos_f, Dim_f, 0, Num_pairs, Base, []),
    lists:reverse(_pipe).

-file("src/viva_tensor/nn/embedding.gleam", 310).
?DOC(false).
-spec rope(viva_tensor@tensor:tensor(), float()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
rope(Input, Base) ->
    case viva_tensor@tensor:shape(Input) of
        [Seq_len, Dim] ->
            case (Dim rem 2) =:= 0 of
                false ->
                    {error, {invalid_shape, <<"rope: dim must be even"/utf8>>}};

                true ->
                    case Base =< +0.0 of
                        true ->
                            {error,
                                {invalid_shape,
                                    <<"rope: base must be > 0"/utf8>>}};

                        false ->
                            gleam@result:'try'(
                                viva_tensor@tensor:try_to_list(Input),
                                fun(Data) ->
                                    Dim_f = erlang:float(Dim),
                                    Positions = case Seq_len =< 0 of
                                        true ->
                                            [];

                                        false ->
                                            gleam@list:range(0, Seq_len - 1)
                                    end,
                                    Rotated = begin
                                        _pipe = Positions,
                                        gleam@list:flat_map(
                                            _pipe,
                                            fun(Pos) ->
                                                Pos_f = erlang:float(Pos),
                                                Row = take_row(Data, Pos, Dim),
                                                rotate_row(
                                                    Row,
                                                    Pos_f,
                                                    Dim_f,
                                                    Dim div 2,
                                                    Base
                                                )
                                            end
                                        )
                                    end,
                                    {ok, {tensor, Rotated, [Seq_len, Dim]}}
                                end
                            )
                    end
            end;

        Other ->
            {error,
                {dimension_error,
                    <<"rope: expected 2D tensor, got shape "/utf8,
                        (viva_tensor@core@error:shape_to_string(Other))/binary>>}}
    end.