open Expr
open Simplify

type domain =
  | Real
  | Complex
  | Positive
  | Negative
  | NonNegative
  | NonPositive
  | Integer
  | Natural
  | Rational
  | Even
  | Odd

type assumption = (string * domain) list

let assume var domain assumptions =
  (var, domain) :: List.remove_assoc var assumptions

let get_domain var assumptions =
  List.assoc_opt var assumptions

let is_compatible domain1 domain2 =
  match (domain1, domain2) with
  | d1, d2 when d1 = d2 -> true
  | Positive, (Real | Complex | NonNegative | Rational) -> true
  | Negative, (Real | Complex | NonPositive | Rational) -> true
  | NonNegative, (Real | Complex | Rational) -> true
  | NonPositive, (Real | Complex | Rational) -> true
  | Integer, (Real | Complex | Rational) -> true
  | Natural, (Integer | Real | Complex | NonNegative | Rational) -> true
  | Even, (Integer | Real | Complex | Rational) -> true
  | Odd, (Integer | Real | Complex | Rational) -> true
  | _ -> false

let refine assumptions _expr _condition =
  Some assumptions

let simplify_with assumptions expr =
  let rec simplify_expr = function
    | Sqrt (Pow (Var v, Const 2.0)) ->
        (match get_domain v assumptions with
         | Some Positive | Some NonNegative -> Var v
         | Some Negative | Some NonPositive -> Neg (Var v)
         | _ -> Abs (Var v))
    | Abs (Var v) as e ->
        (match get_domain v assumptions with
         | Some Positive | Some NonNegative -> Var v
         | Some Negative | Some NonPositive -> Neg (Var v)
         | _ -> e)
    | Pow (Var v, Const n) as e when Float.is_integer n && int_of_float n mod 2 = 0 ->
        (match get_domain v assumptions with
         | Some Positive | Some NonNegative | Some Negative | Some NonPositive -> e
         | _ -> e)
    | Add (e1, e2) -> Add (simplify_expr e1, simplify_expr e2)
    | Sub (e1, e2) -> Sub (simplify_expr e1, simplify_expr e2)
    | Mul (e1, e2) -> Mul (simplify_expr e1, simplify_expr e2)
    | Div (e1, e2) -> Div (simplify_expr e1, simplify_expr e2)
    | Pow (e1, e2) -> Pow (simplify_expr e1, simplify_expr e2)
    | Neg e -> Neg (simplify_expr e)
    | Sin e -> Sin (simplify_expr e)
    | Cos e -> Cos (simplify_expr e)
    | Tan e -> Tan (simplify_expr e)
    | Sinh e -> Sinh (simplify_expr e)
    | Cosh e -> Cosh (simplify_expr e)
    | Tanh e -> Tanh (simplify_expr e)
    | Asin e -> Asin (simplify_expr e)
    | Acos e -> Acos (simplify_expr e)
    | Atan e -> Atan (simplify_expr e)
    | Atan2 (e1, e2) -> Atan2 (simplify_expr e1, simplify_expr e2)
    | Exp e -> Exp (simplify_expr e)
    | Ln e -> Ln (simplify_expr e)
    | Log (e1, e2) -> Log (simplify_expr e1, simplify_expr e2)
    | Sqrt e -> Sqrt (simplify_expr e)
    | Abs e -> Abs (simplify_expr e)
    | e -> e
  in
  simplify (simplify_expr expr)

let verify_inequality _expr _assumptions =
  None