open List
(* ========================================
     Definitions
   ======================================== *)

type id = string
type term = Var of id | Func of (id * term list)

type subst = id * term list
type equation = term * term

(* convert a term to a string *)
let rec term_to_string t =
   match t with
   | Var id -> id
   | Func(f, []) -> f
   | Func(f, tl) -> f ^ "(" ^
      let rec aux2 vl =
        match vl with
        | [] -> ""
        | h::[] -> term_to_string h
        | h::t -> (term_to_string h) ^ "," ^ (aux2 t) 
      in (aux2 tl) ^ ")"

let print_term t =
  print_string (term_to_string t)

(* Compare two terms *)
let rec compare_term t1 t2 =
   match t1, t2 with
   | Var(id1), Var(id2) -> String.compare id1 id2
   | Var(_), Func(_,_) -> -1
   | Func(_,_), Var(_) -> 1
   | Func(f, fs), Func(g, gs) -> let comp = String.compare f g in if comp > 0 then 1 
   else if comp < 0 then -1 
   else List.compare compare_term fs gs

(* Look if there's a var to be substituted in the term in the substitution environment *)
let rec indom t sl = 
  match t with
  | Var id -> List.exists (fun (a,_) -> a = id ) sl
  | Func(_, tl) -> List.exists (fun a -> indom a sl) tl

(* occurs checks if given var is in term *)
let rec occurs id t =
  match t with
  | Var i -> i = id 
  | Func(_, tl) -> List.exists (fun a -> occurs id a) tl


(* extends_varname adds suffix to all var names of a term *)
let rec extends_varname t ext =
  match t with
  | Var id -> Var(id ^ ext)
  | Func(f, tl) -> Func(f, List.map (fun a -> extends_varname a ext) tl)


(* vars gives a list of all vars in a term *)
let rec vars t =
  match t with
  | Var id -> [id]
  | Func(_, tl) -> List.fold_left (fun a b -> (vars b)@a) [] tl

  (* ========================================
      Substitution
     ======================================== *)

(* apply applies a substitution to a var *)
let apply id sub = let (_,s) = try List.find (fun (a,_) -> a = id ) sub with Not_found -> (id,Var(id)) in s

(* subst applies all possible substition from an environment to a term *)
let rec substit trm sub =
  match trm with
  | Var id -> apply id sub
  | Func(f,tl) -> Func(f,List.map (fun a -> substit a sub) tl)  

  (* ========================================
      Unification
     ======================================== *)

(* Solves an equation list by returning solution list *)
let rec solve eq sub = 
  match eq with
  | [] -> Some sub
  | (Var(x), term)::t -> if Var(x) = term then (solve t sub) else (elim x term t sub) (* If x = x it's a useless equation *)
  | (term, Var(x))::t -> (elim x term t sub) (*It's useless to check if term = Var(x) because it would be the same case as above *)
  | (Func(f, fs), Func(g, gs))::t -> if f = g then (solve ((List.combine fs gs)@t) sub) else (*failwith (Printf.sprintf "f=%s g=%s" f g)*) None (* If f and g are not equal, the equation can't be solved *)

and elim id term eq sub =
  if occurs id term then (*failwith (Printf.sprintf "id=%s is in term" id)*) None (* If that's the case, we would have something like x = f(x) which can't be solved *)
  else let sigma_xy = [(id, term)] in
  solve (List.map (fun (a,b) -> (substit a sigma_xy, substit b sigma_xy)) eq) ((List.map (fun (i, t) -> (i, substit t sigma_xy)) sub)@sigma_xy) (* We apply the sigma_xy substitution to the equations and the solution list and we add it to the solution list *)