open Expr

let rec to_latex ?(implicit_mult=true) = function
  | Const f ->
      if Float.is_integer f then
        string_of_int (int_of_float f)
      else
        string_of_float f
  | SymConst Pi -> "\\pi"
  | SymConst E -> "e"
  | Var s -> s
  | Add (e1, e2) -> to_latex ~implicit_mult e1 ^ " + " ^ to_latex ~implicit_mult e2
  | Sub (e1, e2) -> to_latex ~implicit_mult e1 ^ " - " ^ to_latex_paren ~implicit_mult e2
  | Mul (e1, e2) ->
      let sep = if implicit_mult && should_implicit e1 e2 then "\\," else "\\cdot" in
      to_latex_mul ~implicit_mult e1 ^ sep ^ to_latex_mul ~implicit_mult e2
  | Div (e1, e2) -> "\\frac{" ^ to_latex ~implicit_mult e1 ^ "}{" ^ to_latex ~implicit_mult e2 ^ "}"
  | Pow (e1, e2) -> to_latex_atom ~implicit_mult e1 ^ "^{" ^ to_latex ~implicit_mult e2 ^ "}"
  | Neg e -> "-" ^ to_latex_atom ~implicit_mult e
  | Sin e -> "\\sin\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"
  | Cos e -> "\\cos\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"
  | Tan e -> "\\tan\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"
  | Sinh e -> "\\sinh\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"
  | Cosh e -> "\\cosh\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"
  | Tanh e -> "\\tanh\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"
  | Asin e -> "\\arcsin\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"
  | Acos e -> "\\arccos\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"
  | Atan e -> "\\arctan\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"
  | Atan2 (e1, e2) -> "\\text{atan2}\\left(" ^ to_latex ~implicit_mult e1 ^ ", " ^ to_latex ~implicit_mult e2 ^ "\\right)"
  | Exp e -> "e^{" ^ to_latex ~implicit_mult e ^ "}"
  | Ln e -> "\\ln\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"
  | Log (base_e, arg) -> "\\log_{" ^ to_latex ~implicit_mult base_e ^ "}\\left(" ^ to_latex ~implicit_mult arg ^ "\\right)"
  | Sqrt e -> "\\sqrt{" ^ to_latex ~implicit_mult e ^ "}"
  | Abs e -> "\\left|" ^ to_latex ~implicit_mult e ^ "\\right|"

and should_implicit e1 e2 =
  match (e1, e2) with
  | (Const _ | SymConst _), (Var _ | Sin _ | Cos _ | Tan _ | Sinh _ | Cosh _ | Tanh _ |
                              Asin _ | Acos _ | Atan _ | Exp _ | Ln _ | Log _ | Sqrt _ | Abs _ | Pow _) -> true
  | _ -> false

and to_latex_atom ?(implicit_mult=true) = function
  | (Const _ | SymConst _ | Var _) as e -> to_latex ~implicit_mult e
  | e -> "\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"

and to_latex_mul ?(implicit_mult=true) = function
  | (Const _ | SymConst _ | Var _ | Pow _ | Sin _ | Cos _ | Tan _ | Sinh _ | Cosh _ | Tanh _
    | Asin _ | Acos _ | Atan _ | Atan2 _ | Exp _ | Ln _ | Log _ | Sqrt _ | Abs _) as e -> to_latex ~implicit_mult e
  | e -> "\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"

and to_latex_paren ?(implicit_mult=true) = function
  | (Add _ | Sub _) as e -> "\\left(" ^ to_latex ~implicit_mult e ^ "\\right)"
  | e -> to_latex ~implicit_mult e

let to_graphviz expr =
  let counter = ref 0 in
  let next_id () =
    incr counter;
    !counter
  in

  let rec build_graph e =
    let id = next_id () in
    match e with
    | Const f ->
        let label = if Float.is_integer f then string_of_int (int_of_float f) else string_of_float f in
        (id, Printf.sprintf "  n%d [label=\"%s\" shape=oval];\n" id label, "")
    | SymConst Pi -> (id, Printf.sprintf "  n%d [label=\"π\" shape=oval];\n" id, "")
    | SymConst E -> (id, Printf.sprintf "  n%d [label=\"e\" shape=oval];\n" id, "")
    | Var v -> (id, Printf.sprintf "  n%d [label=\"%s\" shape=oval];\n" id v, "")
    | Add (e1, e2) -> build_binary id "+" e1 e2
    | Sub (e1, e2) -> build_binary id "-" e1 e2
    | Mul (e1, e2) -> build_binary id "*" e1 e2
    | Div (e1, e2) -> build_binary id "/" e1 e2
    | Pow (e1, e2) -> build_binary id "^" e1 e2
    | Neg e -> build_unary id "-" e
    | Sin e -> build_unary id "sin" e
    | Cos e -> build_unary id "cos" e
    | Tan e -> build_unary id "tan" e
    | Sinh e -> build_unary id "sinh" e
    | Cosh e -> build_unary id "cosh" e
    | Tanh e -> build_unary id "tanh" e
    | Asin e -> build_unary id "asin" e
    | Acos e -> build_unary id "acos" e
    | Atan e -> build_unary id "atan" e
    | Atan2 (e1, e2) -> build_binary id "atan2" e1 e2
    | Exp e -> build_unary id "exp" e
    | Ln e -> build_unary id "ln" e
    | Log (e1, e2) -> build_binary id "log" e1 e2
    | Sqrt e -> build_unary id "sqrt" e
    | Abs e -> build_unary id "abs" e

  and build_unary id op e =
    let (eid, enodes, eedges) = build_graph e in
    let node = Printf.sprintf "  n%d [label=\"%s\" shape=diamond];\n" id op in
    let edge = Printf.sprintf "  n%d -> n%d;\n" id eid in
    (id, node ^ enodes, edge ^ eedges)

  and build_binary id op e1 e2 =
    let (id1, nodes1, edges1) = build_graph e1 in
    let (id2, nodes2, edges2) = build_graph e2 in
    let node = Printf.sprintf "  n%d [label=\"%s\" shape=box];\n" id op in
    let edge1 = Printf.sprintf "  n%d -> n%d [label=\"L\"];\n" id id1 in
    let edge2 = Printf.sprintf "  n%d -> n%d [label=\"R\"];\n" id id2 in
    (id, node ^ nodes1 ^ nodes2, edge1 ^ edge2 ^ edges1 ^ edges2)
  in

  let (_, nodes, edges) = build_graph expr in
  "digraph expr {\n" ^ nodes ^ edges ^ "}\n"