open OUnit2
open Leibniz

let test_lexer_basic _ =
  let tokens = Lexer.tokenize "2 + 3" in
  assert_equal 4 (List.length tokens)

let test_lexer_implicit_mult _ =
  let tokens = Lexer.tokenize "2x" in
  let token_list = List.map (fun t -> t.Lexer.token) tokens in
  assert_equal 4 (List.length token_list);
  match token_list with
  | [Lexer.Num 2.0; Lexer.Star; Lexer.Var "x"; Lexer.EOF] -> ()
  | _ -> assert_failure "implicit multiplication not working"

let test_parser_basic _ =
  let expr = Parser.parse "x + 1" in
  assert_equal (Expr.Add (Expr.Var "x", Expr.Const 1.0)) expr

let test_parser_implicit_mult _ =
  let expr = Parser.parse "2x" in
  assert_equal (Expr.Mul (Expr.Const 2.0, Expr.Var "x")) expr

let test_parser_symbolic_constants _ =
  let expr = Parser.parse "pi" in
  assert_equal (Expr.SymConst Expr.Pi) expr

let test_simplify_basic _ =
  let expr = Expr.Add (Expr.Const 0.0, Expr.Var "x") in
  let result = Simplify.simplify expr in
  assert_equal (Expr.Var "x") result

let test_simplify_fixed_point _ =
  let expr = Expr.Add (Expr.Mul (Expr.Const 0.0, Expr.Var "x"), Expr.Var "y") in
  let result = Simplify.simplify expr in
  assert_equal (Expr.Var "y") result

let test_simplify_collect_like_terms _ =
  let expr = Expr.Add (Expr.Var "x", Expr.Var "x") in
  let result = Simplify.simplify expr in
  assert_equal (Expr.Mul (Expr.Const 2.0, Expr.Var "x")) result

let test_diff_basic _ =
  let expr = Expr.Pow (Expr.Var "x", Expr.Const 2.0) in
  let result = Diff.diff "x" expr in
  let expected = Simplify.simplify (Expr.Mul (Expr.Const 2.0, Expr.Var "x")) in
  assert_equal expected result

let test_diff_sin _ =
  let expr = Expr.Sin (Expr.Var "x") in
  let result = Diff.diff "x" expr in
  assert_equal (Expr.Cos (Expr.Var "x")) result

let test_diff_product_rule _ =
  let expr = Expr.Mul (Expr.Var "x", Expr.Sin (Expr.Var "x")) in
  let _result = Diff.diff "x" expr in
  ()

let test_eval_basic _ =
  let expr = Expr.Add (Expr.Var "x", Expr.Const 1.0) in
  let result = Eval.eval [("x", 2.0)] expr in
  assert_equal 3.0 result

let test_eval_symbolic_constants _ =
  let expr = Expr.SymConst Expr.Pi in
  let result = Eval.eval [] expr in
  assert_bool "pi evaluation" (abs_float (result -. 3.14159265) < 0.0001)

let test_canonical_equality _ =
  let e1 = Expr.Add (Expr.Var "x", Expr.Var "y") in
  let e2 = Expr.Add (Expr.Var "y", Expr.Var "x") in
  assert_bool "commutativity" (Canonical.equal e1 e2)

let test_substitute_basic _ =
  let expr = Expr.Add (Expr.Var "x", Expr.Const 1.0) in
  let result = Substitute.substitute "x" (Expr.Const 2.0) expr in
  assert_equal (Expr.Add (Expr.Const 2.0, Expr.Const 1.0)) result

let test_integrate_basic _ =
  let expr = Expr.Var "x" in
  match Integrate.integrate "x" expr with
  | Some result ->
      let expected = Expr.Div (Expr.Pow (Expr.Var "x", Expr.Const 2.0), Expr.Const 2.0) in
      assert_bool "integration result not equal" (Canonical.equal expected result)
  | None -> assert_failure "integration failed"

let test_integrate_sin _ =
  let expr = Expr.Sin (Expr.Var "x") in
  match Integrate.integrate "x" expr with
  | Some result ->
      assert_equal (Expr.Neg (Expr.Cos (Expr.Var "x"))) result
  | None -> assert_failure "integration of sin failed"

let test_taylor_sin _ =
  let expr = Expr.Sin (Expr.Var "x") in
  let _result = Series.maclaurin "x" expr 5 in
  ()

let test_gradient_basic _ =
  let expr = Expr.Add (Expr.Pow (Expr.Var "x", Expr.Const 2.0),
                       Expr.Pow (Expr.Var "y", Expr.Const 2.0)) in
  let grad = Multivariate.gradient ["x"; "y"] expr in
  assert_equal 2 (List.length grad)

let test_numerical_bisection _ =
  let expr = Expr.Sub (Expr.Pow (Expr.Var "x", Expr.Const 2.0), Expr.Const 4.0) in
  match Numerical.bisection expr "x" 0.0 3.0 0.001 100 with
  | Some root -> assert_bool "root near 2" (abs_float (root -. 2.0) < 0.01)
  | None -> assert_failure "bisection failed"

let suite =
  "leibniz tests" >::: [
    "lexer basic" >:: test_lexer_basic;
    "lexer implicit mult" >:: test_lexer_implicit_mult;
    "parser basic" >:: test_parser_basic;
    "parser implicit mult" >:: test_parser_implicit_mult;
    "parser symbolic constants" >:: test_parser_symbolic_constants;
    "simplify basic" >:: test_simplify_basic;
    "simplify fixed point" >:: test_simplify_fixed_point;
    "simplify collect like terms" >:: test_simplify_collect_like_terms;
    "diff basic" >:: test_diff_basic;
    "diff sin" >:: test_diff_sin;
    "diff product rule" >:: test_diff_product_rule;
    "eval basic" >:: test_eval_basic;
    "eval symbolic constants" >:: test_eval_symbolic_constants;
    "canonical equality" >:: test_canonical_equality;
    "substitute basic" >:: test_substitute_basic;
    "integrate basic" >:: test_integrate_basic;
    "integrate sin" >:: test_integrate_sin;
    "taylor sin" >:: test_taylor_sin;
    "gradient basic" >:: test_gradient_basic;
    "numerical bisection" >:: test_numerical_bisection;
  ]

let () = run_test_tt_main suite