viva_math

-module(viva_math@matrix_dense).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/matrix_dense.gleam").
-export([zeros/2, identity/1, from_list/3, get/3, row/2, to_rows/1, byte_size/1, add/2, sub/2, hadamard/2, scale/2, transpose/1, mul/2, frobenius/1, trace/1]).
-export_type([dense_mat/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(
    " Dense matrix backed by `BitArray` (binary row-major).\n"
    "\n"
    " Counterpart to `viva_math/matrix.MatN`. Where `MatN` stores rows as\n"
    " nested `List(Float)`, `DenseMat` keeps a single contiguous binary\n"
    " buffer with one IEEE-754 little-endian double per element. This:\n"
    "\n"
    " - **Saves memory**: 8 bytes/element vs ~24 bytes/element for cons-cell\n"
    "   lists on 64-bit BEAM.\n"
    " - **Speeds up indexed access**: O(1) random read of any cell via\n"
    "   `byte_offset = (row · cols + col) · 8`.\n"
    " - **Avoids list reversal**: `transpose` and `matmul` rebuild a binary\n"
    "   instead of allocating thousands of cons cells.\n"
    "\n"
    " ## Scope\n"
    "\n"
    " Suitable for matrices up to a few thousand rows × cols. For tensor-scale\n"
    " work (batched GEMM, GPU acceleration) defer to `viva_tensor`. This\n"
    " module is the middle ground between the trivially-correct `MatN` and\n"
    " the heavy `viva_tensor`.\n"
    "\n"
    " ## Limitations\n"
    "\n"
    " - Float-only (`Float` ↔ 64-bit IEEE 754 little-endian).\n"
    " - Sequential algorithms (no parallelism, no SIMD).\n"
    " - Conversion to/from `MatN` available for interop.\n"
).

-type dense_mat() :: {dense_mat, integer(), integer(), bitstring()}.

-file("src/viva_math/matrix_dense.gleam", 49).
-spec build_zeros(integer(), bitstring()) -> bitstring().
build_zeros(N, Acc) ->
    case N =< 0 of
        true ->
            Acc;

        false ->
            build_zeros(N - 1, <<Acc/bitstring, +0.0:64/float-little>>)
    end.

-file("src/viva_math/matrix_dense.gleam", 39).
?DOC(" Zero matrix of given shape. Errors on non-positive dimensions.\n").
-spec zeros(integer(), integer()) -> {ok, dense_mat()} | {error, nil}.
zeros(Rows, Cols) ->
    case (Rows =< 0) orelse (Cols =< 0) of
        true ->
            {error, nil};

        false ->
            Zero_data = build_zeros(Rows * Cols, <<>>),
            {ok, {dense_mat, Rows, Cols, Zero_data}}
    end.

-file("src/viva_math/matrix_dense.gleam", 67).
-spec build_identity(integer(), integer(), integer(), bitstring()) -> bitstring().
build_identity(N, I, J, Acc) ->
    case I >= N of
        true ->
            Acc;

        false ->
            V = case I =:= J of
                true ->
                    1.0;

                false ->
                    +0.0
            end,
            Next_acc = <<Acc/bitstring, V:64/float-little>>,
            case (J + 1) >= N of
                true ->
                    build_identity(N, I + 1, 0, Next_acc);

                false ->
                    build_identity(N, I, J + 1, Next_acc)
            end
    end.

-file("src/viva_math/matrix_dense.gleam", 57).
?DOC(" Identity matrix of size n.\n").
-spec identity(integer()) -> {ok, dense_mat()} | {error, nil}.
identity(N) ->
    case N =< 0 of
        true ->
            {error, nil};

        false ->
            Data = build_identity(N, 0, 0, <<>>),
            {ok, {dense_mat, N, N, Data}}
    end.

-file("src/viva_math/matrix_dense.gleam", 99).
-spec pack_floats(list(float()), bitstring()) -> bitstring().
pack_floats(Xs, Acc) ->
    case Xs of
        [] ->
            Acc;

        [X | Rest] ->
            pack_floats(Rest, <<Acc/bitstring, X:64/float-little>>)
    end.

-file("src/viva_math/matrix_dense.gleam", 85).
?DOC(" Build from row-major list of values. Errors if shape doesn't match.\n").
-spec from_list(integer(), integer(), list(float())) -> {ok, dense_mat()} |
    {error, nil}.
from_list(Rows, Cols, Values) ->
    case ((Rows =< 0) orelse (Cols =< 0)) orelse (erlang:length(Values) /= (Rows
    * Cols)) of
        true ->
            {error, nil};

        false ->
            Data = pack_floats(Values, <<>>),
            {ok, {dense_mat, Rows, Cols, Data}}
    end.

-file("src/viva_math/matrix_dense.gleam", 111).
?DOC(" Random-access read in O(1). Errors on out-of-bounds.\n").
-spec get(dense_mat(), integer(), integer()) -> {ok, float()} | {error, nil}.
get(M, Row, Col) ->
    case (((Row < 0) orelse (Row >= erlang:element(2, M))) orelse (Col < 0))
    orelse (Col >= erlang:element(3, M)) of
        true ->
            {error, nil};

        false ->
            Index = (Row * erlang:element(3, M)) + Col,
            Bit_offset = Index * 64,
            case erlang:element(4, M) of
                <<_:Bit_offset, Value:64/float-little, _/bitstring>> ->
                    {ok, Value};

                _ ->
                    {error, nil}
            end
    end.

-file("src/viva_math/matrix_dense.gleam", 133).
-spec extract_row(dense_mat(), integer(), integer(), list(float())) -> list(float()).
extract_row(M, Row_idx, Col_idx, Acc) ->
    case Col_idx >= erlang:element(3, M) of
        true ->
            lists:reverse(Acc);

        false ->
            case get(M, Row_idx, Col_idx) of
                {ok, V} ->
                    extract_row(M, Row_idx, Col_idx + 1, [V | Acc]);

                {error, _} ->
                    lists:reverse(Acc)
            end
    end.

-file("src/viva_math/matrix_dense.gleam", 126).
?DOC(" Get an entire row as a list.\n").
-spec row(dense_mat(), integer()) -> {ok, list(float())} | {error, nil}.
row(M, Row_idx) ->
    case (Row_idx < 0) orelse (Row_idx >= erlang:element(2, M)) of
        true ->
            {error, nil};

        false ->
            {ok, extract_row(M, Row_idx, 0, [])}
    end.

-file("src/viva_math/matrix_dense.gleam", 154).
-spec build_rows(dense_mat(), integer(), list(list(float()))) -> list(list(float())).
build_rows(M, I, Acc) ->
    case I >= erlang:element(2, M) of
        true ->
            lists:reverse(Acc);

        false ->
            case row(M, I) of
                {ok, R} ->
                    build_rows(M, I + 1, [R | Acc]);

                {error, _} ->
                    lists:reverse(Acc)
            end
    end.

-file("src/viva_math/matrix_dense.gleam", 150).
?DOC(" Convert to nested-list representation for interop with `MatN`.\n").
-spec to_rows(dense_mat()) -> list(list(float())).
to_rows(M) ->
    build_rows(M, 0, []).

-file("src/viva_math/matrix_dense.gleam", 170).
?DOC(" Number of bytes consumed by the data buffer. Useful for benchmarking.\n").
-spec byte_size(dense_mat()) -> integer().
byte_size(M) ->
    erlang:bit_size(erlang:element(4, M)) div 8.

-file("src/viva_math/matrix_dense.gleam", 228).
-spec zip_data(
    bitstring(),
    bitstring(),
    fun((float(), float()) -> float()),
    bitstring()
) -> bitstring().
zip_data(A, B, F, Acc) ->
    case {A, B} of
        {<<X:64/float-little, Arest/bitstring>>,
            <<Y:64/float-little, Brest/bitstring>>} ->
            zip_data(
                Arest,
                Brest,
                F,
                <<Acc/bitstring, ((F(X, Y))):64/float-little>>
            );

        {_, _} ->
            Acc
    end.

-file("src/viva_math/matrix_dense.gleam", 211).
-spec zip_with(dense_mat(), dense_mat(), fun((float(), float()) -> float())) -> dense_mat().
zip_with(A, B, F) ->
    New_data = zip_data(erlang:element(4, A), erlang:element(4, B), F, <<>>),
    {dense_mat, erlang:element(2, A), erlang:element(3, A), New_data}.

-file("src/viva_math/matrix_dense.gleam", 182).
?DOC(" Element-wise add. Errors on shape mismatch.\n").
-spec add(dense_mat(), dense_mat()) -> {ok, dense_mat()} | {error, nil}.
add(A, B) ->
    case (erlang:element(2, A) /= erlang:element(2, B)) orelse (erlang:element(
        3,
        A
    )
    /= erlang:element(3, B)) of
        true ->
            {error, nil};

        false ->
            {ok, zip_with(A, B, fun(X, Y) -> X + Y end)}
    end.

-file("src/viva_math/matrix_dense.gleam", 190).
?DOC(" Element-wise subtract.\n").
-spec sub(dense_mat(), dense_mat()) -> {ok, dense_mat()} | {error, nil}.
sub(A, B) ->
    case (erlang:element(2, A) /= erlang:element(2, B)) orelse (erlang:element(
        3,
        A
    )
    /= erlang:element(3, B)) of
        true ->
            {error, nil};

        false ->
            {ok, zip_with(A, B, fun(X, Y) -> X - Y end)}
    end.

-file("src/viva_math/matrix_dense.gleam", 198).
?DOC(" Element-wise Hadamard product.\n").
-spec hadamard(dense_mat(), dense_mat()) -> {ok, dense_mat()} | {error, nil}.
hadamard(A, B) ->
    case (erlang:element(2, A) /= erlang:element(2, B)) orelse (erlang:element(
        3,
        A
    )
    /= erlang:element(3, B)) of
        true ->
            {error, nil};

        false ->
            {ok, zip_with(A, B, fun(X, Y) -> X * Y end)}
    end.

-file("src/viva_math/matrix_dense.gleam", 220).
-spec map_data(bitstring(), fun((float()) -> float()), bitstring()) -> bitstring().
map_data(Data, F, Acc) ->
    case Data of
        <<X:64/float-little, Rest/bitstring>> ->
            map_data(Rest, F, <<Acc/bitstring, ((F(X))):64/float-little>>);

        _ ->
            Acc
    end.

-file("src/viva_math/matrix_dense.gleam", 206).
?DOC(" Scalar multiplication.\n").
-spec scale(dense_mat(), float()) -> dense_mat().
scale(M, S) ->
    New_data = map_data(erlang:element(4, M), fun(X) -> X * S end, <<>>),
    {dense_mat, erlang:element(2, M), erlang:element(3, M), New_data}.

-file("src/viva_math/matrix_dense.gleam", 253).
-spec transpose_loop(dense_mat(), integer(), integer(), bitstring()) -> bitstring().
transpose_loop(M, New_row, New_col, Acc) ->
    case New_row >= erlang:element(3, M) of
        true ->
            Acc;

        false ->
            V = case get(M, New_col, New_row) of
                {ok, X} ->
                    X;

                {error, _} ->
                    +0.0
            end,
            Next_acc = <<Acc/bitstring, V:64/float-little>>,
            case (New_col + 1) >= erlang:element(2, M) of
                true ->
                    transpose_loop(M, New_row + 1, 0, Next_acc);

                false ->
                    transpose_loop(M, New_row, New_col + 1, Next_acc)
            end
    end.

-file("src/viva_math/matrix_dense.gleam", 248).
?DOC(" Transpose. O(rows · cols) but with contiguous writes.\n").
-spec transpose(dense_mat()) -> dense_mat().
transpose(M) ->
    New_data = transpose_loop(M, 0, 0, <<>>),
    {dense_mat, erlang:element(3, M), erlang:element(2, M), New_data}.

-file("src/viva_math/matrix_dense.gleam", 309).
-spec dot_row_col(
    dense_mat(),
    dense_mat(),
    integer(),
    integer(),
    integer(),
    float()
) -> float().
dot_row_col(A, B, I, J, K, Acc) ->
    case K >= erlang:element(3, A) of
        true ->
            Acc;

        false ->
            Av = case get(A, I, K) of
                {ok, X} ->
                    X;

                {error, _} ->
                    +0.0
            end,
            Bv = case get(B, K, J) of
                {ok, X@1} ->
                    X@1;

                {error, _} ->
                    +0.0
            end,
            dot_row_col(A, B, I, J, K + 1, Acc + (Av * Bv))
    end.

-file("src/viva_math/matrix_dense.gleam", 289).
-spec matmul_loop(dense_mat(), dense_mat(), integer(), integer(), bitstring()) -> bitstring().
matmul_loop(A, B, I, J, Acc) ->
    case I >= erlang:element(2, A) of
        true ->
            Acc;

        false ->
            V = dot_row_col(A, B, I, J, 0, +0.0),
            Next_acc = <<Acc/bitstring, V:64/float-little>>,
            case (J + 1) >= erlang:element(3, B) of
                true ->
                    matmul_loop(A, B, I + 1, 0, Next_acc);

                false ->
                    matmul_loop(A, B, I, J + 1, Next_acc)
            end
    end.

-file("src/viva_math/matrix_dense.gleam", 279).
?DOC(
    " Matrix product A · B. Errors on shape mismatch.\n"
    "\n"
    " Naïve triple-loop O(n³). For larger matrices defer to `viva_tensor` which\n"
    " dispatches to BLAS / cuBLAS.\n"
).
-spec mul(dense_mat(), dense_mat()) -> {ok, dense_mat()} | {error, nil}.
mul(A, B) ->
    case erlang:element(3, A) /= erlang:element(2, B) of
        true ->
            {error, nil};

        false ->
            New_data = matmul_loop(A, B, 0, 0, <<>>),
            {ok,
                {dense_mat,
                    erlang:element(2, A),
                    erlang:element(3, B),
                    New_data}}
    end.

-file("src/viva_math/matrix_dense.gleam", 338).
-spec sum_squared(bitstring(), float()) -> float().
sum_squared(Data, Acc) ->
    case Data of
        <<X:64/float-little, Rest/bitstring>> ->
            sum_squared(Rest, Acc + (X * X));

        _ ->
            Acc
    end.

-file("src/viva_math/matrix_dense.gleam", 334).
?DOC(" Frobenius norm √(Σ aᵢⱼ²).\n").
-spec frobenius(dense_mat()) -> float().
frobenius(M) ->
    math:sqrt(sum_squared(erlang:element(4, M), +0.0)).

-file("src/viva_math/matrix_dense.gleam", 353).
-spec trace_loop(dense_mat(), integer(), float()) -> float().
trace_loop(M, I, Acc) ->
    case I >= erlang:element(2, M) of
        true ->
            Acc;

        false ->
            case get(M, I, I) of
                {ok, V} ->
                    trace_loop(M, I + 1, Acc + V);

                {error, _} ->
                    Acc
            end
    end.

-file("src/viva_math/matrix_dense.gleam", 346).
?DOC(" Trace of a square matrix.\n").
-spec trace(dense_mat()) -> {ok, float()} | {error, nil}.
trace(M) ->
    case erlang:element(2, M) /= erlang:element(3, M) of
        true ->
            {error, nil};

        false ->
            {ok, trace_loop(M, 0, +0.0)}
    end.