open Expr
open Diff
open Simplify

let gradient vars expr =
  List.map (fun v -> diff v expr) vars

let jacobian vars exprs =
  List.map (fun expr -> gradient vars expr) exprs

let hessian vars expr =
  let grad = gradient vars expr in
  List.map (fun deriv -> List.map (fun v -> diff v deriv) vars) grad

let divergence vars vector_field =
  if List.length vars <> List.length vector_field then
    failwith "divergence: dimension mismatch"
  else
    let terms = List.map2 (fun v e -> diff v e) vars vector_field in
    let sum = List.fold_left (fun acc t -> Add (acc, t)) (Const 0.0) terms in
    simplify sum

let curl vector_field =
  match vector_field with
  | [fx; fy; fz] ->
      let dfz_dy = diff "y" fz in
      let dfy_dz = diff "z" fy in
      let dfx_dz = diff "z" fx in
      let dfz_dx = diff "x" fz in
      let dfy_dx = diff "x" fy in
      let dfx_dy = diff "y" fx in
      [simplify (Sub (dfz_dy, dfy_dz));
       simplify (Sub (dfx_dz, dfz_dx));
       simplify (Sub (dfy_dx, dfx_dy))]
  | _ -> failwith "curl: requires exactly 3 components"

let laplacian vars expr =
  let hess = hessian vars expr in
  let diagonal = List.mapi (fun i row -> List.nth row i) hess in
  let sum = List.fold_left (fun acc t -> Add (acc, t)) (Const 0.0) diagonal in
  simplify sum