NB. ===========================================================================
NB. futhark_test.ijs - Test Suite for Futhark-J Bridge
NB. ===========================================================================
NB.
NB. Run with: load 'futhark_test.ijs'
NB. ===========================================================================

load '/path/to/futhark.ijs'

NB. ===========================================================================
NB. Test Framework
NB. ===========================================================================

TESTS_RUN =: 0
TESTS_PASSED =: 0
TESTS_FAILED =: 0

assert_equal =: 4 : 0
  'name expected' =. x
  actual =. y
  TESTS_RUN =: TESTS_RUN + 1
  if. expected -: actual do.
    TESTS_PASSED =: TESTS_PASSED + 1
    echo '  PASS: ', name
  else.
    TESTS_FAILED =: TESTS_FAILED + 1
    echo '  FAIL: ', name
    echo '    Expected: ', ": expected
    echo '    Actual:   ', ": actual
  end.
)

assert_close =: 4 : 0
  'name expected tol' =. x
  actual =. y
  TESTS_RUN =: TESTS_RUN + 1
  diff =. | expected - actual
  if. +./ , diff < tol do.
    TESTS_PASSED =: TESTS_PASSED + 1
    echo '  PASS: ', name
  else.
    TESTS_FAILED =: TESTS_FAILED + 1
    echo '  FAIL: ', name
    echo '    Expected: ', ": expected
    echo '    Actual:   ', ": actual
  end.
)

assert_true =: 4 : 0
  name =. x
  cond =. y
  TESTS_RUN =: TESTS_RUN + 1
  if. cond do.
    TESTS_PASSED =: TESTS_PASSED + 1
    echo '  PASS: ', name
  else.
    TESTS_FAILED =: TESTS_FAILED + 1
    echo '  FAIL: ', name
  end.
)

NB. ===========================================================================
NB. Pre-compile all test functions (avoids multiple compilations during tests)
NB. ===========================================================================

echo 'Pre-compiling Futhark functions...'

NB. Scalar operations
t_sum =: 'entry f (xs: []f64) : f64 = reduce (+) 0 xs' fut
t_prod =: 'entry f (xs: []f64) : f64 = reduce (*) 1 xs' fut
t_min =: 'entry f (xs: []f64) : f64 = reduce f64.min f64.highest xs' fut
t_max =: 'entry f (xs: []f64) : f64 = reduce f64.max f64.lowest xs' fut
t_avg =: 'entry f (xs: []f64) : f64 = reduce (+) 0 xs / f64.i64 (length xs)' fut

NB. Array operations
t_double =: 'entry f (xs: []f64) : []f64 = map (*2) xs' fut
t_square =: 'entry f (xs: []f64) : []f64 = map (\x -> x * x) xs' fut
t_neg =: 'entry f (xs: []f64) : []f64 = map (\x -> 0 - x) xs' fut
t_cumsum =: 'entry f (xs: []f64) : []f64 = scan (+) 0 xs' fut
t_runmax =: 'entry f (xs: []f64) : []f64 = scan f64.max f64.lowest xs' fut
t_id =: 'entry f (xs: []f64) : []f64 = map (\x -> x) xs' fut

echo 'Compilation complete.'

NB. ===========================================================================
NB. Run Tests
NB. ===========================================================================

echo ''
echo '+===============================================================+'
echo '|          FUTHARK-J BRIDGE TEST SUITE                         |'
echo '+===============================================================+'
echo ''

NB. -----------------------------------------------------------------
echo '=== Compilation Tests ==='
echo ''

test_code =. 'entry test (xs: []f64) : []f64 = map (*2) xs'
so_file =. fut_compile test_code
'Compilation produces .so file' assert_true (# so_file) > 0
'Compiled file exists' assert_true fexist so_file

t1 =. 6!:2 'so_file2 =. fut_compile test_code'
'Cached compilation is fast' assert_true t1 < 0.01

NB. -----------------------------------------------------------------
echo ''
echo '=== Scalar Return Operations ==='
echo ''

('Sum of 1-5';15) assert_equal t_sum 1 2 3 4 5
('Sum of single element';42) assert_equal t_sum ,42
('Product of 1-5';120) assert_equal t_prod 1 2 3 4 5
('Min of 5 3 8 1 9';1) assert_equal t_min 5 3 8 1 9
('Max of 5 3 8 1 9';9) assert_equal t_max 5 3 8 1 9
('Average of 1-5';3;1e_10) assert_close t_avg 1 2 3 4 5

NB. -----------------------------------------------------------------
echo ''
echo '=== Array Return Operations ==='
echo ''

('Double 1 2 3 4 5';2 4 6 8 10) assert_equal t_double 1 2 3 4 5
('Square 1 2 3 4 5';1 4 9 16 25) assert_equal t_square 1 2 3 4 5
('Negate 1 2 3';_1 _2 _3) assert_equal t_neg 1 2 3
('Cumsum 1 2 3 4 5';1 3 6 10 15) assert_equal t_cumsum 1 2 3 4 5
('Running max';3 3 4 4 5 9 9 9 9 9) assert_equal t_runmax 3 1 4 1 5 9 2 6 5 3

NB. -----------------------------------------------------------------
echo ''
echo '=== Edge Cases ==='
echo ''

('Single element sum';42) assert_equal t_sum ,42
('Single element double';84) assert_equal 0{t_double ,42
('Negative numbers';_2 _4 _6) assert_equal t_double _1 _2 _3
('Zero in array';0 2 0 4 0) assert_equal t_double 0 1 0 2 0

small_vals =. 1e_10 * 1 2 3
('Small values preserved';small_vals;1e_15) assert_close t_id small_vals

NB. -----------------------------------------------------------------
echo ''
echo '=== Large Array Handling ==='
echo ''

large =. ? 10000 # 0
j_sum_large =. +/ large
('Large array sum';j_sum_large;1e_6) assert_close t_sum large

j_double_large =. 2 * large
('Large array double';j_double_large;1e_10) assert_close t_double large

NB. -----------------------------------------------------------------
echo ''
echo '=== Type Mapping Tests ==='
echo ''

('f64 maps to d';'d') assert_equal fut_type_to_j 'f64'
('f32 maps to f';'f') assert_equal fut_type_to_j 'f32'
('i64 maps to x';'x') assert_equal fut_type_to_j 'i64'
('i32 maps to i';'i') assert_equal fut_type_to_j 'i32'

NB. -----------------------------------------------------------------
echo ''
echo '=== Metadata Parsing Tests ==='
echo ''

code1 =. 'entry myfunction (xs: []f64) : f64 = reduce (+) 0 xs'
('Parse entry name';'myfunction') assert_equal fut_parse_entry code1
'Detects scalar return' assert_true 0 = fut_returns_array code1

code2 =. 'entry foo (xs: []f64) : []f64 = map (*2) xs'
'Detects array return' assert_true 1 = fut_returns_array code2

('Parse f64 type';'f64') assert_equal fut_parse_dtype code1
code3 =. 'entry foo (xs: []i32) : []i32 = map (*2) xs'
('Parse i32 type';'i32') assert_equal fut_parse_dtype code3

NB. -----------------------------------------------------------------
echo ''
echo '=== Multi-dimensional Array Tests ==='
echo ''

NB. 2D transpose
t_transp =: 'entry f (xs: [][]f64) : [][]f64 = transpose xs' fut
mat =: 3 4 $ 1 2 3 4 5 6 7 8 9 10 11 12
transp_result =: t_transp mat
('2D transpose shape';4 3) assert_equal $ transp_result
('2D transpose values';|: mat) assert_equal transp_result

NB. 2D flatten
t_flat =: 'entry f (xs: [][]f64) : []f64 = flatten xs' fut
flat_result =: t_flat mat
('Flatten 2D to 1D';, mat) assert_equal flat_result

NB. -----------------------------------------------------------------
echo ''
echo '=== Different Input/Output Size Tests ==='
echo ''

NB. Take (shrink)
t_take3 =: 'entry f (xs: []f64) : []f64 = take 3 xs' fut
('Take first 3';0 1 2) assert_equal t_take3 i. 10

NB. Drop (shrink)
t_drop5 =: 'entry f (xs: []f64) : []f64 = drop 5 xs' fut
('Drop first 5';5 6 7 8 9) assert_equal t_drop5 i. 10

NB. Replicate (grow)
t_rep2 =: 'entry f (xs: []f64) : []f64 = flatten (map (\x -> replicate 2 x) xs)' fut
('Replicate 2x';1 1 2 2 3 3) assert_equal t_rep2 1 2 3

NB. Concat (grow)
t_concat =: 'entry f (xs: []f64) : []f64 = concat xs xs' fut
('Concat self';1 2 3 1 2 3) assert_equal t_concat 1 2 3

NB. -----------------------------------------------------------------
echo ''
echo '=== Extended Parsing Tests ==='
echo ''

code_2d =. 'entry f (xs: [][]f64) : [][]f64 = transpose xs'
('Parse 2D input rank';2) assert_equal fut_input_rank code_2d
('Parse 2D output rank';2) assert_equal fut_output_rank code_2d

code_flat =. 'entry f (xs: [][]f64) : []f64 = flatten xs'
('Parse flatten input rank';2) assert_equal fut_input_rank code_flat
('Parse flatten output rank';1) assert_equal fut_output_rank code_flat

('Extract basetype from []f64';'f64') assert_equal fut_extract_basetype '[]f64'
('Extract basetype from [][]f64';'f64') assert_equal fut_extract_basetype '[][]f64'
('Extract basetype from f64';'f64') assert_equal fut_extract_basetype 'f64'

NB. 3D parsing tests
code_3d =. 'entry f (xs: [][][]f64) : [][][]f64 = xs'
('Parse 3D input rank';3) assert_equal fut_input_rank code_3d
('Parse 3D output rank';3) assert_equal fut_output_rank code_3d

code_3d_to_1d =. 'entry f (xs: [][][]f64) : []f64 = flatten (flatten xs)'
('Parse 3D->1D input rank';3) assert_equal fut_input_rank code_3d_to_1d
('Parse 3D->1D output rank';1) assert_equal fut_output_rank code_3d_to_1d

NB. -----------------------------------------------------------------
echo ''
echo '=== 3D Array Tests ==='
echo ''

NB. Create 3D test array
arr3d =: 2 3 4 $ i. 24

NB. 3D identity
t_id3d =: 'entry f (xs: [][][]f64) : [][][]f64 = xs' fut
id3d_result =: t_id3d arr3d
('3D identity shape';$ arr3d) assert_equal $ id3d_result
('3D identity values';arr3d) assert_equal id3d_result

NB. 3D map
t_map3d =: 'entry f (xs: [][][]f64) : [][][]f64 = map (map (map (*2))) xs' fut
map3d_result =: t_map3d arr3d
('3D map shape';$ arr3d) assert_equal $ map3d_result
('3D map values';2*arr3d) assert_equal map3d_result

NB. 3D flatten to 1D
t_flat3d_1d =: 'entry f (xs: [][][]f64) : []f64 = flatten (flatten xs)' fut
flat3d_1d_result =: t_flat3d_1d arr3d
('3D->1D flatten'; ,arr3d) assert_equal flat3d_1d_result

NB. 3D flatten to 2D
t_flat3d_2d =: 'entry f (xs: [][][]f64) : [][]f64 = flatten xs' fut
flat3d_2d_result =: t_flat3d_2d arr3d
('3D->2D flatten shape';6 4) assert_equal $ flat3d_2d_result
('3D->2D flatten values'; ,arr3d) assert_equal , flat3d_2d_result

NB. 3D scalar reduction
t_sum3d =: 'entry f (xs: [][][]f64) : f64 = flatten (flatten xs) |> reduce (+) 0' fut
('3D sum';+/,arr3d) assert_equal t_sum3d arr3d

NB. -----------------------------------------------------------------
echo ''
echo '=== J Native Verb Comparison Tests ==='
echo ''

NB. Compare Futhark operations with equivalent J verbs
NB. These tests verify semantic equivalence

NB. 1D operations
test_1d =: 1 2 3 4 5 6 7 8 9 10

NB. Sum: Futhark reduce (+) 0 vs J +/
('Futhark sum = J +/';+/ test_1d) assert_equal t_sum test_1d

NB. Product: Futhark reduce (*) 1 vs J */
('Futhark product = J */';*/ test_1d) assert_equal t_prod test_1d

NB. Double: Futhark map (*2) vs J 2&*
('Futhark double = J 2&*';2*test_1d) assert_equal t_double test_1d

NB. Square: Futhark map (\x -> x*x) vs J *:
('Futhark square = J *:';*: test_1d) assert_equal t_square test_1d

NB. Cumsum: Futhark scan (+) 0 vs J +/\
('Futhark cumsum = J +/\\';+/\ test_1d) assert_equal t_cumsum test_1d

NB. 2D operations
test_2d =: 3 4 $ i. 12

NB. Transpose: Futhark transpose vs J |:
('Futhark transpose = J |:';|: test_2d) assert_equal t_transp test_2d

NB. Flatten: Futhark flatten vs J ,
('Futhark flatten = J ,';, test_2d) assert_equal t_flat test_2d

NB. 3D operations
test_3d =: 2 3 4 $ i. 24

NB. 3D identity vs J ]
('Futhark id = J ]';] test_3d) assert_equal t_id3d test_3d

NB. 3D flatten to 1D vs J ,
('Futhark flat3d = J ,';, test_3d) assert_equal t_flat3d_1d test_3d

NB. 3D sum vs J +/,
('Futhark sum3d = J +/,';+/, test_3d) assert_equal t_sum3d test_3d

NB. -----------------------------------------------------------------
echo ''
echo '=== Dynamic Helper Function Tests ==='
echo ''

NB. Test fut_make_new_sig builds correct function names
sig_1d =. 'f64' fut_make_new_sig 1
('1D sig has _1d';1) assert_equal +./ '_1d' E. sig_1d
('1D sig has futhark_new';1) assert_equal +./ 'futhark_new' E. sig_1d

sig_2d =. 'f64' fut_make_new_sig 2
('2D sig has _2d';1) assert_equal +./ '_2d' E. sig_2d
('2D sig has futhark_new';1) assert_equal +./ 'futhark_new' E. sig_2d

sig_3d =. 'f64' fut_make_new_sig 3
('3D sig has _3d';1) assert_equal +./ '_3d' E. sig_3d
('3D sig has futhark_new';1) assert_equal +./ 'futhark_new' E. sig_3d

NB. Test different types
sig_i32 =. 'i32' fut_make_new_sig 2
('i32 sig has _i32';1) assert_equal +./ 'i32_2d' E. sig_i32

NB. -----------------------------------------------------------------
echo ''
echo '=== Dyadic (Two-Input) Function Tests ==='
echo ''

NB. Dot product: scalar result from two inputs
t_dot =: 'entry f (xs: []f64) (ys: []f64) : f64 = reduce (+) 0 (map2 (*) xs ys)' fut

NB. Test small vectors
test_a =: 1 2 3
test_b =: 4 5 6
('Dyadic dot product (1 2 3).(4 5 6)';32) assert_equal test_a t_dot test_b

NB. Test larger vectors
test_x =: i. 100
test_y =: 100 ?@$ 0
('Dyadic dot = J +/@:*';1) assert_equal 1e_10 > | (test_x t_dot test_y) - test_x +/@:* test_y

NB. Element-wise multiplication: array result from two inputs
t_mul =: 'entry f (xs: []f64) (ys: []f64) : []f64 = map2 (*) xs ys' fut
('Dyadic mul = J *';test_a * test_b) assert_equal test_a t_mul test_b

NB. Element-wise addition: array result from two inputs
t_add =: 'entry f (xs: []f64) (ys: []f64) : []f64 = map2 (+) xs ys' fut
('Dyadic add = J +';test_a + test_b) assert_equal test_a t_add test_b

NB. Element-wise subtraction
t_sub =: 'entry f (xs: []f64) (ys: []f64) : []f64 = map2 (-) xs ys' fut
('Dyadic sub = J -';test_a - test_b) assert_equal test_a t_sub test_b

NB. Element-wise max
t_max2 =: 'entry f (xs: []f64) (ys: []f64) : []f64 = map2 f64.max xs ys' fut
('Dyadic max = J >.';test_a >. test_b) assert_equal test_a t_max2 test_b

NB. Element-wise min
t_min2 =: 'entry f (xs: []f64) (ys: []f64) : []f64 = map2 f64.min xs ys' fut
('Dyadic min = J <.';test_a <. test_b) assert_equal test_a t_min2 test_b

NB. Input count detection
('Single input count';1) assert_equal fut_count_inputs 'entry f (xs: []f64) : []f64 = xs'
('Double input count';2) assert_equal fut_count_inputs 'entry f (xs: []f64) (ys: []f64) : f64 = 0'

NB. -----------------------------------------------------------------
echo ''
echo '=== Scalar Input Tests (v1.6.0) ==='
echo ''

NB. Scalar input, array output - prime factorization
t_factors =: 'entry f (n: i64) : []i64 = let (_, _, result) = loop (x, i, acc) = (n, 2i64, ([] : []i64)) while x > 1 do if x % i == 0 then (x / i, i, acc ++ [i]) else (x, i + 1, acc) in result' fut

('Factors of 12';2 2 3) assert_equal t_factors 12
('Factors of 100';2 2 5 5) assert_equal t_factors 100
('Factors of 17 (prime)';,17) assert_equal t_factors 17
('Factors of 1 (empty)';0) assert_equal # t_factors 1

NB. Scalar input, scalar output - Fibonacci
t_fib =: 'entry f (n: i64) : i64 = let (result, _, _) = loop (a, b, i) = (0i64, 1i64, 0i64) while i < n do (b, a + b, i + 1) in result' fut

('Fib 0';0) assert_equal t_fib 0
('Fib 1';1) assert_equal t_fib 1
('Fib 10';55) assert_equal t_fib 10
('Fib 20';6765) assert_equal t_fib 20

NB. Scalar input parsing tests
code_scalar_in =. 'entry f (n: i64) : []i64 = [n]'
('Scalar input rank';0) assert_equal fut_input_rank code_scalar_in
('Scalar input dtype';'i64') assert_equal fut_input_dtype code_scalar_in

code_scalar_in_out =. 'entry f (n: i64) : i64 = n * 2'
('Scalar in/out input rank';0) assert_equal fut_input_rank code_scalar_in_out
('Scalar in/out output rank';0) assert_equal fut_output_rank code_scalar_in_out

NB. Test ): vs ) : parsing (no space before colon)
code_no_space =. 'entry f (n: i64): []i64 = [n]'
('Parse return type no space';'[]i64') assert_equal fut_parse_return_type code_no_space
('Parse output rank no space';1) assert_equal fut_output_rank code_no_space

NB. -----------------------------------------------------------------
echo ''
echo '=== Fixed-Size Array Tests (v1.7.0) ==='
echo ''

NB. Test fixed-size dimension parsing
code_fixed =. 'entry f (ns: []i64) : [][3]i64 = map (\n -> [n, n*2, n*3]) ns'
('Fixed-size output rank';2) assert_equal fut_output_rank code_fixed
('Fixed-size output dtype';'i64') assert_equal fut_output_dtype code_fixed
('Fixed-size input rank';1) assert_equal fut_input_rank code_fixed

NB. Test basetype extraction from fixed-size arrays
('Extract basetype from [][3]i64';'i64') assert_equal fut_extract_basetype '[][3]i64'
('Extract basetype from [][32]i64';'i64') assert_equal fut_extract_basetype '[][32]i64'
('Extract basetype from [5][3]f32';'f32') assert_equal fut_extract_basetype '[5][3]f32'

NB. Test bracket counting with fixed dimensions
('Bracket count [][3]i64';2) assert_equal fut_count_brackets '[][3]i64'
('Bracket count [][32]i64';2) assert_equal fut_count_brackets '[][32]i64'
('Bracket count [5][3][7]f64';3) assert_equal fut_count_brackets '[5][3][7]f64'

NB. 1D to 2D with fixed inner dimension
t_triple =: 'entry f (ns: []i64) : [][3]i64 = map (\n -> [n, n*2, n*3]) ns' fut
triple_result =: t_triple 1 2 3 4 5
('1D->2D fixed shape';5 3) assert_equal $ triple_result
('1D->2D fixed row 1';1 2 3) assert_equal 0 { triple_result
('1D->2D fixed row 5';5 10 15) assert_equal 4 { triple_result

NB. -----------------------------------------------------------------
echo ''
echo '=== Helper Function Parsing Tests (v1.7.0) ==='
echo ''

NB. Test that helper functions don't interfere with entry point parsing
code_with_def =: 0 : 0
  def helper (x: i64) : i64 = x * 2

  entry f (ns: []i64) : []i64 =
    map helper ns
)
('Def helper: input type';'[]i64') assert_equal fut_parse_input_type code_with_def
('Def helper: input rank';1) assert_equal fut_input_rank code_with_def
('Def helper: output type';'[]i64') assert_equal fut_parse_return_type code_with_def
('Def helper: output rank';1) assert_equal fut_output_rank code_with_def

NB. Test more complex code with helper returning different type than entry
code_factorize =: 0 : 0
  def smallest_factor (n: i64) : i64 =
    if n <= 1 then 1
    else if n % 2 == 0 then 2
    else n

  entry f (n: i64): []i64 =
    let factors = [smallest_factor n]
    in factors
)
('Complex def: input rank';0) assert_equal fut_input_rank code_factorize
('Complex def: output type';'[]i64') assert_equal fut_parse_return_type code_factorize
('Complex def: output rank';1) assert_equal fut_output_rank code_factorize

NB. Test batch factorization (1D input, 2D output with helper)
code_batch =: 0 : 0
  def factorize (n: i64) : [32]i64 =
    let (_, _, factors, _) =
      loop (x, i, acc, c) = (n, 2i64, replicate 32 0i64, 0i64)
      while x > 1 && c < 32 do
        if x % i == 0
          then (x / i, i, acc with [c] = i, c + 1)
          else (x, i + 1, acc, c)
    in factors

  entry f (ns: []i64) : [][32]i64 =
    map factorize ns
)
('Batch factorize: input type';'[]i64') assert_equal fut_parse_input_type code_batch
('Batch factorize: input rank';1) assert_equal fut_input_rank code_batch
('Batch factorize: output type';'[][32]i64') assert_equal fut_parse_return_type code_batch
('Batch factorize: output rank';2) assert_equal fut_output_rank code_batch

NB. Compile and test batch factorization
t_batch_factor =: code_batch fut
batch_result =: t_batch_factor 12 100 30
('Batch factor shape';3 32) assert_equal $ batch_result
('Batch factor row 1 (12=2*2*3)';2 2 3) assert_equal 3 {. 0 { batch_result
('Batch factor row 2 (100=2*2*5*5)';2 2 5 5) assert_equal 4 {. 1 { batch_result
('Batch factor row 3 (30=2*3*5)';2 3 5) assert_equal 3 {. 2 { batch_result

NB. -----------------------------------------------------------------
NB. Summary
echo ''
echo '==============================================================='
echo ''
echo '  Tests run:    ', ": TESTS_RUN
echo '  Tests passed: ', ": TESTS_PASSED
echo '  Tests failed: ', ": TESTS_FAILED
echo ''

test_result =: TESTS_FAILED = 0

echo '  ', > test_result { 'SOME TESTS FAILED';'ALL TESTS PASSED!'

echo ''
echo '==============================================================='