open Driver
open Test_common

let output_dir =
  let name = "test_dual_floats" in
  if not (Sys.file_exists name) then Sys.mkdir name 0o755;
  name

let backends_to_test =
  [
    Config.GradualGuarded;
    Config.GradualMonoClosures;
    Config.GradualMonoValues;
    Config.GradualGuardedDFO;
    Config.GradualMonoDFO;
  ]

let test = test ~output_dir

(* simple examples of using floats and lambdas with ? or float annotations *)
let simple_float_arithmetic_test =
  test __POS__ "simple float arithmetic"
    {|
  let x : ? = 3.5 in
  let y : float = 2.0 in
  let z : ? = (x :: float) +. y in
  print_float (z :: float)
|}
    "5.5"

let simple_float_lambda_test =
  test __POS__ "simple float lambda"
    {|
  let simple_calc : ? =
  (fun (x : float) (y : ?) ->
     (let result : float = x +. y in
     result))
  in

  let result1 : ? = simple_calc 5.0 3.0 in
  print_float result1;
  0
|}
    "8"

let simple_float_tuple_test =
  test __POS__ "simple float tuple"
    {|
  let tup : float * ? = (1.0, 2.0 :: ?) in
  let x : float = tup.(0) in
  let y : ? = tup.(1) in
  let sum : float = x +. y in
  print_float sum
|}
    "3"

let simple_float_vector_test =
  test __POS__ "simple float vector"
    {|
  let v : vec[?] = vector 3 (1.0 :: ?) in
  vector_set v 1 (2.0 :: ?);
  vector_set v 2 (3.0 :: ?);
  (let sum : float = (vector_get v 0) +. (vector_get v 1) +. (vector_get v 2) in
  print_float sum)
|}
    "6"

(* ===== Tuple Coercion Tests ===== *)

let tuple_dynamic_float_test =
  test __POS__ "tuple dynamic float"
    {|
  let x : ? = 3.14 in
  let t : ? * int = (x, 42) in
  print_float (t.(0) :: float);
  print_char '\n';
  print_int (t.(1))
|}
    "3.14\n42"

let tuple_mixed_types_test =
  test __POS__ "tuple mixed types"
    {|
  let a : float = 1.5 in
  let b : ? = 2.5 in
  let t : float * ? = (a, b) in
  print_float (t.(0));
  print_char '\n';
  print_float (t.(1) :: float)
|}
    "1.5\n2.5"

let tuple_dynamic_to_static_test =
  test __POS__ "tuple dynamic to static"
    {|
  let t : ? * ? * int = (1.0, 2.0, 3) in
  let sum : float = (t.(0) :: float) +. (t.(1) :: float) in
  print_float sum
|}
    "3"

let tuple_static_to_dynamic_test =
  test __POS__ "tuple static to dynamic"
    {|
  let t : float * float = (1.5, 2.5) in
  let x : ? = t.(0) in
  let y : ? = t.(1) in
  print_float (x :: float);
  print_char '\n';
  print_float (y :: float)
|}
    "1.5\n2.5"

let tuple_nested_coercion_test =
  test __POS__ "tuple nested coercion"
    {|
  let inner : float * ? = (1.0, 2.0) in
  let outer : ? * int = (inner, 42) in
  let retrieved : float * ? = outer.(0) in
  print_float (retrieved.(0));
  print_char '\n';
  print_float (retrieved.(1) :: float)
|}
    "1\n2"

(* ===== Vector Coercion Tests ===== *)

let vector_dynamic_float_test =
  test __POS__ "vector dynamic float"
    {|
  let v : vec[?] = vector 3 (1.5 :: ?) in
  let _ = vector_set v 1 (2.5 :: ?) in
  let _ = vector_set v 2 (3.5 :: ?) in
  print_float (vector_get v 0 :: float);
  print_char '\n';
  print_float (vector_get v 1 :: float);
  print_char '\n';
  print_float (vector_get v 2 :: float)
|}
    "1.5\n2.5\n3.5"

let vector_static_to_dynamic_test =
  test __POS__ "vector static to dynamic"
    {|
  let v : vec[float] = vector 2 1.0 in
  vector_set v 1 2.5;
  (let x : ? = vector_get v 0 in
  let y : ? = vector_get v 1 in
  print_float (x :: float);
  print_char '\n';
  print_float (y :: float))
|}
    "1\n2.5"

let vector_mixed_operations_test =
  test __POS__ "vector mixed operations"
    {|
  let v : vec[?] = vector 3 (0.0 :: ?) in
  let _ = vector_set v 0 (1.5 :: ?) in
  let _ = vector_set v 1 (2.5 :: ?) in
  let sum : float = (vector_get v 0 :: float) +. (vector_get v 1 :: float) in
  print_float sum
|}
    "4"

let vector_length_dynamic_test =
  test __POS__ "vector length dynamic"
    {|
  let v : vec[?] = vector 5 (3.14 :: ?) in
  print_int (vector_length v)
|}
    "5"

(* ===== Lambda Coercion Tests ===== *)

let lambda_dynamic_arg_test =
  test __POS__ "lambda dynamic arg"
    {|
  let f : ? -> int = fun x : ? -> (x :: int) + 1 in
  print_int (f (5 :: ?))
|}
    "6"

let lambda_dynamic_return_test =
  test __POS__ "lambda dynamic return"
    {|
  let f : int -> ? = fun x : int -> (x + 2 :: ?) in
  let result : ? = f 3 in
  print_int (result :: int)
|}
    "5"

let lambda_fully_dynamic_test =
  test __POS__ "lambda fully dynamic"
    {|
  let f : ? -> ? = fun x : ? -> 
    let y : int = x :: int in
    (y * 2 :: ?)
  in
  print_int (f (4 :: ?) :: int)
|}
    "8"

let lambda_higher_order_test =
  test __POS__ "lambda higher order"
    {|
  let apply : (? -> ?) -> ? -> ? = 
    fun (f : (? -> ?)) (x : ?) -> f x
  in
  let double : ? -> ? = fun x : ? -> 
    let y : int = x :: int in
    (y * 2 :: ?)
  in
  print_int (apply double (3 :: ?) :: int)
|}
    "6"

let lambda_closure_test =
  test __POS__ "lambda closure"
    {|
  let make_adder : int -> (? -> int) = 
    fun n : int -> fun x : ? -> n + (x :: int)
  in
  let add5 : ? -> int = make_adder 5 in
  print_int (add5 (3 :: ?))
|}
    "8"

(* ===== Lambda Coercion Tests (floats version )===== *)

let lambda_dynamic_arg_float_test =
  test __POS__ "lambda dynamic arg (floats)"
    {|
  let f : ? -> float = fun x : ? -> (x :: float) +. 1.0 in
  print_float (f (5.0 :: ?))
|}
    "6"

let lambda_dynamic_return_float_test =
  test __POS__ "lambda dynamic return (floats)"
    {|
  let f : float -> ? = fun x : float -> (x +. 2.0 :: ?) in
  let result : ? = f 3.0 in
  print_float (result :: float)
|}
    "5"

let lambda_fully_dynamic_float_test =
  test __POS__ "lambda fully dynamic (floats)"
    {|
  let f : ? -> ? = fun x : ? -> 
    let y : float = x :: float in
    (y *. 2.0 :: ?)
  in
  print_float (f (4.5 :: ?) :: float)
|}
    "9"

let lambda_higher_order_float_test =
  test __POS__ "lambda higher order (floats)"
    {|
  let apply : (? -> ?) -> ? -> ? = 
    fun (f : (? -> ?)) (x : ?) -> f x
  in
  let double : ? -> ? = fun x : ? -> 
    let y : float = x :: float in
    (y *. 2.0 :: ?)
  in
  print_float (apply double (3.0 :: ?) :: float)
|}
    "6"

let lambda_closure_float_test =
  test __POS__ "lambda closure (floats)"
    {|
  let make_adder : float -> (? -> float) = 
    fun n : float -> fun x : ? -> n +. (x :: float)
  in
  let add5 : ? -> float = make_adder 5.0 in
  print_float (add5 (3.0 :: ?))
|}
    "8"

(* ===== Vector of Tuples Tests ===== *)

let vector_of_tuples_basic_test =
  test __POS__ "vector of tuples basic"
    {|
  let v : vec[int * ?] = vector 2 (1, (2 :: ?)) in
  let _ = vector_set v 1 (3, (4 :: ?)) in
  let t1 : int * ? = vector_get v 0 in
  let t2 : int * ? = vector_get v 1 in
  print_int (t1.(0));
  print_char '\n';
  print_int (t1.(1) :: int);
  print_char '\n';
  print_int (t2.(0));
  print_char '\n';
  print_int (t2.(1) :: int)
|}
    "1\n2\n3\n4"

let vector_of_tuples_dynamic_test =
  test __POS__ "vector of tuples dynamic"
    {|
  let v : vec[? * ?] = vector 2 ((1 :: ?), (2 :: ?)) in
  let _ = vector_set v 1 ((3 :: ?), (4 :: ?)) in
  let t : ? * ? = vector_get v 1 in
  print_int (t.(0) :: int);
  print_char '\n';
  print_int (t.(1) :: int)
|}
    "3\n4"

let vector_of_tuples_computation_test =
  test __POS__ "vector of tuples computation"
    {|
  let v : vec[int * int] = vector 3 (0, 0) in
  let _ = vector_set v 0 (1, 2) in
  let _ = vector_set v 1 (3, 4) in
  let _ = vector_set v 2 (5, 6) in
  let sum : int = 
    let t0 = vector_get v 0 in
    let t1 = vector_get v 1 in
    let t2 = vector_get v 2 in
    t0.(0) + t0.(1) + t1.(0) + t1.(1) + t2.(0) + t2.(1)
  in
  print_int sum
|}
    "21"

let vector_of_tuples_mixed_test =
  test __POS__ "vector of tuples mixed"
    {|
  let v : vec[? * int * ?] = vector 2 ((1 :: ?), 2, (3 :: ?)) in
  let t : ? * int * ? = vector_get v 0 in
  print_int (t.(0) :: int);
  print_char '\n';
  print_int (t.(1));
  print_char '\n';
  print_int (t.(2) :: int)
|}
    "1\n2\n3"

(* ===== Vector of Tuples Tests (float version) ===== *)
let vector_of_tuples_basic_float_test =
  test __POS__ "vector of tuples basic"
    {|
  let v : vec[float * ?] = vector 2 (1.0, (2.0 :: ?)) in
  let _ = vector_set v 1 (3.0, (4.0 :: ?)) in
  let t1 : float * ? = vector_get v 0 in
  let t2 : float * ? = vector_get v 1 in
  print_float (t1.(0));
  print_char '\n';
  print_float (t1.(1) :: float);
  print_char '\n';
  print_float (t2.(0));
  print_char '\n';
  print_float (t2.(1) :: float)
|}
    "1\n2\n3\n4"

let vector_of_tuples_dynamic_float_test =
  test __POS__ "vector of tuples dynamic"
    {|
  let v : vec[? * ?] = vector 2 ((1.0 :: ?), (2.0 :: ?)) in
  let _ = vector_set v 1 ((3.0 :: ?), (4.0 :: ?)) in
  let t : ? * ? = vector_get v 1 in
  print_float (t.(0) :: float);
  print_char '\n';
  print_float (t.(1) :: float)
|}
    "3\n4"

let vector_of_tuples_computation_float_test =
  test __POS__ "vector of tuples computation"
    {|
  let v : vec[float * float] = vector 3 (0.0, 0.0) in
  let _ = vector_set v 0 (1.0, 2.0) in
  let _ = vector_set v 1 (3.0, 4.0) in
  let _ = vector_set v 2 (5.0, 6.0) in
  let sum : float = 
    let t0 = vector_get v 0 in
    let t1 = vector_get v 1 in
    let t2 = vector_get v 2 in
    t0.(0) +. t0.(1) +. t1.(0) +. t1.(1) +. t2.(0) +. t2.(1)
  in
  print_float sum
|}
    "21"

let vector_of_tuples_mixed_float_test =
  test __POS__ "vector of tuples mixed"
    {|
  let v : vec[? * float * ?] = vector 2 ((1.0 :: ?), 2.0, (3.0 :: ?)) in
  let t : ? * float * ? = vector_get v 0 in
  print_float (t.(0) :: float);
  print_char '\n';
  print_float (t.(1));
  print_char '\n';
  print_float (t.(2) :: float)
|}
    "1\n2\n3"

(* ===== Tuple of Tuples Tests ===== *)

let tuple_of_tuples_basic_test =
  test __POS__ "tuple of tuples basic"
    {|
  let inner1 : float * float = (1.0, 2.0) in
  let inner2 : float * float = (3.0, 4.0) in
  let outer : (float * float) * (float * float) = (inner1, inner2) in
  let t1 : float * float = outer.(0) in
  let t2 : float * float = outer.(1) in
  print_float (t1.(0));
  print_char '\n';
  print_float (t1.(1));
  print_char '\n';
  print_float (t2.(0));
  print_char '\n';
  print_float (t2.(1))
|}
    "1\n2\n3\n4"

let tuple_of_tuples_dynamic_test =
  test __POS__ "tuple of tuples dynamic"
    {|
  let inner1 : ? * ? = ((1.0 :: ?), (2.0 :: ?)) in
  let inner2 : ? * ? = ((3.0 :: ?), (4.0 :: ?)) in
  let outer : (? * ?) * (? * ?) = (inner1, inner2) in
  let t1 : ? * ? = outer.(0) in
  print_float (t1.(0) :: float);
  print_char '\n';
  print_float (t1.(1) :: float)
|}
    "1\n2"

let tuple_of_tuples_mixed_test =
  test __POS__ "tuple of tuples mixed"
    {|
  let inner1 : float * ? = (1.0, (2.0 :: ?)) in
  let inner2 : ? * float = ((3.0 :: ?), 4.0) in
  let outer : (float * ?) * (? * float) = (inner1, inner2) in
  let t1 : float * ? = outer.(0) in
  let t2 : ? * float = outer.(1) in
  print_float (t1.(0));
  print_char '\n';
  print_float (t1.(1) :: float);
  print_char '\n';
  print_float (t2.(0) :: float);
  print_char '\n';
  print_float (t2.(1))
|}
    "1\n2\n3\n4"

let tuple_of_tuples_nested_test =
  test __POS__ "tuple of tuples nested"
    {|
  let t : ((float * ?) * int) * (? * (float * float)) = 
    (((1.0, (2.0 :: ?)), 3), ((4.0 :: ?), (5.0, 6.0))) in
  let left : (float * ?) * int = t.(0) in
  let left_inner : float * ? = left.(0) in
  print_float (left_inner.(0));
  print_char '\n';
  print_float (left_inner.(1) :: float);
  print_char '\n';
  print_int (left.(1))
|}
    "1\n2\n3"

(* ===== Complex Combination Tests ===== *)

let lambda_with_vector_test =
  test __POS__ "lambda with vector"
    {|
  let map_vec : ((? -> ?) -> (vec[?] -> vec[?])) = 
  fun (f : ? -> ?) -> (fun (v : vec[?]) ->
      let result : vec[?] = vector (vector_length v) (0.0 :: ?) in
      let _ = vector_set result 0 (f (vector_get v 0)) in
      let _ = vector_set result 1 (f (vector_get v 1)) in
      result)
  in
  let double : (? -> ?) = fun (x : ?) -> 
  let y : float = x :: float in
  (y *. 2.0 :: ?)
  in
  let v : vec[?] = vector 2 (1.5 :: ?) in
  let _ = vector_set v 1 (2.5 :: ?) in
  let result : vec[?] = (map_vec double) v in
  print_float (vector_get result 0 :: float);
  print_char '\n';
  print_float (vector_get result 1 :: float)
|}
    "3\n5"

let tuple_in_lambda_test =
  test __POS__ "tuple in lambda"
    {|
  let swap : ? * ? -> ? * ? = 
    fun (t : ? * ?) -> (t.(1), t.(0))
  in
  let input : ? * ? = ((1.0 :: ?), (2.0 :: ?)) in
  let output : ? * ? = swap input in
  print_float (output.(0) :: float);
  print_char '\n';
  print_float (output.(1) :: float)
|}
    "2\n1"

let vector_of_lambdas_test =
  test __POS__ "vector of lambdas"
    {|
  let v : vec[? -> ?] = vector 2 (fun x : ? -> x) in
  let f1 : ? -> ? = fun x : ? -> 
    let y : float = x :: float in
    (y +. 1.0 :: ?)
  in
  let f2 : ? -> ? = fun x : ? -> 
    let y : float = x :: float in
    (y *. 2.0 :: ?)
  in
  let _ = vector_set v 0 f1 in
  let _ = vector_set v 1 f2 in
  let g1 : ? -> ? = vector_get v 0 in
  let g2 : ? -> ? = vector_get v 1 in
  print_float (g1 (5.0 :: ?) :: float);
  print_char '\n';
  print_float (g2 (5.0 :: ?) :: float)
|}
    "6\n10"

let nested_structure_test =
  test __POS__ "nested structure"
    {|
  let v : vec[? * (int * ?)] = vector 2 ((0 :: ?), (0, (0 :: ?))) in
  let _ = vector_set v 0 ((1 :: ?), (2, (3 :: ?))) in
  let _ = vector_set v 1 ((4 :: ?), (5, (6 :: ?))) in
  let t : ? * (int * ?) = vector_get v 1 in
  let inner : int * ? = t.(1) in
  print_int (t.(0) :: int);
  print_char '\n';
  print_int (inner.(0));
  print_char '\n';
  print_int (inner.(1) :: int)
|}
    "4\n5\n6"

let nested_structure_float_test =
  test __POS__ "nested structure (float)"
    {|
  let v : vec[? * (float * ?)] = vector 2 ((0.0 :: ?), (0.0, (0.0 :: ?))) in
  let _ = vector_set v 0 ((1.0 :: ?), (2.0, (3.0 :: ?))) in
  let _ = vector_set v 1 ((4.0 :: ?), (5.0, (6.0 :: ?))) in
  let t : ? * (float * ?) = vector_get v 1 in
  let inner : float * ? = t.(1) in
  print_float (t.(0) :: float);
  print_char '\n';
  print_float (inner.(0));
  print_char '\n';
  print_float (inner.(1) :: float)
|}
    "4\n5\n6"

let suite =
  [
    ( "simple_float_tests",
      [
        simple_float_arithmetic_test;
        simple_float_lambda_test;
        simple_float_tuple_test;
        simple_float_vector_test;
      ] );
    ( "tuple_coercion",
      [
        tuple_dynamic_float_test;
        tuple_mixed_types_test;
        tuple_dynamic_to_static_test;
        tuple_static_to_dynamic_test;
        tuple_nested_coercion_test;
      ] );
    ( "vector_coercion",
      [
        vector_dynamic_float_test;
        vector_static_to_dynamic_test;
        vector_mixed_operations_test;
        vector_length_dynamic_test;
      ] );
    ( "lambda_coercion",
      [
        lambda_dynamic_arg_test;
        lambda_dynamic_return_test;
        lambda_fully_dynamic_test;
        lambda_higher_order_test;
        lambda_closure_test;
      ] );
    ( "lambda_coercion (floats version)",
      [
        lambda_dynamic_arg_float_test;
        lambda_dynamic_return_float_test;
        lambda_fully_dynamic_float_test;
        lambda_higher_order_float_test;
        lambda_closure_float_test;
      ] );
    ( "vector_of_tuples",
      [
        vector_of_tuples_basic_test;
        vector_of_tuples_basic_float_test;
        vector_of_tuples_dynamic_test;
        vector_of_tuples_computation_test;
        vector_of_tuples_mixed_test;
      ] );
    ( "vector_of_tuples (floats version)",
      [
        vector_of_tuples_basic_float_test;
        vector_of_tuples_dynamic_float_test;
        vector_of_tuples_computation_float_test;
        vector_of_tuples_mixed_float_test;
      ] );
    ( "tuple_of_tuples",
      [
        tuple_of_tuples_basic_test;
        tuple_of_tuples_dynamic_test;
        tuple_of_tuples_mixed_test;
        tuple_of_tuples_nested_test;
      ] );
    ( "complex_combinations",
      [
        lambda_with_vector_test;
        tuple_in_lambda_test;
        vector_of_lambdas_test;
        nested_structure_test;
        nested_structure_float_test;
      ] );
  ]

let () =
  Alcotest.run "test_coercion" @@ to_alcotest_test_suite backends_to_test suite
