open Expr
open Integrate

let fourier_transform expr var =
  let omega = "omega" in
  let integrand = Mul (expr, Exp (Mul (Mul (Const (-1.0), SymConst E), Mul (Var omega, Var var)))) in
  match integrate var integrand with
  | Some result -> Some result
  | None ->
      match expr with
      | Const c -> Some (Mul (Const c, Const 0.0))
      | Exp (Mul (Const a, Var v)) when v = var && a < 0.0 ->
          Some (Div (Const 1.0, Add (Const a, Mul (SymConst E, Var omega))))
      | Sin (Mul (Const a, Var v)) when v = var ->
          Some (Div (Const a, Sub (Pow (Var omega, Const 2.0), Pow (Const a, Const 2.0))))
      | Cos (Mul (Const a, Var v)) when v = var ->
          Some (Div (Var omega, Sub (Pow (Var omega, Const 2.0), Pow (Const a, Const 2.0))))
      | _ -> None

let inverse_fourier_transform _expr _omega =
  None

let laplace_transform expr t =
  let s = "s" in
  match expr with
  | Const c -> Some (Div (Const c, Var s))
  | Var v when v = t -> Some (Div (Const 1.0, Pow (Var s, Const 2.0)))
  | Pow (Var v, Const n) when v = t && Float.is_integer n && n >= 0.0 ->
      let rec factorial n =
        if n <= 1.0 then 1.0
        else n *. factorial (n -. 1.0)
      in
      Some (Div (Const (factorial n), Pow (Var s, Const (n +. 1.0))))
  | Exp (Mul (Const a, Var v)) when v = t ->
      Some (Div (Const 1.0, Sub (Var s, Const a)))
  | Sin (Mul (Const a, Var v)) when v = t ->
      Some (Div (Const a, Add (Pow (Var s, Const 2.0), Pow (Const a, Const 2.0))))
  | Cos (Mul (Const a, Var v)) when v = t ->
      Some (Div (Var s, Add (Pow (Var s, Const 2.0), Pow (Const a, Const 2.0))))
  | Sinh (Mul (Const a, Var v)) when v = t ->
      Some (Div (Const a, Sub (Pow (Var s, Const 2.0), Pow (Const a, Const 2.0))))
  | Cosh (Mul (Const a, Var v)) when v = t ->
      Some (Div (Var s, Sub (Pow (Var s, Const 2.0), Pow (Const a, Const 2.0))))
  | Mul (Exp (Mul (Const a, Var v1)), Sin (Mul (Const b, Var v2)))
    when v1 = t && v2 = t ->
      Some (Div (Const b,
        Add (Pow (Sub (Var s, Const a), Const 2.0), Pow (Const b, Const 2.0))))
  | Mul (Exp (Mul (Const a, Var v1)), Cos (Mul (Const b, Var v2)))
    when v1 = t && v2 = t ->
      Some (Div (Sub (Var s, Const a),
        Add (Pow (Sub (Var s, Const a), Const 2.0), Pow (Const b, Const 2.0))))
  | _ -> None

let inverse_laplace_transform expr s =
  let t = "t" in
  match expr with
  | Div (Const c, Var v) when v = s -> Some (Const c)
  | Div (Const 1.0, Pow (Var v, Const n)) when v = s && Float.is_integer n && n > 0.0 ->
      let rec factorial n =
        if n <= 1.0 then 1.0
        else n *. factorial (n -. 1.0)
      in
      Some (Div (Pow (Var t, Const (n -. 1.0)), Const (factorial (n -. 1.0))))
  | Div (Const 1.0, Sub (Var v, Const a)) when v = s ->
      Some (Exp (Mul (Const a, Var t)))
  | Div (Const a, Add (Pow (Var v, Const 2.0), Pow (Const b, Const 2.0))) when v = s ->
      Some (Mul (Const (a /. b), Sin (Mul (Const b, Var t))))
  | Div (Var v, Add (Pow (Var v2, Const 2.0), Pow (Const b, Const 2.0)))
    when v = s && v2 = s ->
      Some (Cos (Mul (Const b, Var t)))
  | _ -> None

let z_transform _expr _n =
  None

let inverse_z_transform _expr _z =
  None