open Expr
open Simplify

type special_function =
  | Gamma of expr
  | Beta of expr * expr
  | Erf of expr
  | Erfc of expr
  | BesselJ of int * expr
  | BesselY of int * expr
  | LegendreP of int * expr
  | HermiteH of int * expr
  | LaguerreL of int * expr
  | ChebyshevT of int * expr
  | ChebyshevU of int * expr

let gamma x = Gamma x
let beta x y = Beta (x, y)
let erf x = Erf x
let erfc x = Erfc x

let bessel_j n x = BesselJ (n, x)
let bessel_y n x = BesselY (n, x)

let legendre_p n x =
  let rec compute n =
    if n = 0 then Const 1.0
    else if n = 1 then x
    else
      let p_n_1 = compute (n - 1) in
      let p_n_2 = compute (n - 2) in
      let n_f = float_of_int n in
      simplify (Div (
        Sub (
          Mul (Const (2.0 *. n_f -. 1.0), Mul (x, p_n_1)),
          Mul (Const (n_f -. 1.0), p_n_2)
        ),
        Const n_f
      ))
  in
  compute n

let hermite_h n x =
  let rec compute n =
    if n = 0 then Const 1.0
    else if n = 1 then Mul (Const 2.0, x)
    else
      let h_n_1 = compute (n - 1) in
      let h_n_2 = compute (n - 2) in
      simplify (Sub (
        Mul (Const 2.0, Mul (x, h_n_1)),
        Mul (Const (2.0 *. float_of_int (n - 1)), h_n_2)
      ))
  in
  compute n

let laguerre_l n x =
  let rec compute n =
    if n = 0 then Const 1.0
    else if n = 1 then Sub (Const 1.0, x)
    else
      let l_n_1 = compute (n - 1) in
      let l_n_2 = compute (n - 2) in
      let n_f = float_of_int n in
      simplify (Div (
        Sub (
          Mul (Const (2.0 *. n_f -. 1.0 -. 1.0), l_n_1),
          Sub (Mul (x, l_n_1), Mul (Const (n_f -. 1.0), l_n_2))
        ),
        Const n_f
      ))
  in
  compute n

let chebyshev_t n x =
  let rec compute n =
    if n = 0 then Const 1.0
    else if n = 1 then x
    else
      let t_n_1 = compute (n - 1) in
      let t_n_2 = compute (n - 2) in
      simplify (Sub (
        Mul (Const 2.0, Mul (x, t_n_1)),
        t_n_2
      ))
  in
  compute n

let chebyshev_u n x =
  let rec compute n =
    if n = 0 then Const 1.0
    else if n = 1 then Mul (Const 2.0, x)
    else
      let u_n_1 = compute (n - 1) in
      let u_n_2 = compute (n - 2) in
      simplify (Sub (
        Mul (Const 2.0, Mul (x, u_n_1)),
        u_n_2
      ))
  in
  compute n

let factorial n =
  let rec fact n acc =
    if n <= 1 then acc
    else fact (n - 1) (acc * n)
  in
  Const (float_of_int (fact n 1))

let binomial n k =
  if k < 0 || k > n then Const 0.0
  else
    let rec binom n k =
      if k = 0 || k = n then 1
      else binom (n - 1) (k - 1) + binom (n - 1) k
    in
    Const (float_of_int (binom n k))