open Symb

(* Alloy abstract syntax tree *)

type relation = [
| `Plus of (int option * (relation * relation))
| `Moins of (int option * (relation * relation))
| `Inter of (int option * (relation * relation))
| `Fleche of ((int * int) option * (relation * relation))
| `Jointure of ((int * int) option * (relation * relation))
| `Restriction_droite of (int option * (relation * relation))
| `Restriction_gauche of (int option * (relation * relation))
| `Tilde of relation
| `Fonction of (ident * relation list)
| `Variable of ident
| `Univ
| `Iden
| `None
]

let rec string_of_relation = function
  | `Plus (_, (a, b)) -> "(" ^ string_of_relation a ^ " + " ^ string_of_relation b ^ ")"
  | `Moins (_, (a, b)) -> "(" ^ string_of_relation a ^ " - " ^ string_of_relation b ^ ")"
  | `Inter (_, (a, b)) -> "(" ^ string_of_relation a ^ " & " ^ string_of_relation b ^ ")"
  | `Fleche (_, (a, b)) -> "(" ^ string_of_relation a ^ "->" ^ string_of_relation b ^ ")"
  | `Jointure (_, (a, b)) -> "(" ^ string_of_relation a ^ "." ^ string_of_relation b ^ ")"
  | `Restriction_droite (_, (a, b)) -> "(" ^ string_of_relation a ^ " :> " ^ string_of_relation b ^ ")"
  | `Restriction_gauche (_, (a, b)) -> "(" ^ string_of_relation a ^ " <: " ^ string_of_relation b ^ ")"
  | `Tilde a -> "~" ^ string_of_relation a
  | `Fonction (f, l) -> f ^ "<" ^ String.concat "," (List.map string_of_relation l) ^ ">"
  | `Variable i -> i
  | `Univ -> "univ"
  | `Iden -> "iden"
  | `None -> "none"

type alloy_formule_atom = [
| `Predicate of (ident * (relation list))
| `Some of relation
| `No of relation
| `Lone of relation
| `One of relation
| `In of (relation * relation)
| `Equality of (relation * relation) 
]

type alloy_formule = [
| alloy_formule_atom
| `Negation of alloy_formule
| `Conjunction of alloy_formule list
| `Disjunction of alloy_formule list
| `Exists of ((ident * relation) list * alloy_formule)
| `Forall of ((ident * relation) list * alloy_formule)
]

let implication f g = `Disjunction [g; `Negation f]

let equivalence f g = `Conjunction [implication f g; implication g f]

type 'a alloy_corps = {
  corps_nom : ident;
  corps_param : (ident * relation (* ident list *) ) list;
  corps_corps : 'a;
}

type alloy_sigtype = [
| `Extends of ident
| `In of relation
]

type alloy_sigmult = [
| `One
| `Abstract
]

type alloy_sigrel = {
  sigrel_name : ident;
  sigrel_function : bool;
  sigrel_type : ident list;
}

type alloy_sig = {
  sig_name : ident;
  sig_mult : alloy_sigmult option;
  sig_type : alloy_sigtype option;
  sig_rels : alloy_sigrel list;
}

type 'token env = {
  env_pred : alloy_formule alloy_corps list;
  env_fun : relation alloy_corps list;
  env_fact : (ident * alloy_formule) list;
  env_sig : alloy_sig list;
  env_assert : (ident * alloy_formule) list;
  env_token: Lexing.lexbuf -> 'token; (* The lexer *)
  env_modul: (Lexing.lexbuf -> 'token) -> Lexing.lexbuf -> 'token env -> 'token env; (* The parser *)
  env_modulenames : ident list;
  env_searchpath : ident list;
}

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

  {
    env_pred = [];
    env_fun = [];
    env_fact = [];
    env_sig = [];
    env_assert = [];
    env_token = (fun _ -> failwith "Alloyast.env_token : lexer is undefined");
    env_modul = (fun _ _ x -> prerr_endline "Alloyast.env_modul : warning : parser is undefined"; x);
    env_modulenames = [];
    env_searchpath = [];
  }

let predicat f e = { e with
  env_pred = f::e.env_pred;
}

let fonction f e = { e with
  env_fun = f::e.env_fun;
}

let signature f e = { e with
  env_sig = f::e.env_sig;
}
  
let fait f e = { e with
  env_fact = f::e.env_fact;
}

let assertion f e = { e with
  env_assert = f::e.env_assert;
}

exception Corps_introuvable of ident

let rec chercher_corps nom = function
  | [] -> raise (Corps_introuvable nom)
  | a::_ when a.corps_nom = nom -> a
  | _::q -> chercher_corps nom q     

let process_channel
    (ch : in_channel)
    (env : 'token env)
    : 'token env =

  let lexbuf = Lexing.from_channel ch in
  env.env_modul env.env_token lexbuf env


(* Process a file. Please note : if the file does not exist, then returns None.
   This, in order to allow the calling function to retry with another file name
   when importing a module. *)

let process_file
    (name : ident)
    (env : 'token env)
    : 'token env option =

  if Sys.file_exists name
  then
    let ch =
      prerr_endline ("File found: " ^ name);
      open_in name in
    let env2 = process_channel ch env
    in Some (
      prerr_endline ("End of file: " ^ name);
      close_in ch;
      env2
    )
  else begin
    prerr_endline ("File not found: " ^ name);
    None
  end


(* Then, process a module. Given the search pathnames, and the module path,
   it successively tries one of them with the path of the given module. If
   it does not work, then try with the module path from which the first path item
   has been removed.

   Note that if a module of the given name has already been processed, then
   reloading another module with the same name will have no effect. *)

let process_module
    (modulepath, modulename : ident list * ident)
    (env: 'token env)
    : 'token env = 

  let searchpath = env.env_searchpath in
  prerr_endline ("Processing module: " ^ modulename);
  if List.mem modulename env.env_modulenames
  then begin
    prerr_endline ("Module with name: " ^ modulename ^ " already loaded. Skipping.");
    env
  end else begin
    prerr_endline ("No module loaded yet with name: " ^ modulename ^ ". Searching for corresponding file.");
    let rec build_fullnames = function
      | [], l -> l
      | a::q, [] -> build_fullnames (q, [a])
      | a::q, (b::_ as l) -> build_fullnames (q, (Filename.concat a b)::l)
    in let fullnames = build_fullnames (List.rev modulepath, [modulename ^ ".als"])
    in let sp = ((Sys.getcwd ())::searchpath)

    in let rec aux_fullnames = function
      | [] -> invalid_arg ("Fol.process_module: module " ^ modulename ^ " not found")
      | a::q -> let rec aux_searchpaths = function
          | [] -> 
              prerr_endline "No matching files found. Downgrading module path.";
              aux_fullnames q
          | p::r -> let f = Filename.concat p a in begin match
              prerr_endline ("Trying file: " ^ f);
              process_file (Filename.concat p a) env
            with
              | None -> aux_searchpaths r
              | Some e -> e
            end
        in aux_searchpaths sp
    in aux_fullnames fullnames
  end


let rec subst_naive_rel x r = function
  | `Plus (k, (a, b)) -> `Plus (k, (subst_naive_rel x r a, subst_naive_rel x r b))
  | `Moins (k, (a, b)) -> `Moins (k, (subst_naive_rel x r a, subst_naive_rel x r b))
  | `Inter (k, (a, b)) -> `Inter (k, (subst_naive_rel x r a, subst_naive_rel x r b))
  | `Fleche (k, (a, b)) -> `Fleche (k, (subst_naive_rel x r a, subst_naive_rel x r b))
  | `Jointure (k, (a, b)) -> `Jointure (k, (subst_naive_rel x r a, subst_naive_rel x r b))
  | `Restriction_droite (k, (a, b)) -> `Restriction_droite (k, (subst_naive_rel x r a, subst_naive_rel x r b))
  | `Restriction_gauche (k, (a, b)) -> `Restriction_gauche (k, (subst_naive_rel x r a, subst_naive_rel x r b))
  | `Tilde s -> `Tilde (subst_naive_rel x r s)
  | `Fonction (f, l) -> `Fonction (f, List.map (subst_naive_rel x r) l)
  | `Variable v when v = x -> r
  | s -> s

and subst_naive_form x r = function
  | `Negation f -> `Negation (subst_naive_form x r f)
  | `Conjunction l -> `Conjunction (List.map (subst_naive_form x r) l)
  | `Disjunction l -> `Disjunction (List.map (subst_naive_form x r) l)
  | `Exists (l, f) when not (List.mem_assoc x l) ->
      `Exists (l, subst_naive_form x r f)
  | `Forall (l, f) when not (List.mem_assoc x l) ->
      `Forall (l, subst_naive_form x r f)
  | `Predicate (f, l) -> `Predicate (f, List.map (subst_naive_rel x r) l)
  | `Some s -> `Some (subst_naive_rel x r s)
  | `No s -> `No (subst_naive_rel x r s)
  | `Lone s -> `Lone (subst_naive_rel x r s)
  | `One s -> `One (subst_naive_rel x r s)
  | `In (a, b) -> `In (subst_naive_rel x r a, subst_naive_rel x r b)
  | `Equality (a, b) -> `Equality (subst_naive_rel x r a, subst_naive_rel x r b)
  | f -> f

let affecte subst l s =
  let l2 = List.rev_map (fun (x, r) -> ((x, nouveau_symbole x), r)) l in
  let s2 = List.fold_left
    (fun a ((x, y), _) -> subst x (`Variable y) a)
    s l2
  in
  List.fold_left (fun a ((_, y), r) -> subst y r a) s2 l2

let affecte_rel l = affecte subst_naive_rel l
let affecte_form l = affecte subst_naive_form l

type arite = {
  arite_max : int;
  arite_rel : int;
  arite_tag : relation;
}

let rec max_arite_couple preds foncts rho f a b = 
  let aa = max_arite_rel preds foncts rho a and ab = max_arite_rel preds foncts rho b in
  if aa.arite_rel <> ab.arite_rel
  then failwith ("arites differentes : " ^ string_of_relation aa.arite_tag ^ " : " ^ string_of_int aa.arite_rel ^  ", " ^ string_of_relation ab.arite_tag ^ " : " ^ string_of_int ab.arite_rel)
  else {
    arite_max = max aa.arite_rel (max aa.arite_max ab.arite_max);
    arite_rel = aa.arite_rel;
    arite_tag = f aa ab;
  }
and max_arite_rest preds foncts rho f a cond =
  let c = max_arite_rel preds foncts rho cond in
  if c.arite_rel <> 1
  then failwith "arite: restriction non unaire"
  else let aa = max_arite_rel preds foncts rho a
  in {
    arite_max = max aa.arite_rel (max aa.arite_max c.arite_max);
    arite_rel = aa.arite_rel;
    arite_tag = f aa c;
  }
and max_arite_rel preds foncts rho = function
  | `Plus (_, (a, b)) ->
      max_arite_couple preds foncts rho (fun aa ab -> `Plus (Some aa.arite_rel, (aa.arite_tag, ab.arite_tag))) a b
  | `Moins (_, (a, b)) ->
      max_arite_couple preds foncts rho (fun aa ab -> `Moins (Some aa.arite_rel, (aa.arite_tag, ab.arite_tag))) a b
  | `Inter (_, (a, b)) ->
      max_arite_couple preds foncts rho (fun aa ab -> `Inter (Some aa.arite_rel, (aa.arite_tag, ab.arite_tag))) a b
  | `Fleche (_, (a, b)) ->
      let aa = max_arite_rel preds foncts rho a and ab = max_arite_rel preds foncts rho b
      in let s = aa.arite_rel + ab.arite_rel
      in {
        arite_max = max s (max aa.arite_max ab.arite_max);
        arite_rel = s;
        arite_tag = `Fleche (Some (aa.arite_rel, ab.arite_rel), (aa.arite_tag, ab.arite_tag));
      }      
  | `Jointure (_, (a, b)) ->
      let aa = max_arite_rel preds foncts rho a and ab = max_arite_rel preds foncts rho b
      in let s = aa.arite_rel + ab.arite_rel - 2
      in if s < 1
      then failwith "jointure de deux unaires"
      else {
        arite_max = max s (max aa.arite_max ab.arite_max);
        arite_rel = s;
        arite_tag = `Jointure (Some (aa.arite_rel, ab.arite_rel), (aa.arite_tag, ab.arite_tag));
      }
  | `Restriction_droite (_, (a, cond)) ->
      max_arite_rest preds foncts rho (fun aa c -> `Restriction_droite (Some aa.arite_rel, (aa.arite_tag, c.arite_tag))) a cond
  | `Restriction_gauche (_, (cond, a)) ->
      max_arite_rest preds foncts rho (fun aa c -> `Restriction_gauche (Some aa.arite_rel, (aa.arite_tag, c.arite_tag))) a cond
  | `Tilde a ->
      let aa = max_arite_rel preds foncts rho a in
      if aa.arite_rel <> 2
      then failwith "tilde d'un non-binaire"
      else { aa with
        arite_tag = `Tilde (aa.arite_tag);
      }
  | `Iden -> {
      arite_max = 2;
      arite_rel = 2;
      arite_tag = `Iden
    }
  | `Fonction (f, l) ->
      (* l'ordre de passage des param doit etre respecte *)
      let al = List.map (max_arite_rel preds foncts rho) l in
      let g =  chercher_corps f foncts in
      let rho2 x =
        let rec aux = function
          | (y, _)::_, a::_ when y = x -> a
          | _::q, _::r -> aux (q, r)
          | _ -> rho x
        in aux (g.corps_param, al)
      in let s : arite = max_arite_rel preds foncts rho2 g.corps_corps
      in {
        arite_max = List.fold_left (fun (m : int) (a : arite) -> max m a.arite_max) s.arite_max al;
        arite_rel = s.arite_rel;
        arite_tag = `Fonction (f, List.map (fun a -> a.arite_tag) al);
      }
  | `Variable v -> rho v
  | `Univ | `None as f -> {
      arite_max = 1;
      arite_rel = 1;
      arite_tag = f;
    }

let rec max_arite_form preds foncts rho = function
  | `Negation f -> 
      let n, g = max_arite_form preds foncts rho f in
      n, `Negation g
  | `Conjunction l ->
      let m = List.rev_map (max_arite_form preds foncts rho) l in
      List.fold_left (fun p (n, _) -> max p n) 0 m, `Conjunction (List.rev_map snd m)
  | `Disjunction l ->
      let m = List.rev_map (max_arite_form preds foncts rho) l in
      List.fold_left (fun p (n, _) -> max p n) 0 m, `Disjunction (List.rev_map snd m)
  | `Exists (l, f) ->
      let al = List.rev_map (fun (x, r) -> (x, max_arite_rel preds foncts rho r)) l
      in let rho2 x =
        try List.assoc x al
        with _ -> rho x
      in let n, g = max_arite_form preds foncts rho2 f
      in 
      List.fold_left (fun m (_, a) -> max m (a.arite_max)) n al, `Exists (List.rev_map (fun (x, a) -> (x, a.arite_tag)) al, g)
  | `Forall (l, f) ->
      let al = List.rev_map (fun (x, r) -> (x, max_arite_rel preds foncts rho r)) l
      in let rho2 x =
        try
          { (List.assoc x al) with
            arite_tag = `Variable x; }
        with _ -> rho x
      in let n, g = max_arite_form preds foncts rho2 f
      in 
      List.fold_left (fun m (_, a) -> max m (a.arite_max)) n al, `Forall (List.rev_map (fun (x, a) -> (x, a.arite_tag)) al, g)
  | `Predicate (f, l) ->
      (* l'ordre de passage des param doit etre respecte *)
      let al = List.map (max_arite_rel preds foncts rho) l in
      let g =  chercher_corps f preds in
      let rho2 x =
        let rec aux = function
          | (y, _)::_, a::_ when y = x -> a
          | _::q, _::r -> aux (q, r)
          | _ -> rho x
        in aux (g.corps_param, al)
      in let s, _ = max_arite_form preds foncts rho2 g.corps_corps
      in
      List.fold_left (fun m a -> max m a.arite_max) s al, `Predicate (f, List.map (fun a -> a.arite_tag) al)
  | `Some r ->
      let a = max_arite_rel preds foncts rho r in
      a.arite_max, `Some a.arite_tag
  | `No r ->
      let a = max_arite_rel preds foncts rho r in
      a.arite_max, `No a.arite_tag
  | `Lone r ->
      let a = max_arite_rel preds foncts rho r in
      a.arite_max, `Lone a.arite_tag
  | `One r ->
      let a = max_arite_rel preds foncts rho r in
      a.arite_max, `One a.arite_tag
  | `In (a, b) ->
      let aa = max_arite_rel preds foncts rho a
      and ab = max_arite_rel preds foncts rho b
      in max aa.arite_max ab.arite_max, `In (aa.arite_tag, ab.arite_tag)
  | `Equality (a, b) ->
      let aa = max_arite_rel preds foncts rho a
      and ab = max_arite_rel preds foncts rho b
      in max aa.arite_max ab.arite_max, `Equality (aa.arite_tag, ab.arite_tag)


let max_arite_sig rho s =
  List.fold_left
    (fun accu b -> 
      let n = List.length b.sigrel_type + 1 in
      fun x -> if x = b.sigrel_name then {
        arite_max = n;
        arite_rel = n;
        arite_tag = `Variable x;
      } else accu x)
    (fun b -> if b = s.sig_name then {
      arite_max = 1;
      arite_rel = 1;
      arite_tag = `Variable b;
    } else rho b)
    s.sig_rels