open Fol
open Alloyast
(* equivalence... *)



(* Relations *)

type relation_corps = [
| `Form of fol_formula
| `Rel of relation
]

and relation = ident -> relation_corps



(* Built-in relations *)

let univ
    : relation =

  (fun _ -> `Form `True)



let none
    : relation =
  (fun _ -> `Form `False)



let iden
    : relation =

  (fun x -> `Rel (fun y -> `Form (`Equality ((`Unknown, `Variable x),
  (`Unknown, `Variable y)))))



(* atome: Creates a relation from an atomic object (which is always a
   variable name introduced by a quantifier). *)

let atome 
    (i : ident) (* Name of the atom *)
    : relation =

  (fun x -> `Form (`Equality ((`Unknown, `Variable x), (`Unknown, `Variable i))))


let rec op o f1 f2 = 
  fun x -> match f1 x, f2 x with
    | `Form u, `Form v -> `Form (o u v)
    | `Rel u, `Rel v -> `Rel (op o u v)
    | _ -> invalid_arg "Fol.op: operation between different arities"


(* si_formule: Alloy implies/else operator
   si_relation: Alloy if/then/else operator *)

let si_formule
    (hyp : fol_formula)
    (vrai : fol_formula)
    (faux : fol_formula)
    : fol_formula =

  `Conjunction [
    implication hyp vrai;
    implication (`Negation hyp) faux;
  ]



let si_relation
    (cond : fol_formula)
    : relation -> relation -> relation =

  op (si_formule cond)
    


(* plus: Alloy + operator *)

let plus 
    : relation -> relation -> relation =

  op (fun a b -> `Disjunction [a; b])



(* moins: Alloy - operator *)

let moins 
    : relation -> relation -> relation =

  op (fun a b -> `Conjunction [a; `Negation b])



(* inter: Alloy & operator *)

let inter
    : relation -> relation -> relation =

  op (fun a b -> `Conjunction [a; b])




(* tilde: Alloy ~ operator *)

let tilde
    (f : relation)
    : relation =
        
        fun x -> `Rel (fun y -> match f y with
          | `Form _ -> invalid_arg "Fol.tilde: cannot invert unary relation"
          | `Rel g -> g x
        )



(* jointure: Alloy . operator *)

let rec jointure
    (f : relation)
    (g : relation)
    : relation =
  
  let t = nouveau_symbole "jointure" in
  match g t with
    |  `Rel h ->
        begin match f t with
          | `Form p -> begin function x ->
              let rec aux_h = function
                | `Form phi -> `Form (`Exists ([t, `Unknown], `Conjunction [p; phi]))
                | `Rel k -> `Rel (fun y -> aux_h (k y))
              in aux_h (h x)
            end
          | _ -> fun x -> 
              begin match f x with
                | `Form _ -> failwith "Fol.jointure: impossible (arity dependency)"
                | `Rel h ->  `Rel (jointure h g)
              end
        end
    | `Form q -> begin function x ->
        match f x with
          | `Form _ -> failwith "Fol.jointure: impossible (both unary)"
          | `Rel h -> begin match h t with
              | `Form p -> `Form (`Exists ([t, `Unknown], `Conjunction [p; q]))
              | `Rel _ -> `Rel (jointure h g)
            end
      end
          



(* fleche: Alloy -> operator *)

let rec fleche 
    (f : relation)
    (g : relation)
    : relation =

  fun x -> match f x with
    | `Form p ->
        let rec aux = function
          | `Form q -> `Form (`Conjunction [p; q])
          | `Rel h -> `Rel (fun z -> aux (h z))
        in
        `Rel (fun y -> aux (g y))
    | `Rel h -> `Rel (fleche h g)



(* restriction_gauche: Alloy <: operator *)

let restriction_gauche
   (cond : relation)
   (f : relation) 
   :  relation =

   begin fun x -> match cond x with
     | `Rel _ -> invalid_arg "Fol.restriction_gauche: non-unary restriction domain"
     | `Form hyp -> let rec aux = function
         | `Form p -> `Form (`Conjunction [hyp; p])
         | `Rel h -> `Rel (fun y -> aux (h y))
       in aux (f x)
   end



(* restriction_droite: Alloy :> operator *)

let rec restriction_droite
    (f : relation)
    (cond : relation)
    : relation =

  fun x -> match f x with
    | `Form p -> begin match cond x with
        | `Rel _ -> invalid_arg "Fol.restriction_droite: non-unary restriction domain"
        | `Form hyp -> `Form (`Conjunction [p; hyp])
      end
    | `Rel g -> `Rel (restriction_droite g cond)




(* Multiplicative formulae (mult r) where r is a relation, and mult is one
   of the following : some, no, lone, one. *)

let mult_some
    : relation -> fol_formula =

  let rec aux accu = function
    | `Form p -> `Exists (accu, p)
    | `Rel f -> 
        let x = nouveau_symbole "mult_some" in aux ((x, `Unknown)::accu) (f x)
  in fun r -> aux [] (`Rel r)

let mult_no 
    (f : relation)
    : fol_formula =

  `Negation (mult_some f)



let mult_lone
    (h : relation) 
    : fol_formula =
  
  let rec aux vars eqs = function
    | `Form p, `Form q ->
        `Forall (vars, implication (`Conjunction [p; q]) (`Conjunction eqs))
    | `Rel f, `Rel g ->
        let x = nouveau_symbole "mult_lone_x" in
        let y = nouveau_symbole "mult_lone_y" in
        aux ((x, `Unknown)::(y, `Unknown)::vars ) ((`Equality ((`Unknown,
        `Variable x), (`Unknown, `Variable y)))::eqs) (f x, g y)
    | _ -> failwith "mult_lone: BUG (this error raised is theoretically a dead code)"
  in aux [] [] (`Rel h, `Rel h)



let mult_one
    (h : relation)
    : fol_formula = 

  `Conjunction [mult_some h; mult_lone h]



(* Inclusion and equality *)

(* inclusion: Relation inclusion *)

let inclusion
    (r1 : relation) (* sub *)
    (r2 : relation) (* super *)
    : fol_formula =

  let rec aux accu = function
    | `Form p, `Form q ->
        `Forall (accu, implication p q)
    | `Rel f, `Rel g ->
        let x = nouveau_symbole "inclusion" in
        aux ((x, `Unknown)::accu) (f x, g x)
    | _ -> invalid_arg "Fol.inclusion: inclusion between different arities"
  in aux [] (`Rel r1, `Rel r2)



(* egalite : Relation equality *)

let egalite
    (r1 : relation)
    (r2 : relation)
    : fol_formula =
  
  `Conjunction [inclusion r1 r2; inclusion r2 r1]



(* funconstr: Additional constraint for functions *)

let relation_fonction
    (precond : ident -> fol_formula)
    (f : ident) 
    : fol_formula =

  let x = nouveau_symbole ("funconstr_" ^ f) in
  `Forall ([x, `Unknown], implication (precond x) (`Predicate ("def__" ^ f, [`Unknown, `Variable x])))


(* disjoints : The relations in the given list are disjoint one to the
   other. Actually, this function computes a list of formulae it appends to
   a stack given in input (avoids list concatenation)

   A sequence (A_i) of relations is disjoint one to the other if and only
   if for any i, A_1+...+A_i-1 and A_i are disjoint.

 *)

let rec disjoints
    (accu : fol_formula list)
    : relation list -> fol_formula list = function
        
      | a::b::q -> disjoints ((mult_no (inter a b))::accu) ((plus a b)::q)
      | _ -> accu



(* plusplus: Alloy ++ operator (overriding relations) *)

let plusplus
    (a : relation) 
    (b : relation)
    : relation =

  fun x -> match a x, b x with
    | `Form p, `Form q -> `Form (`Disjunction [p; q])
    | `Rel f, `Rel g -> `Rel (si_relation (mult_some g) g f)
    | _ -> invalid_arg "Fol.plusplus"




(* Processing an Alloy module. *)

(* The environment where signatures, predicates and facts will be stored. *)

exception Introuvable of string



(* Generic empty environment.
   Must be a function because of env_token and env_modul functions. *)

let empty_env
    (_ : 'a)
    : 'token env =

  {
    env_noms_rels = [];
    env_rels = (fun _ -> begin
      prerr_endline "Empty env";
      None end);
    env_types = (fun _ -> []);
    env_peres = [];
    env_fils = (fun _ -> []);
    env_abstract = [];
    env_contraintes = [];
    env_predicats = [];
    env_fonctions = [];
    env_theoremes = [];
    env_searchpath = [];
    env_modulenames = [];
    env_token = (fun _ -> failwith "Fol.env_token: lexer is undefined");
    env_modul = (fun _ _ x -> prerr_endline "Fol.env_modul: warning: parser is undefined"; x);
  }


(* Function to copy an environment, forgetting about parser and lexer.
   This (unused in practice) could be useful for multi-parsing...
*)

let copy_env
    (env : 'token1 env)
    : 'token2 env =

  { (empty_env ()) with
    env_noms_rels = env.env_noms_rels;
    env_rels = env.env_rels;
    env_types = env.env_types;
    env_peres = env.env_peres;
    env_fils = env.env_fils;
    env_abstract = env.env_abstract;
    env_contraintes = env.env_contraintes;
    env_predicats = env.env_predicats;
    env_fonctions = env.env_fonctions;
    env_theoremes = env.env_theoremes;
    env_searchpath = env.env_searchpath;
    env_modulenames = env.env_modulenames;
  }

(* Generic skeleton to process a signature. It returns a modified environment. *)

let definition nom = 
  let rec aux accu = function
    | 1 ->
        (fun x -> `Form (`Predicate (nom, ((`Unknown, `Variable x)::accu))))
    | n when n > 1 ->
        (fun x -> `Rel (aux ((`Unknown, `Variable x)::accu) (n - 1)))
    | _ -> failwith "definition: negative argument"
  in aux []

let deffonct nom = 
  fun x -> `Rel (fun y -> `Form (`Conjunction [
    `Equality ((`Unknown, `Variable y), (`Unknown, `Function (nom, [`Unknown, `Variable x])));
    `Predicate ("def__" ^ nom, [`Unknown, `Variable x]);
  ]))

let unsafe_get_relation env nom = match env.env_rels nom with
  | None -> failwith ("unsafe_get_relation: relation not found: " ^ nom)
  | Some r -> r

let unsafe_get_sig env nom x = match unsafe_get_relation env nom x with
  | `Form f -> f
  | _ -> failwith ("unsafe_get_sig: not a signature: " ^ nom)

type sigtype = 
  | Sig_abstract
  | Sig_one


let sig_generic 
    (nom : ident) (* The name of the signature. *)
    (sigtype: sigtype option) (* Is the signature abstract ? one ? *)
    (rels0 : (ident * ident list) list) (* The relations
                                           defined inside the signature,
                                           with the types of their
                                           arguments (the first excluded =
                                           signature being defined) *)
    (foncts : ident list) (* The functional relations defined inside the
                             signature *)
    (env : 'token env) (* The input environment, to which the signature will be
                   appended to construct the output environment *)
    : 'token env =

  prerr_endline ("General signature processing for: " ^ nom);
  let rels = List.map (fun (i, tl) -> (i, nom::tl)) rels0 in
  let def x =
    if sigtype = Some Sig_one
    then `Equality ((`Unknown, `Variable x), (`Unknown, `Constant nom))
    else `Predicate (nom, [`Unknown, `Variable x])
  in
  let def2 x = `Form (def x) in

  { env with

    env_noms_rels = begin
      prerr_endline ("sig " ^ nom ^ ": BEGIN processing relation names");
      let rec aux accu = function
        | [] -> 
            prerr_endline ("sig " ^ nom ^ ": END processing relation names");
            accu
        | (i, _) :: q when List.mem i accu ->
            prerr_endline ("relation " ^ i ^ ": already existing");
            aux accu q
        | (i, _) :: q ->
            prerr_endline ("adding relation: " ^ i);
            aux (i::accu) q
      in     
      aux env.env_noms_rels rels
    end;

    env_rels = begin
      prerr_endline ("sig " ^ nom ^ " : processing relations");
      let rels1 = (fun n -> if n = nom then
        Some def2
      else env.env_rels n) in
      let rels2 = fun n -> begin 
(*        prerr_endline ("Searching relation " ^ n ^ " in the definition of the signature " ^ nom); *)
        match rels1 n with
          | Some _ as res -> res
          | None when List.mem n foncts ->
              Some (
                prerr_endline ("function found: " ^ n);
                deffonct n
              )
          | None -> begin try
              let k = List.length (List.assoc n rels) in
              Some (
                prerr_endline ("non-function found, arity " ^ string_of_int k ^ " : " ^ n);
                definition n k
              )
            with Not_found -> None
            end
      end
      in
      rels2;
    end;

    env_types = begin
      prerr_endline ("sig " ^ nom ^ " : processing relation types");
      let types1 = (fun n -> if n = nom then [[nom]] else env.env_types n) in
      let types2 = (fun n -> let l = types1 n in try
        (List.assoc n rels)::l
      with Not_found -> l)
      in
      types2
    end;

    env_abstract = (if sigtype = Some Sig_abstract then (nom::env.env_abstract) else env.env_abstract);
    
    env_contraintes = begin
      prerr_endline ("sig " ^ nom ^ ": BEGIN processing functionality constraints");
      let rec aux accu = function
        | [] -> 
            prerr_endline ("sig " ^ nom ^ ": END processing functionality constraints");
            accu
        | n::q ->
            prerr_endline ("functionality constraint for: " ^ n);
            aux ((relation_fonction def n)::accu) q
      in aux env.env_contraintes foncts
    end;

  }



(* [abstract] sig Nom *)

let sig_master 
    (nom : ident) (* Signature name *)
    (sigtype : sigtype option) (* Is the signature abstract ? one ? *)
    (rels0, foncts : (ident * ident list) list * (* Relations *)
      ident list) (* Functional relations *)
    (env : 'token env) (* input environment *)
    : 'token env =

  prerr_endline ("Processing master signature: " ^ nom);
  let env2 = (sig_generic nom sigtype rels0 foncts env) in {env2 with
    env_peres =
      nom::env2.env_peres;
  }


(* sig Nom in Parent1 + Parent2 & Parent3 ... *)

let sig_in
    (nom : ident) (* Signature name *)
    (sigtype: sigtype option) (* Is the signature abstract ? one ? *)
    (parent : relation) (* The in clause, result of processing
                                       the signature expression *)
    (rels0, foncts : (ident * ident list) list * (* Relations *)
      ident list) (* Functional relations *)
    (env : 'token env) (* input environment *)
    : 'token env =

  prerr_endline ("Processing subset signature: " ^ nom);
  let env2 = (sig_generic nom sigtype rels0 foncts env) in {
    env2 with
    env_contraintes =
      begin match env2.env_rels nom with
        | None -> failwith "sig_in: impossible"
        | Some r -> (inclusion r parent)::env2.env_contraintes
      end;
  }


(* [abstract] sig Nom extends Parent *)

let sig_extends
    (nom : ident) (* Signature name *)
    (parent : ident) (* Extended signature *)
    (sigtype : sigtype option) (* Is the signature abstract ? *)
    (rels0, foncts : (ident * ident list) list * (* Relations *)
      ident list) (* Functional relations *)
    (env : 'token env) (* Input environment *)
    : 'token env =

  prerr_endline ("Processing extension signature: " ^ nom);  
  match env.env_rels parent with 
    | None -> failwith ("sig_extends: relation " ^ parent ^ " not found")
    | Some r ->
        let env2 = (sig_in nom sigtype r (rels0, foncts) env) in {
          env2 with
            env_fils = (fun n -> if n = parent then (nom::(env2.env_fils n)) else
              env2.env_fils n);
        }



(* Adding a predicate (Alloy "pred") into the env_predicats field of the
   environment. *)

let predicat
    (nom : ident) (* name *)
    (vars : ident list) (* names of the arguments *)
    (corps : 'token env -> fol_formula) (* body of the formula, using the names in
                                    vars *)
    (env : 'token env) (* source environment, into which the predicate will be
                   added, and also into which the values of the arguments
                   will be stored whenever the predicate is applied *)
    : 'token env =

  prerr_endline ("Processing predicate: " ^ nom);
  { env with
    env_predicats = (
      (nom, fun l ->
        let rec build_formula accu = function
          | ([], []) -> corps accu
          | v::w, r::s ->
              let accu2 = { accu with
                env_rels = (begin fun x -> 
(*                  prerr_endline ("Searching relation " ^ x ^ " as argument of the predicate " ^ nom); *)
                  if x = v then 
                  Some r
                 else accu.env_rels x end)
              } in build_formula accu2 (w, s)
          | _ -> invalid_arg ("pred " ^ nom ^ ": wrong number of arguments")
              in build_formula env (vars, l)
      )::env.env_predicats
    )
  }



(* Samely, adding a function (Alloy "fun") into the env_fonctions field of
   the environment *)

let fonction
    (nom : ident) (* name *)
    (vars : ident list) (* names of the arguments *)
    (corps : 'token env -> relation) (* body of the formula, using the names in
                                    vars *)
    (env : 'token env) (* source environment, into which the predicate will be
                   added, and also into which the values of the arguments
                   will be stored whenever the predicate is applied *)
    : 'token env =

  prerr_endline ("Processing function: " ^ nom);
  { env with
    env_fonctions = (
      (nom, fun l ->
        let rec build_relation accu = function
          | ([], []) -> corps accu
          | v::w, r::s ->
              let accu2 = { accu with
                env_rels = (begin fun x ->                  
(*                  prerr_endline ("Searching relation " ^ x ^ " as argument of the function " ^ nom); *)
                  if x = v then
                  Some r
                else accu.env_rels x end)
              } in build_relation accu2 (w, s)
          | _ -> invalid_arg ("fun " ^ nom ^ ": wrong number of arguments")
              in build_relation env (vars, l)
      )::env.env_fonctions
    )
  }


(* Adding an axiom (Alloy "fact") into the env_contraintes field of the
   environment. *)

let axiome
    (nom : ident)
    (corps : 'token env -> fol_formula)
    (env : 'token env)
    : 'token env =

  prerr_endline ("Processing axiom: " ^ nom);
  { env with
    env_contraintes = (corps env)::env.env_contraintes }


(* Samely, adding a theorem (Alloy "assert") into the env_theorem field *)

let theoreme 
    (nom : ident)
    (corps : 'token env -> fol_formula)
    (env : 'token env)
    : 'token env =

  prerr_endline ("Processing theorem: " ^ nom);
  { env with
    env_theoremes = (nom, corps env)::env.env_theoremes }



(* Finalization : build typing constraints related to signatures. *)

type resultat = {
  res_axiomes : fol_formula list; (* axioms. *)
  res_theoremes : (ident * fol_formula) list; (* theorems *)
}

let finaliser 
    (env : 'token env)
    : resultat = 

  prerr_endline "BEGIN Finalizing";
  let envf = {
    res_axiomes = (

      prerr_endline "Processing abstract signatures";
      let rec aux_abstract accu = function
        | [] -> accu
        | (a::l) -> 
            prerr_endline ("Processing abstract sig: " ^ a);
            let f = inclusion (unsafe_get_relation env a) (fun x -> `Form (
              `Disjunction (`False::(List.map (fun s -> unsafe_get_sig env s x)
                (env.env_fils a))))) in
            aux_abstract (f::accu) l
      in let ax1 = aux_abstract env.env_contraintes env.env_abstract
      in

      prerr_endline "Processing disjoint signatures";
      let rec aux_fils_disjoints accu = function
        | [] -> accu
        | a::q ->
            prerr_endline ("Processing disjoint subsigs for: " ^ a);
            let f = env.env_fils a in
            let accu2 = 
              aux_fils_disjoints (disjoints accu (List.map (unsafe_get_relation env) f)) f
            in aux_fils_disjoints accu2 q
      in let ax19 = aux_fils_disjoints ax1 env.env_peres
      in
      let ax2 = disjoints ax19 (List.map (unsafe_get_relation env) env.env_peres) in

      prerr_endline "Processing relation typesystem";
      let rec aux_typage accu = function
        | [] -> accu
        | a::q ->
            prerr_endline ("Processing type of relation: " ^ a);
            let l = List.map (List.map (unsafe_get_relation env)) (env.env_types a) in
            let l2 = List.map (fun l -> List.fold_left fleche (List.hd l) (List.tl l)) l in
            let r = List.fold_left plus (List.hd l2) (List.tl l2) in
            let f = inclusion (unsafe_get_relation env a) r in
            aux_typage (f::accu) q
      in let ax3 = aux_typage ax2 env.env_noms_rels
      in
      
      prerr_endline "Processing disjunction of master sigs";

      (
        inclusion univ
           (fun x -> `Form (`Disjunction (`False::(List.map (fun s ->
            unsafe_get_sig env s x) env.env_peres))))
      ) :: ax3
      

    );

    res_theoremes = begin
      env.env_theoremes;
    end;
  } in begin
    prerr_endline "END Finalizing. Enjoy the formulae !";
    envf
  end
  




(* Callings *)

(* Calling a predicate *)

let appel_predicat
    (nom : ident)
    (rels : relation list)
    (env : 'token env)
    : fol_formula =

  prerr_endline ("Calling predicate: " ^ nom);
  List.assoc nom env.env_predicats rels



(* Samely, calling a function *)

let appel_fonction
    (nom : ident)
    (rels : relation list)
    (env : 'token env)
    : relation =

  prerr_endline ("Calling function: " ^ nom);
  List.assoc nom env.env_fonctions rels



(* Calling a relation *)

let appel_relation
    (nom : ident)
    (env : 'token env)
    : relation =

  prerr_endline ("Calling relation: " ^ nom);
  match env.env_rels nom with
    | None -> invalid_arg ("relation " ^ nom ^ ": not found")
    | Some r -> r


(* Forall and exists *)

let quantifier
    quant vars corps env =
  
  let noms = List.map fst vars in

  let nouvvars = List.map (fun n -> (n, nouveau_symbole ("quant_" ^ n))) noms in

  let vars_quant = List.map (fun (_, v) -> (v, `Unknown)) nouvvars in
  
  let precond = `Conjunction (List.map (fun (v, r) -> match r (List.assoc v nouvvars) with
    | `Form f -> f
    | _ -> invalid_arg "Fol.quantifier: cannot quantify over relations") vars) in
  
  let env2 = { env with
    env_rels = fun n -> begin
(*      prerr_endline ("Searching relation " ^ n ^ " in a quantification."); *)
      if List.mem n noms
      then Some (atome (List.assoc n nouvvars))
      else env.env_rels n
    end;
  } in

  quant vars_quant precond (corps env2)


let quant_forall
    (vars : (ident * relation) list)
    :  ('token env -> fol_formula)
    -> 'token env
    -> fol_formula =

  quantifier (fun vars_quant precond corps ->
    `Forall (vars_quant, implication precond corps)) vars


let quant_exists
    (vars : (ident * relation) list)
    : ('token env -> fol_formula)
    -> 'token env
    -> fol_formula =

  quantifier (fun vars_quant precond corps ->
    `Exists (vars_quant, `Conjunction [precond; corps])) vars



(* Processing a file. It's weird, because it could depend on the parser.
   Actually even because of the type system. But finally not.
   I have to make env dependent on the type of the tokens. 
 *)