open Expr
open Simplify
open Diff
open Substitute

let rec integrate var = function
  | Const c -> Some (Mul (Const c, Var var))
  | SymConst _ as s -> Some (Mul (s, Var var))
  | Var v when v = var -> Some (Div (Pow (Var var, Const 2.0), Const 2.0))
  | Var v -> Some (Mul (Var v, Var var))
  | Add (e1, e2) ->
      (match (integrate var e1, integrate var e2) with
       | Some i1, Some i2 -> Some (simplify (Add (i1, i2)))
       | _ -> None)
  | Sub (e1, e2) ->
      (match (integrate var e1, integrate var e2) with
       | Some i1, Some i2 -> Some (simplify (Sub (i1, i2)))
       | _ -> None)
  | Mul (Const c, e) | Mul (e, Const c) ->
      (match integrate var e with
       | Some i -> Some (simplify (Mul (Const c, i)))
       | None -> None)
  | Mul (SymConst s, e) | Mul (e, SymConst s) ->
      (match integrate var e with
       | Some i -> Some (simplify (Mul (SymConst s, i)))
       | None -> None)
  | Pow (Var v, Const n) when v = var && n <> -1.0 ->
      let exp = n +. 1.0 in
      Some (simplify (Div (Pow (Var var, Const exp), Const exp)))
  | Div (Const 1.0, Var v) when v = var ->
      Some (Ln (Abs (Var var)))
  | Div (e, Var v) when v = var ->
      (match e with
       | Const c -> Some (simplify (Mul (Const c, Ln (Abs (Var var)))))
       | _ -> None)
  | Sin (Var v) when v = var ->
      Some (Neg (Cos (Var var)))
  | Cos (Var v) when v = var ->
      Some (Sin (Var var))
  | Tan (Var v) when v = var ->
      Some (Neg (Ln (Abs (Cos (Var var)))))
  | Div (Const 1.0, Pow (Cos (Var v), Const 2.0)) when v = var ->
      Some (Tan (Var var))
  | Sinh (Var v) when v = var ->
      Some (Cosh (Var var))
  | Cosh (Var v) when v = var ->
      Some (Sinh (Var var))
  | Tanh (Var v) when v = var ->
      Some (Ln (Cosh (Var var)))
  | Div (Const 1.0, Sqrt (Sub (Const 1.0, Pow (Var v, Const 2.0)))) when v = var ->
      Some (Asin (Var var))
  | Neg (Div (Const 1.0, Sqrt (Sub (Const 1.0, Pow (Var v, Const 2.0))))) when v = var ->
      Some (Acos (Var var))
  | Div (Const 1.0, Add (Const 1.0, Pow (Var v, Const 2.0))) when v = var ->
      Some (Atan (Var var))
  | Exp (Var v) when v = var ->
      Some (Exp (Var var))
  | Pow (Const a, Var v) when v = var ->
      Some (simplify (Div (Pow (Const a, Var var), Ln (Const a))))
  | Pow (SymConst E, Var v) when v = var ->
      Some (Exp (Var var))
  | e ->
      match try_u_substitution var e with
      | Some result -> Some result
      | None -> try_by_parts var e

and try_u_substitution var expr =
  let rec find_inner = function
    | Sin u | Cos u | Tan u | Sinh u | Cosh u | Tanh u
    | Asin u | Acos u | Atan u | Exp u | Ln u | Sqrt u | Abs u -> Some u
    | Pow (u, _) -> Some u
    | Add (e1, e2) | Sub (e1, e2) | Mul (e1, e2) | Div (e1, e2) ->
        (match find_inner e1 with
         | Some _ as r -> r
         | None -> find_inner e2)
    | _ -> None
  in
  match find_inner expr with
  | Some u when u <> Var var ->
      let u_prime = diff var u in
      let expr_simplified = simplify expr in
      let test_expr = simplify (Div (expr_simplified, u_prime)) in
      let substituted = substitute var (Var "u_temp") test_expr in
      (match substituted with
       | e when not (contains_var var e) ->
           (match integrate "u_temp" e with
            | Some integrated ->
                let result = substitute "u_temp" u integrated in
                Some (simplify result)
            | None -> None)
       | _ -> None)
  | _ -> None

and try_by_parts var = function
  | Mul (e1, e2) ->
      let priority = function
        | Ln _ -> 5
        | Asin _ | Acos _ | Atan _ -> 4
        | Var _ | Pow (Var _, _) -> 3
        | Sin _ | Cos _ | Tan _ | Sinh _ | Cosh _ | Tanh _ -> 2
        | Exp _ -> 1
        | _ -> 0
      in
      let (u, dv) =
        if priority e1 >= priority e2 then (e1, e2) else (e2, e1)
      in
      (match integrate var dv with
       | Some v ->
           let du = diff var u in
           (match integrate var (simplify (Mul (v, du))) with
            | Some second_integral ->
                Some (simplify (Sub (Mul (u, v), second_integral)))
            | None -> None)
       | None -> None)
  | _ -> None

and contains_var var = function
  | Const _ | SymConst _ -> false
  | Var v -> v = var
  | Add (e1, e2) | Sub (e1, e2) | Mul (e1, e2) | Div (e1, e2) | Pow (e1, e2) ->
      contains_var var e1 || contains_var var e2
  | Neg e | Sin e | Cos e | Tan e | Sinh e | Cosh e | Tanh e
  | Asin e | Acos e | Atan e | Exp e | Ln e | Sqrt e | Abs e ->
      contains_var var e
  | Atan2 (e1, e2) | Log (e1, e2) ->
      contains_var var e1 || contains_var var e2

let integrate_definite var lower upper expr =
  match integrate var expr with
  | None -> None
  | Some antideriv ->
      let upper_val = Eval.eval [(var, upper)] antideriv in
      let lower_val = Eval.eval [(var, lower)] antideriv in
      Some (upper_val -. lower_val)