mlx

-module(erlang_matrix).
-export([
    zeros/1, ones/1, 
    add/2, multiply/2, matmul/2,
    sum/1, mean/1, transpose/1,
    dot/2, sigmoid/1, softmax/1,
    create_random/1
]).

%% Create a zeros matrix
zeros([Rows, Cols]) ->
    [[0.0 || _ <- lists:seq(1, Cols)] || _ <- lists:seq(1, Rows)].

%% Create a ones matrix  
ones([Rows, Cols]) ->
    [[1.0 || _ <- lists:seq(1, Cols)] || _ <- lists:seq(1, Rows)].

%% Create random matrix
create_random([Rows, Cols]) ->
    [[rand:uniform() || _ <- lists:seq(1, Cols)] || _ <- lists:seq(1, Rows)].

%% Element-wise addition
add(A, B) when is_list(A), is_list(B) ->
    lists:zipwith(fun(RowA, RowB) ->
        lists:zipwith(fun(X, Y) -> X + Y end, RowA, RowB)
    end, A, B).

%% Element-wise multiplication
multiply(A, B) when is_list(A), is_list(B) ->
    lists:zipwith(fun(RowA, RowB) ->
        lists:zipwith(fun(X, Y) -> X * Y end, RowA, RowB)
    end, A, B).

%% Matrix multiplication
matmul(A, B) ->
    BT = transpose(B),
    [[dot_vectors(RowA, ColB) || ColB <- BT] || RowA <- A].

dot_vectors(V1, V2) ->
    lists:sum(lists:zipwith(fun(X, Y) -> X * Y end, V1, V2)).

%% Vector dot product
dot(V1, V2) when is_list(V1), is_list(V2) ->
    dot_vectors(V1, V2).

%% Transpose
transpose([[]|_]) -> [];
transpose(Matrix) ->
    [lists:map(fun hd/1, Matrix) | transpose(lists:map(fun tl/1, Matrix))].

%% Sum all elements
sum(Matrix) when is_list(hd(Matrix)) ->
    lists:sum([lists:sum(Row) || Row <- Matrix]);
sum(Vector) ->
    lists:sum(Vector).

%% Mean of all elements
mean(Matrix) when is_list(hd(Matrix)) ->
    Total = sum(Matrix),
    Count = length(Matrix) * length(hd(Matrix)),
    Total / Count;
mean(Vector) ->
    sum(Vector) / length(Vector).

%% Sigmoid activation
sigmoid(Matrix) when is_list(hd(Matrix)) ->
    [[1.0 / (1.0 + math:exp(-X)) || X <- Row] || Row <- Matrix];
sigmoid(Vector) ->
    [1.0 / (1.0 + math:exp(-X)) || X <- Vector].

%% Softmax (simplified for vectors)
softmax(Vector) ->
    ExpVals = [math:exp(X) || X <- Vector],
    Sum = lists:sum(ExpVals),
    [X / Sum || X <- ExpVals].