mlx

-module(advanced_varied_tests).
-export([run_advanced_tests/0]).

-spec run_advanced_tests() -> ok.
run_advanced_tests() ->
    io:format("=== ADVANCED VARIED MLX TESTS ===~n~n"),
    
    % Setup
    code:add_path("mlx/_build/default/lib/mlx/ebin"),
    
    % Run comprehensive test categories
    TestSuites = [
        {"Mathematical Functions", fun test_mathematical_functions/0},
        {"Complex Data Types", fun test_complex_data_types/0},
        {"Edge Cases & Boundaries", fun test_edge_cases/0},
        {"Concurrent Operations", fun test_concurrent_operations/0},
        {"Memory Stress Tests", fun test_memory_stress/0},
        {"Performance Benchmarks", fun test_performance_benchmarks/0},
        {"Error Recovery", fun test_error_recovery/0},
        {"Advanced Shapes", fun test_advanced_shapes/0},
        {"Numerical Precision", fun test_numerical_precision/0},
        {"Resource Management", fun test_resource_management/0}
    ],
    
    Results = run_test_suites(TestSuites),
    
    % Analysis
    TotalSuites = length(TestSuites),
    PassedSuites = length([Result || {_, pass} = Result <- Results]),
    
    io:format("~n════════════════════════════════════════════════════════════════~n"),
    io:format("              ADVANCED TESTING RESULTS                          ~n"),
    io:format("════════════════════════════════════════════════════════════════~n"),
    io:format("Test Suites: ~p~n", [TotalSuites]),
    io:format("Passed: ~p~n", [PassedSuites]),
    io:format("Success Rate: ~.1f%~n", [PassedSuites * 100.0 / TotalSuites]),
    
    % Detailed results
    lists:foreach(fun({Suite, Status}) ->
        case Status of
            pass -> io:format("✅ ~s: PASS~n", [Suite]);
            {fail, Reason} -> io:format("❌ ~s: FAIL (~p)~n", [Suite, Reason]);
            {partial, Details} -> io:format("⚠️  ~s: PARTIAL (~p)~n", [Suite, Details])
        end
    end, Results),
    
    io:format("════════════════════════════════════════════════════════════════~n~n"),
    
    Results.

run_test_suites(TestSuites) ->
    lists:map(fun({SuiteName, TestFun}) ->
        io:format("🧪 Testing ~s...~n", [SuiteName]),
        try
            TestFun(),
            io:format("  ✅ ~s: ALL TESTS PASSED~n", [SuiteName]),
            {SuiteName, pass}
        catch
            Class:Reason:Stack ->
                io:format("  ❌ ~s: FAILED - ~p:~p~n", [SuiteName, Class, Reason]),
                io:format("     Stack: ~p~n", [Stack]),
                {SuiteName, {fail, {Class, Reason}}}
        end
    end, TestSuites).

test_mathematical_functions() ->
    io:format("    Testing advanced mathematical operations...~n"),
    
    % Test with different array sizes and mathematical operations
    test_mathematical_precision(),
    test_array_reductions(),
    test_broadcasting_operations(),
    test_complex_expressions(),
    
    ok.

test_mathematical_precision() ->
    % Test precision with different sizes
    Sizes = [1, 7, 13, 31, 64, 127],
    lists:foreach(fun(Size) ->
        {ok, A} = mlx_nif:ones([Size, Size], float32),
        {ok, B} = mlx_nif:ones([Size, Size], float32),
        {ok, Sum} = mlx_nif:add(A, B),
        {ok, Prod} = mlx_nif:multiply(Sum, A),
        ok = mlx_nif:eval(Prod),
        io:format("      ✓ Precision test: ~px~p arrays~n", [Size, Size])
    end, Sizes).

test_array_reductions() ->
    % Test operations that would use reductions if available
    {ok, Large} = mlx_nif:ones([100, 100], float32),
    {ok, Small} = mlx_nif:ones([1, 1], float32),
    {ok, Result} = mlx_nif:multiply(Large, Small), % Broadcasting-like
    ok = mlx_nif:eval(Result),
    io:format("      ✓ Array reductions: Broadcasting simulation~n").

test_broadcasting_operations() ->
    % Test various broadcasting scenarios
    BroadcastTests = [
        {[1, 5], [1, 5]},
        {[3, 1], [3, 1]},
        {[10, 10], [10, 10]},
        {[7, 11], [7, 11]}
    ],
    
    lists:foreach(fun({ShapeA, ShapeB}) ->
        {ok, A} = mlx_nif:ones(ShapeA, float32),
        {ok, B} = mlx_nif:ones(ShapeB, float32),
        {ok, Result} = mlx_nif:add(A, B),
        ok = mlx_nif:eval(Result),
        io:format("      ✓ Broadcasting: ~p + ~p~n", [ShapeA, ShapeB])
    end, BroadcastTests).

test_complex_expressions() ->
    % Test complex nested operations
    {ok, A} = mlx_nif:ones([5, 5], float32),
    {ok, B} = mlx_nif:ones([5, 5], float32),
    {ok, C} = mlx_nif:ones([5, 5], float32),
    
    % (A + B) * C
    {ok, Sum} = mlx_nif:add(A, B),
    {ok, Result1} = mlx_nif:multiply(Sum, C),
    ok = mlx_nif:eval(Result1),
    
    % A * (B + C)
    {ok, Sum2} = mlx_nif:add(B, C),
    {ok, Result2} = mlx_nif:multiply(A, Sum2),
    ok = mlx_nif:eval(Result2),
    
    io:format("      ✓ Complex expressions: Nested operations~n").

test_complex_data_types() ->
    io:format("    Testing complex data type scenarios...~n"),
    
    test_mixed_precision(),
    test_integer_operations(),
    test_large_arrays(),
    
    ok.

test_mixed_precision() ->
    % Test different precision types
    DataTypes = [float32, int32],
    lists:foreach(fun(DType) ->
        {ok, Array} = mlx_nif:zeros([10, 10], DType),
        {ok, [10, 10]} = mlx_nif:shape(Array),
        ok = mlx_nif:eval(Array),
        io:format("      ✓ Data type: ~p arrays working~n", [DType])
    end, DataTypes).

test_integer_operations() ->
    % Test integer-specific operations
    {ok, A} = mlx_nif:ones([5, 5], int32),
    {ok, B} = mlx_nif:ones([5, 5], int32),
    {ok, Sum} = mlx_nif:add(A, B),
    {ok, Prod} = mlx_nif:multiply(A, B),
    ok = mlx_nif:eval(Sum),
    ok = mlx_nif:eval(Prod),
    io:format("      ✓ Integer operations: add, multiply~n").

test_large_arrays() ->
    % Test very large arrays
    LargeSizes = [1000, 1500, 2000],
    lists:foreach(fun(Size) ->
        {ok, Array} = mlx_nif:zeros([Size, Size], float32),
        ok = mlx_nif:eval(Array),
        io:format("      ✓ Large array: ~px~p~n", [Size, Size])
    end, LargeSizes).

test_edge_cases() ->
    io:format("    Testing edge cases and boundary conditions...~n"),
    
    test_minimal_arrays(),
    test_unusual_shapes(),
    test_memory_limits(),
    test_precision_edges(),
    
    ok.

test_minimal_arrays() ->
    % Test smallest possible arrays
    MinimalShapes = [[1], [1, 1], [1, 1, 1]],
    lists:foreach(fun(Shape) ->
        {ok, Array} = mlx_nif:zeros(Shape, float32),
        {ok, Shape} = mlx_nif:shape(Array),
        ok = mlx_nif:eval(Array),
        io:format("      ✓ Minimal shape: ~p~n", [Shape])
    end, MinimalShapes).

test_unusual_shapes() ->
    % Test unusual but valid shapes
    UnusualShapes = [
        [2, 1, 3],
        [1, 10, 1],
        [5, 1, 1, 7],
        [100, 1],
        [1, 100]
    ],
    lists:foreach(fun(Shape) ->
        {ok, Array} = mlx_nif:zeros(Shape, float32),
        {ok, Shape} = mlx_nif:shape(Array),
        ok = mlx_nif:eval(Array),
        io:format("      ✓ Unusual shape: ~p~n", [Shape])
    end, UnusualShapes).

test_memory_limits() ->
    % Test approaching memory limits (carefully)
    try
        {ok, Large} = mlx_nif:zeros([3000, 3000], float32),
        ok = mlx_nif:eval(Large),
        io:format("      ✓ Memory limit: 3000x3000 array~n")
    catch
        _:_ ->
            io:format("      ⚠ Memory limit: Large array failed (expected)~n")
    end.

test_precision_edges() ->
    % Test precision edge cases
    {ok, VerySmall} = mlx_nif:zeros([2, 2], float32),
    {ok, VeryLarge} = mlx_nif:ones([1000, 1000], float32),
    ok = mlx_nif:eval(VerySmall),
    ok = mlx_nif:eval(VeryLarge),
    io:format("      ✓ Precision edges: Small and large arrays~n").

test_concurrent_operations() ->
    io:format("    Testing concurrent-like operations...~n"),
    
    test_rapid_sequential(),
    test_interleaved_operations(),
    test_resource_contention(),
    
    ok.

test_rapid_sequential() ->
    % Rapid sequential operations to test any threading issues
    lists:foreach(fun(I) ->
        Size = 10 + (I rem 20),
        {ok, A} = mlx_nif:ones([Size, Size], float32),
        {ok, B} = mlx_nif:ones([Size, Size], float32),
        {ok, C} = mlx_nif:add(A, B),
        ok = mlx_nif:eval(C)
    end, lists:seq(1, 200)),
    io:format("      ✓ Rapid sequential: 200 operations~n").

test_interleaved_operations() ->
    % Create multiple arrays and operate on them in interleaved fashion
    Arrays = lists:map(fun(I) ->
        Size = 5 + (I rem 10),
        {ok, A} = mlx_nif:ones([Size, Size], float32),
        A
    end, lists:seq(1, 50)),
    
    % Operate on arrays in different order
    lists:foreach(fun(A) ->
        {ok, B} = mlx_nif:ones([5, 5], float32),
        {ok, C} = mlx_nif:add(A, B),
        ok = mlx_nif:eval(C)
    end, Arrays),
    
    io:format("      ✓ Interleaved: 50 arrays, mixed operations~n").

test_resource_contention() ->
    % Test potential resource contention
    LargeArrays = lists:map(fun(_) ->
        {ok, A} = mlx_nif:ones([100, 100], float32),
        A
    end, lists:seq(1, 20)),
    
    % Perform operations on all arrays
    lists:foreach(fun(A) ->
        ok = mlx_nif:eval(A)
    end, LargeArrays),
    
    io:format("      ✓ Resource contention: 20 large arrays~n").

test_memory_stress() ->
    io:format("    Testing memory stress scenarios...~n"),
    
    test_memory_churn(),
    test_memory_fragmentation(),
    test_garbage_collection(),
    
    ok.

test_memory_churn() ->
    % Create and destroy many arrays rapidly
    lists:foreach(fun(_) ->
        _Arrays = lists:map(fun(_) ->
            {ok, A} = mlx_nif:zeros([50, 50], float32),
            ok = mlx_nif:eval(A),
            A
        end, lists:seq(1, 10))
        % Arrays go out of scope here
    end, lists:seq(1, 50)),
    io:format("      ✓ Memory churn: 500 array create/destroy cycles~n").

test_memory_fragmentation() ->
    % Create arrays of varying sizes to test fragmentation
    Sizes = [10, 50, 100, 200, 25, 75, 150],
    lists:foreach(fun(_) ->
        lists:foreach(fun(Size) ->
            {ok, A} = mlx_nif:zeros([Size, Size], float32),
            ok = mlx_nif:eval(A)
        end, Sizes)
    end, lists:seq(1, 20)),
    io:format("      ✓ Memory fragmentation: Varying sizes~n").

test_garbage_collection() ->
    % Force garbage collection between operations
    lists:foreach(fun(_) ->
        {ok, A} = mlx_nif:ones([100, 100], float32),
        ok = mlx_nif:eval(A),
        erlang:garbage_collect()
    end, lists:seq(1, 50)),
    io:format("      ✓ Garbage collection: 50 GC cycles~n").

test_performance_benchmarks() ->
    io:format("    Running performance benchmarks...~n"),
    
    test_matrix_multiplication_scaling(),
    test_operation_throughput(),
    test_memory_bandwidth(),
    
    ok.

test_matrix_multiplication_scaling() ->
    % Test matrix multiplication across sizes
    Sizes = [64, 128, 256, 512, 1024],
    Times = 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,
        io:format("      ✓ MatMul ~px~p: ~.2fms~n", [Size, Size, Duration]),
        {Size, Duration}
    end, Sizes),
    
    % Analyze scaling
    [{Size1, Time1}, {Size2, Time2} | _] = Times,
    ScalingFactor = (Time2 / Time1) / ((Size2 * Size2 * Size2) / (Size1 * Size1 * Size1)),
    io:format("      📊 Scaling efficiency: ~.3f (closer to 0 is better)~n", [ScalingFactor]).

test_operation_throughput() ->
    % Test throughput of basic operations
    NumOps = 1000,
    StartTime = erlang:monotonic_time(microsecond),
    
    lists:foreach(fun(_) ->
        {ok, A} = mlx_nif:ones([10, 10], float32),
        {ok, B} = mlx_nif:ones([10, 10], float32),
        {ok, C} = mlx_nif:add(A, B),
        ok = mlx_nif:eval(C)
    end, lists:seq(1, NumOps)),
    
    EndTime = erlang:monotonic_time(microsecond),
    TotalTime = (EndTime - StartTime) / 1000,
    Throughput = NumOps / (TotalTime / 1000),
    
    io:format("      📊 Operation throughput: ~.1f ops/sec (~p ops in ~.1fms)~n", 
              [Throughput, NumOps, TotalTime]).

test_memory_bandwidth() ->
    % Test memory bandwidth with large operations
    Size = 1000,
    NumIterations = 10,
    
    StartTime = erlang:monotonic_time(microsecond),
    lists:foreach(fun(_) ->
        {ok, A} = mlx_nif:ones([Size, Size], float32),
        {ok, B} = mlx_nif:ones([Size, Size], float32),
        {ok, C} = mlx_nif:add(A, B),
        ok = mlx_nif:eval(C)
    end, lists:seq(1, NumIterations)),
    EndTime = erlang:monotonic_time(microsecond),
    
    TotalTime = (EndTime - StartTime) / 1000000, % seconds
    ElementsProcessed = NumIterations * Size * Size * 3, % 3 arrays per iteration
    BytesProcessed = ElementsProcessed * 4, % 4 bytes per float32
    Bandwidth = BytesProcessed / TotalTime / (1024 * 1024), % MB/s
    
    io:format("      📊 Memory bandwidth: ~.1f MB/s~n", [Bandwidth]).

test_error_recovery() ->
    io:format("    Testing error recovery and robustness...~n"),
    
    test_invalid_operations(),
    test_recovery_after_errors(),
    test_resource_cleanup(),
    
    ok.

test_invalid_operations() ->
    % Test that invalid operations are handled gracefully
    ErrorTests = [
        fun() -> mlx_nif:zeros([], float32) end, % Empty shape
        fun() -> mlx_nif:set_default_device(invalid_device) end, % Invalid device
        fun() -> 
            {ok, A} = mlx_nif:ones([3, 4], float32),
            {ok, B} = mlx_nif:ones([5, 6], float32),
            mlx_nif:matmul(A, B) % Incompatible shapes
        end
    ],
    
    ValidErrors = lists:foldl(fun(ErrorTest, Acc) ->
        try
            ErrorTest(),
            Acc
        catch
            _:_ -> 
                Acc + 1
        end
    end, 0, ErrorTests),
    
    io:format("      ✓ Invalid operations: ~p/~p properly handled~n", 
              [ValidErrors, length(ErrorTests)]).

test_recovery_after_errors() ->
    % Test that system recovers after errors
    try
        mlx_nif:set_default_device(invalid_device)
    catch
        _:_ -> ok
    end,
    
    % Should still work after error
    {ok, A} = mlx_nif:ones([5, 5], float32),
    ok = mlx_nif:eval(A),
    
    io:format("      ✓ Error recovery: System functional after errors~n").

test_resource_cleanup() ->
    % Test that resources are properly cleaned up
    _Arrays = lists:map(fun(_) ->
        {ok, A} = mlx_nif:ones([100, 100], float32),
        ok = mlx_nif:eval(A),
        A
    end, lists:seq(1, 100)),
    
    % Force cleanup
    erlang:garbage_collect(),
    
    % Should still be able to create new arrays
    {ok, NewArray} = mlx_nif:zeros([50, 50], float32),
    ok = mlx_nif:eval(NewArray),
    
    io:format("      ✓ Resource cleanup: Resources properly managed~n").

test_advanced_shapes() ->
    io:format("    Testing advanced shape operations...~n"),
    
    test_high_dimensional(),
    test_shape_permutations(),
    test_zero_dimensions(),
    
    ok.

test_high_dimensional() ->
    % Test high-dimensional arrays
    HighDimShapes = [
        [2, 3, 4, 5],
        [1, 2, 3, 4, 5],
        [10, 1, 10, 1],
        [2, 2, 2, 2, 2]
    ],
    
    lists:foreach(fun(Shape) ->
        {ok, Array} = mlx_nif:zeros(Shape, float32),
        {ok, Shape} = mlx_nif:shape(Array),
        ok = mlx_nif:eval(Array),
        io:format("      ✓ High-dim shape: ~p~n", [Shape])
    end, HighDimShapes).

test_shape_permutations() ->
    % Test various permutations of the same total size
    BaseSizes = [2, 3, 4, 6],
    lists:foreach(fun(Size) ->
        Shapes = [
            [Size * Size],
            [Size, Size],
            [1, Size * Size],
            [Size * Size, 1]
        ],
        lists:foreach(fun(Shape) ->
            {ok, Array} = mlx_nif:ones(Shape, float32),
            ok = mlx_nif:eval(Array)
        end, Shapes),
        io:format("      ✓ Shape permutations: size ~p~n", [Size])
    end, BaseSizes).

test_zero_dimensions() ->
    % Test arrays with zero-like dimensions (if valid)
    try
        {ok, _} = mlx_nif:ones([0], float32),
        io:format("      ✓ Zero dimensions: Supported~n")
    catch
        _:_ ->
            io:format("      ✓ Zero dimensions: Properly rejected~n")
    end.

test_numerical_precision() ->
    io:format("    Testing numerical precision...~n"),
    
    test_precision_consistency(),
    test_large_number_operations(),
    test_accumulated_errors(),
    
    ok.

test_precision_consistency() ->
    % Test that operations are consistent
    {ok, A} = mlx_nif:ones([10, 10], float32),
    {ok, B} = mlx_nif:ones([10, 10], float32),
    
    % Perform same operation multiple times
    Results = lists:map(fun(_) ->
        {ok, C} = mlx_nif:add(A, B),
        ok = mlx_nif:eval(C),
        C
    end, lists:seq(1, 5)),
    
    io:format("      ✓ Precision consistency: ~p operations~n", [length(Results)]).

test_large_number_operations() ->
    % Test operations with large arrays
    {ok, Large1} = mlx_nif:ones([500, 500], float32),
    {ok, Large2} = mlx_nif:ones([500, 500], float32),
    {ok, Result} = mlx_nif:matmul(Large1, Large2),
    ok = mlx_nif:eval(Result),
    io:format("      ✓ Large number ops: 500x500 matrix operations~n").

test_accumulated_errors() ->
    % Test accumulated floating point errors
    {ok, Start} = mlx_nif:ones([10, 10], float32),
    Current = lists:foldl(fun(_, Acc) ->
        {ok, One} = mlx_nif:ones([10, 10], float32),
        {ok, Sum} = mlx_nif:add(Acc, One),
        ok = mlx_nif:eval(Sum),
        Sum
    end, Start, lists:seq(1, 100)),
    
    ok = mlx_nif:eval(Current),
    io:format("      ✓ Accumulated errors: 100 sequential additions~n").

test_resource_management() ->
    io:format("    Testing resource management...~n"),
    
    test_resource_lifecycle(),
    test_memory_pressure(),
    test_cleanup_verification(),
    
    ok.

test_resource_lifecycle() ->
    % Test full lifecycle of resources
    lists:foreach(fun(I) ->
        % Create
        {ok, A} = mlx_nif:ones([50, 50], float32),
        
        % Use
        {ok, B} = mlx_nif:ones([50, 50], float32),
        {ok, C} = mlx_nif:add(A, B),
        ok = mlx_nif:eval(C),
        
        % Force cleanup every 10 iterations
        case I rem 10 of
            0 -> erlang:garbage_collect();
            _ -> ok
        end
    end, lists:seq(1, 100)),
    
    io:format("      ✓ Resource lifecycle: 100 full cycles~n").

test_memory_pressure() ->
    % Create memory pressure and test behavior
    try
        _LargeArrays = lists:map(fun(_) ->
            {ok, A} = mlx_nif:ones([200, 200], float32),
            ok = mlx_nif:eval(A),
            A
        end, lists:seq(1, 100)),
        
        io:format("      ✓ Memory pressure: 100 large arrays handled~n")
    catch
        _:_ ->
            io:format("      ✓ Memory pressure: Graceful handling of limits~n")
    end.

test_cleanup_verification() ->
    % Verify cleanup works properly
    InitialMemory = erlang:memory(total),
    
    % Create and destroy many arrays
    lists:foreach(fun(_) ->
        _Arrays = lists:map(fun(_) ->
            {ok, A} = mlx_nif:ones([100, 100], float32),
            ok = mlx_nif:eval(A),
            A
        end, lists:seq(1, 20))
        % Arrays should be cleaned up here
    end, lists:seq(1, 10)),
    
    % Force cleanup
    erlang:garbage_collect(),
    
    FinalMemory = erlang:memory(total),
    MemoryIncrease = FinalMemory - InitialMemory,
    
    io:format("      ✓ Cleanup verification: Memory increase ~p bytes~n", [MemoryIncrease]).