open Expr
open Simplify

let rec diff var = function
  | Const _ -> Const 0.0
  | SymConst _ -> Const 0.0
  | Var v -> if v = var then Const 1.0 else Const 0.0
  | Add (e1, e2) -> simplify (Add (diff var e1, diff var e2))
  | Sub (e1, e2) -> simplify (Sub (diff var e1, diff var e2))
  | Mul (e1, e2) ->
      simplify (Add (Mul (diff var e1, e2), Mul (e1, diff var e2)))
  | Div (e1, e2) ->
      let num = Sub (Mul (diff var e1, e2), Mul (e1, diff var e2)) in
      let den = Pow (e2, Const 2.0) in
      simplify (Div (num, den))
  | Pow (e, Const n) ->
      simplify (Mul (Mul (Const n, Pow (e, Const (n -. 1.0))), diff var e))
  | Pow (e1, e2) ->
      let term1 = Mul (e2, Mul (Pow (e1, Sub (e2, Const 1.0)), diff var e1)) in
      let term2 = Mul (Pow (e1, e2), Mul (Ln e1, diff var e2)) in
      simplify (Add (term1, term2))
  | Neg e -> simplify (Neg (diff var e))
  | Sin e -> simplify (Mul (Cos e, diff var e))
  | Cos e -> simplify (Neg (Mul (Sin e, diff var e)))
  | Tan e ->
      let sec2 = Div (Const 1.0, Pow (Cos e, Const 2.0)) in
      simplify (Mul (sec2, diff var e))
  | Sinh e -> simplify (Mul (Cosh e, diff var e))
  | Cosh e -> simplify (Mul (Sinh e, diff var e))
  | Tanh e ->
      let sech2 = Sub (Const 1.0, Pow (Tanh e, Const 2.0)) in
      simplify (Mul (sech2, diff var e))
  | Asin e ->
      let denom = Sqrt (Sub (Const 1.0, Pow (e, Const 2.0))) in
      simplify (Div (diff var e, denom))
  | Acos e ->
      let denom = Sqrt (Sub (Const 1.0, Pow (e, Const 2.0))) in
      simplify (Neg (Div (diff var e, denom)))
  | Atan e ->
      let denom = Add (Const 1.0, Pow (e, Const 2.0)) in
      simplify (Div (diff var e, denom))
  | Atan2 (y, x) ->
      let num = Sub (Mul (x, diff var y), Mul (y, diff var x)) in
      let denom = Add (Pow (x, Const 2.0), Pow (y, Const 2.0)) in
      simplify (Div (num, denom))
  | Exp e -> simplify (Mul (Exp e, diff var e))
  | Ln e -> simplify (Div (diff var e, e))
  | Log (base_e, arg) ->
      let denom = Mul (arg, Ln base_e) in
      simplify (Div (diff var arg, denom))
  | Sqrt e ->
      let denom = Mul (Const 2.0, Sqrt e) in
      simplify (Div (diff var e, denom))
  | Abs e ->
      let sgn = Div (e, Abs e) in
      simplify (Mul (sgn, diff var e))

let rec diff_n var n expr =
  if n <= 0 then expr
  else diff_n var (n - 1) (diff var expr)

let partial vars expr =
  List.fold_left (fun e v -> diff v e) expr vars