mlx

-module(test_runner).
-export([run_basic_tests/0]).

run_basic_tests() ->
    io:format("=== MLX Test Suite Runner ===~n~n"),
    
    % Add the MLX library to the code path
    code:add_path("mlx/_build/default/lib/mlx/ebin"),
    
    % Test results will be collected
    
    % Test 1: Basic array operations
    io:format("Test 1: Basic Array Operations...~n"),
    Test1 = try
        {ok, A} = mlx:zeros([3, 3]),
        {ok, B} = mlx:ones([3, 3]),
        {ok, C} = mlx:add(A, B),
        {ok, [3, 3]} = mlx:shape(C),
        ok = mlx:eval(C),
        io:format("  ✓ Array creation, addition, and evaluation: PASS~n"),
        pass
    catch
        Class:Reason:_ ->
            io:format("  ✗ Array operations: FAIL - ~p:~p~n", [Class, Reason]),
            fail
    end,
    
    % Test 2: Matrix multiplication
    io:format("Test 2: Matrix Multiplication...~n"),
    Test2 = try
        {ok, A2} = mlx:ones([3, 4]),
        {ok, B2} = mlx:ones([4, 2]),
        {ok, C2} = mlx:matmul(A2, B2),
        {ok, [3, 2]} = mlx:shape(C2),
        ok = mlx:eval(C2),
        io:format("  ✓ Matrix multiplication: PASS~n"),
        pass
    catch
        Class2:Reason2:_ ->
            io:format("  ✗ Matrix multiplication: FAIL - ~p:~p~n", [Class2, Reason2]),
            fail
    end,
    
    % Test 3: Device switching
    io:format("Test 3: Device Management...~n"),
    Test3 = try
        ok = mlx:set_default_device(cpu),
        {ok, A3} = mlx:zeros([10, 10]),
        ok = mlx:eval(A3),
        % Try GPU (may not be available)
        case mlx:set_default_device(gpu) of
            ok ->
                {ok, B3} = mlx:ones([10, 10]),
                ok = mlx:eval(B3),
                ok = mlx:set_default_device(cpu),
                io:format("  ✓ Device switching (GPU available): PASS~n");
            {error, _} ->
                io:format("  ✓ Device switching (GPU not available): PASS~n")
        end,
        pass
    catch
        Class3:Reason3:_ ->
            io:format("  ✗ Device management: FAIL - ~p:~p~n", [Class3, Reason3]),
            fail
    end,
    
    % Test 4: Large array performance
    io:format("Test 4: Large Array Performance...~n"),
    Test4 = try
        Size = 500,  % Smaller than 1000 for faster testing
        StartTime = erlang:monotonic_time(millisecond),
        {ok, A4} = mlx:ones([Size, Size]),
        {ok, B4} = mlx:ones([Size, Size]),
        {ok, C4} = mlx:add(A4, B4),
        ok = mlx:eval(C4),
        EndTime = erlang:monotonic_time(millisecond),
        Duration = EndTime - StartTime,
        io:format("  ✓ Large array operation (~px~p) took ~pms: PASS~n", [Size, Size, Duration]),
        pass
    catch
        Class4:Reason4:_ ->
            io:format("  ✗ Large array performance: FAIL - ~p:~p~n", [Class4, Reason4]),
            fail
    end,
    
    % Test 5: MLX NIF version check
    io:format("Test 5: MLX Version Check...~n"),
    Test5 = try
        {ok, mlx_loaded} = mlx_nif:version(),
        {ok, {test_passed, Shape5, Dtype5}} = mlx_nif:test_basic(),
        io:format("  ✓ MLX version check and basic test: PASS~n"),
        io:format("    Basic test result: Shape=~p, Dtype=~p~n", [Shape5, Dtype5]),
        pass
    catch
        Class5:Reason5:_ ->
            io:format("  ✗ MLX version check: FAIL - ~p:~p~n", [Class5, Reason5]),
            fail
    end,
    
    % Calculate results
    AllTests = [Test1, Test2, Test3, Test4, Test5],
    PassedTests = length([T || T <- AllTests, T =:= pass]),
    TotalTests = length(AllTests),
    
    io:format("~n=== Test Results ===~n"),
    io:format("Passed: ~p/~p tests~n", [PassedTests, TotalTests]),
    
    case PassedTests of
        TotalTests ->
            io:format("🎉 ALL TESTS PASSED! MLX is working correctly.~n"),
            success;
        _ ->
            io:format("⚠️  Some tests failed. MLX may have issues.~n"),
            {partial_success, PassedTests, TotalTests}
    end.