open Diff
open Eval
open Multivariate

let newton_raphson expr var initial tolerance max_iter =
  let f_prime = diff var expr in
  let rec iterate x n =
    if n >= max_iter then None
    else
      let fx = eval [(var, x)] expr in
      if abs_float fx < tolerance then Some x
      else
        let fpx = eval [(var, x)] f_prime in
        if abs_float fpx < 1e-10 then None
        else
          let x_next = x -. fx /. fpx in
          if abs_float (x_next -. x) < tolerance then Some x_next
          else iterate x_next (n + 1)
  in
  iterate initial 0

let bisection expr var left right tolerance max_iter =
  let rec iterate a b n =
    if n >= max_iter then None
    else
      let fa = eval [(var, a)] expr in
      let fb = eval [(var, b)] expr in
      if fa *. fb > 0.0 then None
      else
        let mid = (a +. b) /. 2.0 in
        let fmid = eval [(var, mid)] expr in
        if abs_float fmid < tolerance || abs_float (b -. a) < tolerance then
          Some mid
        else if fa *. fmid < 0.0 then
          iterate a mid (n + 1)
        else
          iterate mid b (n + 1)
  in
  iterate left right 0

let trapezoidal expr var lower upper n =
  let h = (upper -. lower) /. float_of_int n in
  let rec sum_interior i acc =
    if i >= n then acc
    else
      let x = lower +. float_of_int i *. h in
      let fx = eval [(var, x)] expr in
      sum_interior (i + 1) (acc +. fx)
  in
  let f_lower = eval [(var, lower)] expr in
  let f_upper = eval [(var, upper)] expr in
  let interior = sum_interior 1 0.0 in
  h *. (f_lower /. 2.0 +. interior +. f_upper /. 2.0)

let simpsons expr var lower upper n =
  let n = if n mod 2 = 1 then n + 1 else n in
  let h = (upper -. lower) /. float_of_int n in
  let rec sum_terms i acc_odd acc_even =
    if i >= n then (acc_odd, acc_even)
    else
      let x = lower +. float_of_int i *. h in
      let fx = eval [(var, x)] expr in
      if i mod 2 = 1 then
        sum_terms (i + 1) (acc_odd +. fx) acc_even
      else if i > 0 then
        sum_terms (i + 1) acc_odd (acc_even +. fx)
      else
        sum_terms (i + 1) acc_odd acc_even
  in
  let f_lower = eval [(var, lower)] expr in
  let f_upper = eval [(var, upper)] expr in
  let (odd, even) = sum_terms 1 0.0 0.0 in
  h /. 3.0 *. (f_lower +. 4.0 *. odd +. 2.0 *. even +. f_upper)

let rec adaptive_quadrature expr var lower upper tolerance =
  let mid = (lower +. upper) /. 2.0 in
  let whole = simpsons expr var lower upper 10 in
  let left_half = simpsons expr var lower mid 10 in
  let right_half = simpsons expr var mid upper 10 in
  let error = abs_float (whole -. (left_half +. right_half)) in
  if error < tolerance then
    left_half +. right_half
  else
    let left = adaptive_quadrature expr var lower mid (tolerance /. 2.0) in
    let right = adaptive_quadrature expr var mid upper (tolerance /. 2.0) in
    left +. right

let gradient_descent expr vars initial learning_rate max_iter =
  let grad_exprs = gradient vars expr in
  let rec iterate point n =
    if n >= max_iter then Some point
    else
      let env = List.combine vars point in
      let grad_vals = List.map (eval env) grad_exprs in
      let new_point = List.map2 (fun p g -> p -. learning_rate *. g) point grad_vals in
      let diff = List.map2 (fun a b -> abs_float (a -. b)) new_point point in
      let max_diff = List.fold_left max 0.0 diff in
      if max_diff < 1e-6 then Some new_point
      else iterate new_point (n + 1)
  in
  iterate initial 0

let invert_matrix matrix =
  let n = List.length matrix in
  let augmented = List.mapi (fun i row ->
    row @ List.init n (fun j -> if i = j then 1.0 else 0.0)
  ) matrix in

  let rec gaussian_elimination mat row =
    if row >= n then mat
    else
      let pivot_row = List.nth mat row in
      let pivot = List.nth pivot_row row in
      if abs_float pivot < 1e-10 then mat
      else
        let normalized = List.map (fun x -> x /. pivot) pivot_row in
        let updated = List.mapi (fun i r ->
          if i = row then normalized
          else
            let factor = List.nth r row in
            List.map2 (fun a b -> a -. factor *. b) r normalized
        ) mat in
        gaussian_elimination updated (row + 1)
  in

  let reduced = gaussian_elimination augmented 0 in
  List.map (fun row -> List.filteri (fun i _ -> i >= n) row) reduced

let newtons_method_opt expr vars initial tolerance max_iter =
  let grad_exprs = gradient vars expr in
  let hess_matrix = hessian vars expr in

  let rec iterate point n =
    if n >= max_iter then None
    else
      let env = List.combine vars point in
      let grad_vals = List.map (eval env) grad_exprs in
      let hess_vals = List.map (fun row ->
        List.map (eval env) row
      ) hess_matrix in

      let hess_inv = invert_matrix hess_vals in
      let delta = List.map (fun row ->
        List.fold_left2 (fun acc h g -> acc +. h *. g) 0.0 row grad_vals
      ) hess_inv in

      let new_point = List.map2 (fun p d -> p -. d) point delta in
      let diff = List.map2 (fun a b -> abs_float (a -. b)) new_point point in
      let max_diff = List.fold_left max 0.0 diff in

      if max_diff < tolerance then Some new_point
      else iterate new_point (n + 1)
  in
  iterate initial 0