open Unification

(* ========================================
     Definitions
   ======================================== *)

type pol = Pos | Neg | Npol
type ray = Var of id | Func of (id * pol * ray list)
(* alternative ray definition using terms *)
(* type ray = PR of id * pol * ray | NR of term *)
type star = ray list
type constellation = star list
type graph = (int * int) * (ray * ray) list

(* List monad *)
let return x = [x] (*plongement dans la monade de liste*)
let (>>=) xs k = List.flatten (List.map k xs)
let guard c x = if c then return x else []

(* ========================================
     Useful functions
   ======================================== *)

(* Convert a pol and an id to a string, adding + or - before the id *)
let pol_to_string pol id =
    if pol = Pos then "+" ^ id 
    else if pol = Neg then "-" ^ id 
    else id

(* Convert a ray (which is polarized) to a term *)
let rec ray_to_term r =
  match r with
  | Var id -> (Var(id) : term)
  | Func(id, pol, raylist) -> (Func(pol_to_string pol id, List.map ray_to_term raylist) : term)

(* Invert polarization of a pol*)
let inv_pol pol = 
    if pol = Pos then Neg
    else if pol = Neg then Pos 
    else pol

(* Invert the polarization of a ray to allow an easier Unification writing *)
let rec inv_pol_ray ray = 
    match ray with
    | Func(id, pol, raylist) -> Func(id, inv_pol pol, List.map inv_pol_ray raylist)
    | _ -> ray

(* Checks if a ray is polarised *)
let rec is_polarised r =
    match r with
    | Var id -> false
    | Func(_, p, r) -> (p <> Npol) || (List.fold_left (fun acc b -> (is_polarised b) || acc) false r)

(* Checks if two rays are dual, meaning that after inverting polarization of one ray, the two rays can be unified *)
let dual_check r1 r2 =
  if (is_polarised r1 && is_polarised r2) then 
  (solve [(extends_varname (ray_to_term (inv_pol_ray r1)) "0"), (extends_varname ((ray_to_term r2)) "1")] []) 
  else None

(* Create an index for a constellation *)
let index_constellation const = 
    List.combine (List.init (List.length const) (fun a -> a)) const

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

(* substit_ray applies all possible substition from an environment to a ray *)
let rec substit_ray ray sub =
  match ray with
  | Var id -> apply_ray id sub
  | Func(f, p, tl) -> Func(f, p, List.map (fun a -> substit_ray a sub) tl) 

(* substit_star applies all possible substition from an environment to a star *)
let substit_star star sub =
  List.map (fun a -> substit_ray a sub) star

(* substit_const applies all possible substition from an environment to a constellation *)
let substit_const const sub =
  List.map (fun a -> substit_star a sub) const

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

(* extends_varname adds suffix to all var names of a star *)
let extends_varname_star const ext =
  List.map (fun a -> extends_varname_ray a ext) const

(* extends_varname adds suffix to all var names of a constellation based on each star number after being indexed *)
let extends_varname_const const =
  List.map (fun (i,a) -> extends_varname_star a (string_of_int i)) (index_constellation const)

(* convert a term to a ray *)
let rec term_to_ray (term : term) =
  match term with
  | Var id -> (Var(id) : ray)
  | Func(f, r) -> Func(f, Npol, List.map (fun a -> term_to_ray a) r)

(* ========================================
     Constellation graph
   ======================================== *)

(* Makes a dgraph from a constellation *)
let dgraph const =
  let indexed_const = index_constellation const in 
  indexed_const >>= fun (i, il) ->
  indexed_const >>= fun (j, jl) ->
  il >>= fun r1 ->
  jl >>= fun r2 ->
  guard (j >= i) ( let uni = dual_check r1 r2 in 
                   if Option.is_some uni then [((i,j),(r1,r2))]
                   else [])

(* Convert a link to a string to be printable *)
let link_to_string dg = 
  let rec aux dgl =
    match dgl with
    | [] -> ""
    | ((i,j),(r1, r2))::[] -> ("(" ^ string_of_int i ^ ", " ^ string_of_int j ^ ")" ^ "," ^ "(" ^ term_to_string (ray_to_term r1) ^ ", " ^ term_to_string (ray_to_term r2) ^ ")")
    | ((i,j),(r1, r2))::t -> ("(" ^ string_of_int i ^ ", " ^ string_of_int j ^ ")" ^ "," ^ "(" ^ term_to_string (ray_to_term r1) ^ ", " ^ term_to_string (ray_to_term r2) ^ ")") ^ "+" ^ (aux t)
  in aux dg ;;

(* Print a dgraph *)
let print_dgraph dg =
  let rec aux dgl =
    match dgl with
    | [] -> ""
    | h::[] -> (link_to_string h)
    | h::t -> (link_to_string h) ^ "\n" ^ aux t 
  in print_string (aux dg);;

let clean_dgraph g =
  List.filter (fun a -> a <> []) g

(* _________ Examples _________ *)
let make_const_pol pol c = Func (c, pol, [])
let make_const c = make_const_pol Npol c

let y = Var("y")
let x = Var("x")
let z = Var("z")
let r = Var("r")
let zero = make_const "0"
let s x = Func("s", Npol, [x])
let add p x y z = Func("add", p, [x;y;z])

(* Convert int to term *)
let rec enat i = 
  if i = 0 then zero else s (enat (i-1))

(* makes the constellation corresponding to an addition *)
let make_const_add n m = 
  [[add Pos zero y y]; [add Neg x y z; add Pos (s x) y (s z)]; [add Neg (enat n) (enat m) r; r]] 

let constellation = make_const_add 1 3 ;;

print_dgraph (dgraph constellation) ;;

(* token is a couple of a family number and a star number in the constellation *)
type token = int * int
type process = token list

(* get a star using its number in the list from a constellation *)
let get_star const i =
  List.nth const i

(* Takes a constellation, a ray and a (ray,ray) list and extracts rays from stars number i (respectively j) that are  not ri (respectively rj) when ri (respectively rj) isn't in the prob list *)
let star_filter const ((i, j),(ri,rj)) prob =
  let (prob_a, prob_b) = List.split prob in
  (if List.mem ri prob_a then []
  else (List.filter (fun a -> a <> ri) (get_star const i))
  )@( 
  if List.mem rj prob_b then []
  else (List.filter (fun a -> a <> rj) (get_star const j))
  )

(* convert the (ray,ray) list part of a link to an equation, converting its rays to terms *)
let link_to_eq prob =
  List.map (fun (ra, rb) -> (ray_to_term (inv_pol_ray ra)), ray_to_term rb) prob

(* removes prob rays from stars *)
let star_postfilter star prob =
  let (prob_a, prob_b) = List.split prob in
  List.filter (fun a -> not(List.mem a prob_a || List.mem a prob_b )) star

(* takes a token, a graph and a constellation and returns the list of tokens to check next and a list of solvable equation *)
let divide_token (fam, n_star) toklist graph const prob fstar checked_stars =
  let links = List.filter (fun ((i, _),(_, _)) -> i = n_star) graph in
    let rec aux l tokl prob_aux star_aux checked_stars_aux =
      match l with
        [] -> Some (tokl,prob_aux,star_aux,checked_stars_aux,fam)
        | ((i, j),(ri,rj))::tl -> 
        if fam > (List.length prob_aux) || prob_aux = [] then (* We check if the family number is the same as the number of equations lists in prob. If it's superior, we add a new list in prob instead of filling the first equation list because it means we're treating a new family *)
          if Option.is_some ((solve (link_to_eq [(ri, rj)])) []) then 
            aux tl ((fam, j)::tokl) ([(ri, rj)]::prob_aux) ( (( star_filter const ((i, j),(ri,rj)) [] ))::star_aux ) (i::checked_stars_aux)
          else None
        else 
          if Option.is_some (solve (link_to_eq ((ri, rj)::(List.hd prob_aux))) []) then (* We made sure prob_aux head would not be empty*)
            aux tl ((fam, j)::tokl) (((ri, rj)::(List.hd prob_aux))::(List.tl prob_aux)) ( (( star_filter const ((i, j),(ri,rj)) (List.hd prob_aux) )@(List.hd star_aux))::(List.tl star_aux) ) (i::checked_stars_aux) (*We use List.hd because the current family we're working on should be the current first*)
          else None
    in if links = [] then Some (toklist,prob,fstar,n_star::checked_stars,fam) 
    else aux links toklist prob fstar checked_stars

(* should be deterministic exec, graph shouldn't be empty *)
let exec const =
  let const_ext = extends_varname_const const in
  let graph = List.flatten (clean_dgraph (dgraph const_ext)) in
  let rec aux (toklist,prob,star,checked_stars,current_fam) =
    begin match toklist with
      | [] -> let gen_token = List.filter (fun ((i,_),(_,_)) -> not( List.mem i checked_stars)) graph in 
          if gen_token = [] then star,prob
          else let ((i,_),(_,_)) = (List.hd gen_token) in aux ([(current_fam+1,i)],prob,star,checked_stars,current_fam+1)
      | h::t -> aux (Option.get (divide_token h t graph const_ext prob star checked_stars))
    end              
  in let ((i,_),(_,_)) = (List.hd graph) 
  in let (constf, prob_tmp) = aux ([(0,i)],[],[],[],0) 
  in let probf = List.rev prob_tmp
  in let indexed_final_const = index_constellation constf
  in List.map (fun (fam_star, a) -> let fam_prob = List.nth probf fam_star 
                in substit_star (star_postfilter a fam_prob) (List.map (fun (i,b) -> (i,term_to_ray b)) (Option.get (solve (link_to_eq fam_prob) [])))) indexed_final_const

(* test constellation non-cyclique déterministe *)
let test = [ [Func("c", Neg, [x]); x] ; [Func("c", Pos, [Func("f", Npol, [y])]) ; Func("c", Npol, [x]) ] ] ;;
print_dgraph (dgraph test);;
exec test ;;

(* TODO :
faire un filtre post application, const_filter
*)

(* not flattened graph ver. :
(* takes a token, a graph and a constellation and returns the list of tokens to check next and a list of solvable equation *)
let divide_token (fam, n_star) graph const =
  let rec aux g toklist prob fstar checked_stars =
    match g with
    | [] -> Some (toklist,prob,fstar,checked_stars,fam)
    | h::t -> let links = List.filter (fun ((i, _),(_, _)) -> i = n_star) h in
        let rec aux2 l tokl prob2 star2 checked_stars2 =
          match l with
          | [] -> Some (tokl,prob2,star2,checked_stars2,fam)
          | ((i, j),(ri,rj))::tl -> 
              if fam > (List.length prob2) || prob2 = [] then (* We check if the family number is the same as the number of equations lists in prob. If it's superior, we add a new list in prob instead of filling the first equation list *)
                if Option.is_some ((solve (link_to_eq [(ri, rj)])) []) then 
                  aux2 tl ((fam, j)::tokl) ([(ri, rj)]::prob2) ( (( star_filter const ((i, j),(ri,rj)) [] ))::star2 ) (i::checked_stars2)
                else None
              else 
              if Option.is_some (solve (link_to_eq ((ri, rj)::(List.hd prob2))) []) then (* *)
                aux2 tl ((fam, j)::tokl) (((ri, rj)::(List.hd prob2))::(List.tl prob2)) ( (( star_filter const ((i, j),(ri,rj)) (List.hd prob2) )@(List.hd star2))::(List.tl star2) ) (i::checked_stars2) (*We use List.hd becasue the current family we're working on should be the current first*)
              else None
        in if links = [] then aux t toklist prob fstar checked_stars else aux2 links toklist prob fstar checked_stars
  in aux graph [] [] [] [] 

(* should be deterministic exec, graph shouldn't be empty *)
let exec const =
  let const_ext = extends_varname_const const in
  let graph = clean_dgraph (dgraph const_ext) in
  let rec aux (toklist,prob,star,checked_stars,current_fam) =
    begin match toklist with
      | [] -> let gen_token = List.filter (fun ((i,_),(_,_)) -> not( List.mem i checked_stars)) (List.flatten graph) in 
          if gen_token = [] then star,prob
          else let ((i,_),(_,_)) = (List.hd gen_token) in aux ([(current_fam+1,i)],prob,star,checked_stars,current_fam+1)
      | h::t -> aux (Option.get (divide_token h graph const_ext))
    end              
  in let ((i,_),(_,_)) = (List.hd (List.hd graph)) 
  in let (constf, prob_tmp) = aux ([(0,i)],[],[],[],0) 
  in let probf = List.rev prob_tmp
  in let indexed_final_const = index_constellation constf
  in List.map (fun (fam_star, a) -> let fam_prob = List.nth probf fam_star 
                in substit_star (star_postfilter a fam_prob) (List.map (fun (i,b) -> (i,term_to_ray b)) (Option.get (solve (link_to_eq fam_prob) [])))) indexed_final_const

*)