-module(focused_test).
-export([test_actual_capabilities/0]).

test_actual_capabilities() ->
    io:format("=== Testing Actual MLX Capabilities ===~n~n"),
    
    % Add the MLX library to the code path
    code:add_path("mlx/_build/default/lib/mlx/ebin"),
    
    Tests = [
        {"Core NIF Functionality", fun test_core_nif/0},
        {"Array Creation & Operations", fun test_array_ops/0},
        {"Matrix Operations", fun test_matrix_ops/0},
        {"Device Management", fun test_device_mgmt/0},
        {"Performance Testing", fun test_performance_scaling/0},
        {"Error Handling", fun test_error_handling/0},
        {"Memory Management", fun test_memory_mgmt/0},
        {"Advanced Array Features", fun test_advanced_arrays/0}
    ],
    
    Results = lists:map(fun({TestName, TestFun}) ->
        io:format("~n--- Testing ~s ---~n", [TestName]),
        try
            TestFun(),
            io:format("✓ ~s: PASS~n", [TestName]),
            {TestName, pass}
        catch
            Class:Reason:Stack ->
                io:format("✗ ~s: FAIL - ~p:~p~n", [TestName, Class, Reason]),
                io:format("  Stack: ~p~n", [Stack]),
                {TestName, {fail, {Class, Reason}}}
        end
    end, Tests),
    
    % Summary
    PassedTests = length([Result || {_, pass} = Result <- Results]),
    TotalTests = length(Results),
    
    io:format("~n=== Final Results ===~n"),
    io:format("Passed: ~p/~p test categories~n", [PassedTests, TotalTests]),
    
    case PassedTests of
        TotalTests ->
            io:format("🎉 ALL ACTUAL CAPABILITIES WORKING!~n");
        _ ->
            io:format("⚠️  Some capabilities failed:~n"),
            [io:format("  - ~s~n", [Name]) || {Name, Status} <- Results, Status =/= pass]
    end,
    
    Results.

test_core_nif() ->
    io:format("  Testing core NIF loading and basic functions...~n"),
    
    % Test NIF version
    {ok, mlx_loaded} = mlx_nif:version(),
    io:format("    ✓ NIF loaded successfully~n"),
    
    % Test basic functionality
    {ok, {test_passed, Shape, Dtype}} = mlx_nif:test_basic(),
    io:format("    ✓ Basic test passed: Shape=~p, Dtype=~p~n", [Shape, Dtype]),
    
    ok.

test_array_ops() ->
    io:format("  Testing array creation and basic operations...~n"),
    
    % Test different data types
    DataTypes = [float32, int32],
    lists:foreach(fun(DType) ->
        {ok, A} = mlx_nif:zeros([3, 4], DType),
        {ok, [3, 4]} = mlx_nif:shape(A),
        io:format("    ✓ ~p arrays: creation and shape~n", [DType])
    end, DataTypes),
    
    % Test arithmetic operations
    {ok, A} = mlx_nif:ones([3, 3], float32),
    {ok, B} = mlx_nif:ones([3, 3], float32),
    {ok, Sum} = mlx_nif:add(A, B),
    {ok, Prod} = mlx_nif:multiply(A, B),
    
    % Verify shapes
    {ok, [3, 3]} = mlx_nif:shape(Sum),
    {ok, [3, 3]} = mlx_nif:shape(Prod),
    
    % Test evaluation
    ok = mlx_nif:eval(Sum),
    ok = mlx_nif:eval(Prod),
    
    io:format("    ✓ Arithmetic operations: add, multiply~n"),
    io:format("    ✓ Lazy evaluation working~n"),
    
    ok.

test_matrix_ops() ->
    io:format("  Testing matrix operations...~n"),
    
    % Test various matrix multiplication sizes
    TestSizes = [
        {[2, 3], [3, 4], [2, 4]},
        {[5, 7], [7, 3], [5, 3]},
        {[10, 8], [8, 12], [10, 12]}
    ],
    
    lists:foreach(fun({ShapeA, ShapeB, ExpectedShape}) ->
        {ok, A} = mlx_nif:ones(ShapeA, float32),
        {ok, B} = mlx_nif:ones(ShapeB, float32),
        {ok, C} = mlx_nif:matmul(A, B),
        {ok, ExpectedShape} = mlx_nif:shape(C),
        ok = mlx_nif:eval(C),
        io:format("    ✓ Matrix multiply: ~p × ~p = ~p~n", [ShapeA, ShapeB, ExpectedShape])
    end, TestSizes),
    
    ok.

test_device_mgmt() ->
    io:format("  Testing device management...~n"),
    
    % Test CPU device
    ok = mlx_nif:set_default_device(cpu),
    {ok, A_cpu} = mlx_nif:zeros([10, 10], float32),
    ok = mlx_nif:eval(A_cpu),
    io:format("    ✓ CPU device working~n"),
    
    % Test GPU device (may not be available)
    case mlx_nif:set_default_device(gpu) of
        ok ->
            {ok, A_gpu} = mlx_nif:ones([10, 10], float32),
            ok = mlx_nif:eval(A_gpu),
            io:format("    ✓ GPU device available and working~n"),
            % Switch back to CPU
            ok = mlx_nif:set_default_device(cpu);
        {error, Reason} ->
            io:format("    ✓ GPU device not available (~p) - CPU fallback working~n", [Reason])
    end,
    
    ok.

test_performance_scaling() ->
    io:format("  Testing performance scaling...~n"),
    
    % Test performance with increasing sizes
    Sizes = [50, 100, 200, 500, 1000],
    
    Results = lists:map(fun(Size) ->
        StartTime = erlang:monotonic_time(microsecond),
        
        {ok, A} = mlx_nif:ones([Size, Size], float32),
        {ok, B} = mlx_nif:ones([Size, Size], float32),
        {ok, C} = mlx_nif:matmul(A, B),
        ok = mlx_nif:eval(C),
        
        EndTime = erlang:monotonic_time(microsecond),
        Duration = (EndTime - StartTime) / 1000, % Convert to milliseconds
        
        io:format("    ✓ ~px~p matrix multiply: ~.2fms~n", [Size, Size, Duration]),
        {Size, Duration}
    end, Sizes),
    
    % Check that performance scales reasonably
    [Size1, Size2, Size3, Size4, Size5] = Sizes,
    [{_, Time1}, {_, Time2}, {_, Time3}, {_, Time4}, {_, Time5}] = Results,
    
    % Larger operations should generally take more time (with some variance)
    io:format("    ✓ Performance scaling analysis:~n"),
    io:format("      ~px~p: ~.2fms (baseline)~n", [Size1, Size1, Time1]),
    io:format("      ~px~p: ~.2fms (~.1fx)~n", [Size5, Size5, Time5, Time5/Time1]),
    
    ok.

test_error_handling() ->
    io:format("  Testing error handling...~n"),
    
    % Test invalid shapes
    case mlx_nif:zeros([-1, 5], float32) of
        {error, _} -> 
            io:format("    ✓ Invalid shape properly rejected~n");
        _ -> 
            io:format("    ? Invalid shape not caught (may be valid in MLX)~n")
    end,
    
    % Test incompatible matrix multiplication
    {ok, A} = mlx_nif:ones([3, 4], float32),
    {ok, B} = mlx_nif:ones([5, 6], float32), % Incompatible with A
    
    case mlx_nif:matmul(A, B) of
        {error, _} -> 
            io:format("    ✓ Incompatible matmul properly rejected~n");
        _ -> 
            io:format("    ? Incompatible matmul not caught (may be broadcasting)~n")
    end,
    
    % Test invalid device
    case mlx_nif:set_default_device(invalid_device) of
        {error, _} -> 
            io:format("    ✓ Invalid device properly rejected~n");
        _ -> 
            io:format("    ? Invalid device not caught~n")
    end,
    
    ok.

test_memory_mgmt() ->
    io:format("  Testing memory management...~n"),
    
    % Create many arrays to test memory handling
    NumArrays = 100,
    Arrays = lists:map(fun(I) ->
        {ok, A} = mlx_nif:zeros([10, 10], float32),
        ok = mlx_nif:eval(A),
        A
    end, lists:seq(1, NumArrays)),
    
    io:format("    ✓ Created ~p arrays without memory issues~n", [NumArrays]),
    
    % Test large array
    {ok, LargeArray} = mlx_nif:zeros([2000, 2000], float32),
    ok = mlx_nif:eval(LargeArray),
    io:format("    ✓ Large array (2000x2000) handled successfully~n"),
    
    % Force garbage collection
    erlang:garbage_collect(),
    io:format("    ✓ Garbage collection completed~n"),
    
    ok.

test_advanced_arrays() ->
    io:format("  Testing advanced array features...~n"),
    
    % Test different shapes and sizes
    Shapes = [
        [1],           % 1D
        [5, 5],        % 2D square
        [3, 7],        % 2D rectangular
        [2, 3, 4],     % 3D
        [2, 2, 2, 2]   % 4D
    ],
    
    lists:foreach(fun(Shape) ->
        {ok, A} = mlx_nif:zeros(Shape, float32),
        {ok, Shape} = mlx_nif:shape(A),
        ok = mlx_nif:eval(A),
        io:format("    ✓ Shape ~p: creation and evaluation~n", [Shape])
    end, Shapes),
    
    % Test chain of operations
    {ok, A} = mlx_nif:ones([5, 5], float32),
    {ok, B} = mlx_nif:ones([5, 5], float32),
    {ok, C} = mlx_nif:add(A, B),
    {ok, D} = mlx_nif:multiply(C, B),
    {ok, E} = mlx_nif:add(D, A),
    ok = mlx_nif:eval(E),
    io:format("    ✓ Chained operations: A + B -> (result * B) + A~n"),
    
    ok.