mlx

%%%-------------------------------------------------------------------
%% @doc Working MLX demo with proper paths
%% @end
%%%-------------------------------------------------------------------

-module(working_demo).

-export([demo/0]).

%% @doc Run MLX demo
-spec demo() -> ok.
demo() ->
    io:format("=== MLX Erlang Demo ===~n~n"),
    
    % Add the compiled MLX modules to the path
    code:add_path("mlx/_build/default/lib/mlx/ebin"),
    
    io:format("1. Testing MLX zeros function:~n"),
    test_zeros(),
    
    io:format("~n2. Testing MLX ones function:~n"),
    test_ones(),
    
    io:format("~n3. Testing basic operations:~n"),
    test_operations(),
    
    io:format("~n4. Testing device management:~n"),
    test_devices(),
    
    io:format("~nDemo complete! MLX is working properly.~n").

test_zeros() ->
    case mlx:zeros([3, 3]) of
        {ok, Array} ->
            case mlx:shape(Array) of
                {ok, Shape} ->
                    io:format("✓ Created zeros array with shape: ~p~n", [Shape]);
                Error ->
                    io:format("✗ Shape error: ~p~n", [Error])
            end;
        Error ->
            io:format("✗ Zeros error: ~p~n", [Error])
    end.

test_ones() ->
    case mlx:ones([2, 4], int32) of
        {ok, Array} ->
            case mlx:shape(Array) of
                {ok, Shape} ->
                    io:format("✓ Created ones array with shape: ~p~n", [Shape]);
                Error ->
                    io:format("✗ Shape error: ~p~n", [Error])
            end;
        Error ->
            io:format("✗ Ones error: ~p~n", [Error])
    end.

test_operations() ->
    case {mlx:ones([3, 3]), mlx:ones([3, 3])} of
        {{ok, A}, {ok, B}} ->
            case mlx:add(A, B) of
                {ok, Sum} ->
                    case mlx:shape(Sum) of
                        {ok, Shape} ->
                            io:format("✓ Addition result shape: ~p~n", [Shape]),
                            case mlx:eval(Sum) of
                                ok ->
                                    io:format("✓ Array evaluated successfully~n");
                                Error ->
                                    io:format("✗ Eval error: ~p~n", [Error])
                            end;
                        Error ->
                            io:format("✗ Shape error: ~p~n", [Error])
                    end;
                Error ->
                    io:format("✗ Addition error: ~p~n", [Error])
            end;
        Error ->
            io:format("✗ Array creation error: ~p~n", [Error])
    end.

test_devices() ->
    case mlx:use_cpu() of
        ok ->
            io:format("✓ Switched to CPU device~n"),
            case mlx:use_gpu() of
                ok ->
                    io:format("✓ Switched to GPU device~n");
                {error, Reason} ->
                    io:format("⚠ Could not switch to GPU: ~p (this is normal if no GPU available)~n", [Reason])
            end;
        Error ->
            io:format("✗ Device switching error: ~p~n", [Error])
    end.