(*
 *  ENIAM_LCGparser, a parser for Logical Categorial Grammar formalism
 *  Copyright (C) 2016 Wojciech Jaworski <wjaworski atSPAMfree mimuw dot edu dot pl>
 *  Copyright (C) 2016 Institute of Computer Science Polish Academy of Sciences
 *
 *  This library is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This library is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *)

open ENIAM_LCGtypes

let rec internal_substitute var_name t = function
  | Atom x -> Atom x
  | AVar x -> if x = var_name then t else AVar x
  | With l -> With (Xlist.map l (internal_substitute var_name t))
  | Zero -> Zero
  | Top -> Top

let rec substitute var_name t = function
  | Tensor l -> Tensor (Xlist.map l (internal_substitute var_name t))
  | Plus l -> Plus (Xlist.map l (substitute var_name t))
  | Imp(s,d,t2) -> Imp(substitute var_name t s,d,substitute var_name t t2)
  | One -> One
  | ImpSet(s,l) -> ImpSet(substitute var_name t s, Xlist.map l (fun (d,s) -> d, substitute var_name t s))
  | WithVar(v,g,e,s) -> if v = var_name then WithVar(v,g,e,s) else WithVar(v,internal_substitute var_name t g,e,substitute var_name t s)
  | Star s -> Star (substitute var_name t s)
  | Bracket(lf,rf,s) -> Bracket(lf,rf,substitute var_name t s)
  | BracketSet d -> BracketSet d
  | Maybe s -> Maybe (substitute var_name t s)

let rec substitute_schema var_name t = function
  | Tensor l -> Tensor l
  | Plus l -> Plus (Xlist.map l (substitute_schema var_name t))
  | Imp(s,d,t2) -> Imp(substitute_schema var_name t s,d,substitute_schema var_name t t2)
  | One -> One
  | ImpSet(s,l) -> ImpSet(substitute_schema var_name t s, List.flatten (Xlist.map l (function
        Both,Tensor[AVar var_name] -> t
      | d,s -> [d, substitute_schema var_name t s])))
  | WithVar(v,g,e,s) -> WithVar(v,g,e,substitute_schema var_name t s)
  | Star s -> Star (substitute_schema var_name t s)
  | Bracket(lf,rf,s) -> Bracket(lf,rf,substitute_schema var_name t s)
  | BracketSet d -> BracketSet d
  | Maybe s -> Maybe (substitute_schema var_name t s)

let rec internal_count_avar var_name = function
    Atom _ -> 0
  | AVar x -> if x = var_name then 1 else 0
  | With l -> Xlist.fold l 0 (fun b t -> internal_count_avar var_name t + b)
  | Zero -> 0
  | Top -> 0

let rec count_avar var_name = function
  | Tensor l -> Xlist.fold l 0 (fun b t -> internal_count_avar var_name t + b)
  | Plus l -> Xlist.fold l 0 (fun b t -> count_avar var_name t + b)
  | Imp(s,d,t2) -> count_avar var_name s + count_avar var_name t2
  | One -> 0
  | ImpSet(s,l) -> count_avar var_name s + Xlist.fold l 0 (fun b (_,t) -> count_avar var_name t + b)
  | WithVar(v,g,e,s) -> if v = var_name then 0 else count_avar var_name s +  internal_count_avar var_name g
  | Star t -> count_avar var_name t
  | Bracket(lf,rf,s) -> count_avar var_name s
  | BracketSet _ -> 0
  | Maybe t -> count_avar var_name t

let rec substitute_substvar v g = function
    Var v as t -> t
  | Tuple l -> Tuple(Xlist.map l (substitute_substvar v g))
  (*   | LetIn(l,s,t) -> LetIn(l,substitute_substvar v g s,substitute_substvar v g t) *)
  | Variant(e,l) -> Variant(e,Xlist.map l (fun (i,t) -> i,substitute_substvar v g t))
  | VariantVar(v2,t) -> if v2 = v then VariantVar(v2,t) else VariantVar(v2,substitute_substvar v g t)
  | SubstVar v2 -> if v2 = v then g else SubstVar v2
  | Case(t,l) -> Case(substitute_substvar v g t,Xlist.map l (fun (x,t) -> x,substitute_substvar v g t))
  | App(s,t) -> App(substitute_substvar v g s,substitute_substvar v g t)
  | Lambda(v2,t) -> Lambda(v2,substitute_substvar v g t)
  | LambdaSet(l,t) -> LambdaSet(l,substitute_substvar v g t)
  | Dot -> Dot
  | Val s -> Val s
  | SetAttr(e,s,t) -> SetAttr(e,substitute_substvar v g s,substitute_substvar v g t)
  | Fix(s,t) -> Fix(substitute_substvar v g s,substitute_substvar v g t)
  | Node t -> Node{t with attrs=Xlist.map t.attrs (fun (e,t) -> e, substitute_substvar v g t);
                          symbol=substitute_substvar v g t.symbol;
                          arg_symbol=substitute_substvar v g t.arg_symbol;
                          args=substitute_substvar v g t.args}
  | Cut t -> Cut(substitute_substvar v g t)
  | t -> failwith ("substitute_substvar: " ^ ENIAM_LCGstringOf.linear_term 0 t)


let empty_node = {
  orth=""; lemma=""; pos=""; weight=0.; id=0; symbol=Dot; arg_symbol=Dot; arg_dir=""; attrs=[]; args=Dot;}

let variable_num_ref = ref 0

let reset_variable_numbers () =
  variable_num_ref := 0

let add_variable_numbers () =
  incr variable_num_ref

let variable_name_ref = ref []

let reset_variable_names () =
  variable_name_ref := []

let rec add_variable_name = function
    [] -> ["a"]
  | "z" :: l -> "a" :: add_variable_name l
  | s :: l -> String.make 1 (Char.chr (Char.code (String.get s 0) + 1)) :: l

let get_variable_name () =
  variable_name_ref := add_variable_name (!variable_name_ref);
  String.concat "" (List.rev (!variable_name_ref)) ^ (string_of_int !variable_num_ref)

let make_arg_symbol l =
  Tuple(Xlist.map l (function
        Atom s -> Val s
      | AVar s -> Val s
      | Top -> Val "T"
      | _ -> failwith "make_arg_symbol"))

let string_of_direction = function
    Forward -> "forward"
  | Backward -> "backward"
  | Both -> "both"

let rec make_term_arg dir = function
    Tensor l ->
    let v = get_variable_name () in
    v, Cut(SetAttr("ARG_DIR",Val (string_of_direction dir),
                   SetAttr("ARG_SYMBOL",make_arg_symbol l,Var v)))
  | Plus l -> let v = get_variable_name () in v, Case(Var v,Xlist.map l (make_term_arg dir))
  (* | Imp(s,d,t2) -> *)
  | One -> get_variable_name (), Dot
  | Maybe s ->
    let v,arg = make_term_arg dir s in
    let w = get_variable_name () in
    w, Fix(Var w,Lambda(v,arg))
  | _ -> failwith "make_term_arg"

let add_args node args =
  {node with args=Tuple(node.args :: args)}

let make_raised_arg_symbol = function
    Imp(Tensor l,_,_) -> make_arg_symbol l
  | _ -> failwith "arg_symbol"

let rec make_raised_term_imp inner_node outer_node arg_symbol = function
  | Imp(s,d,t2) ->
    let v = get_variable_name () in
    let arg_symbol = make_raised_arg_symbol t2 in
    Lambda(v,make_raised_term_imp (App(Var v,inner_node)) outer_node arg_symbol s)
  | ImpSet(s,[d,t2]) ->
    let v = get_variable_name () in
    let arg_symbol = make_raised_arg_symbol t2 in
    LambdaSet([v],make_raised_term_imp (App(Var v,inner_node)) outer_node arg_symbol s)
  | Tensor l ->
    if outer_node.lemma="" then inner_node else
    Node (add_args outer_node [Cut(SetAttr("ARG_SYMBOL",arg_symbol,inner_node))])
  | _ -> failwith "make_raised_term_imp"

let is_raised = function
    [_,Imp(_,_,_)] -> true
  | _ -> false

let rec make_term_imp node outer_node = function
  | Imp(s,d,t2) ->
    if is_raised [d,t2] then make_raised_term_imp (Node node) outer_node Dot (Imp(s,d,t2)) else
    let v,arg = make_term_arg d t2 in
    Lambda(v,make_term_imp (add_args node [arg]) outer_node s)
  | ImpSet(s,l) ->
    if is_raised l then make_raised_term_imp (Node node) outer_node Dot (ImpSet(s,l)) else
    let vars,args = List.split (Xlist.map l (fun (d,t) -> make_term_arg d t)) in
    LambdaSet(vars,make_term_imp (add_args node args) outer_node s)
  | Tensor l -> Node node
  | _ -> failwith "make_term_imp"

let rec make_term_withvar node outer_node = function
    WithVar(category,_,_,t) -> VariantVar(category,make_term_withvar node outer_node t)
  | Bracket(_,_,t) -> make_term_withvar node outer_node t
  | t -> make_term_imp node outer_node t

let make_term node = make_term_withvar node empty_node
let make_raised_term node outer_node = make_term_withvar node outer_node

let rec make_symbol = function
  | Tensor l ->  Tuple(Xlist.map l (function
        Atom s -> Val s
      | AVar s -> SubstVar s
      | _ -> failwith "make_symbol"))
  | Plus l -> failwith "make_symbol"
  | Imp(s,d,t2) -> make_symbol s
  | One -> failwith "make_symbol"
  | ImpSet(s,l) -> make_symbol s
  | WithVar(v,g,e,s) -> make_symbol s
  | Star t -> failwith "make_symbol"
  | Bracket(lf,rf,s) -> make_symbol s
  | BracketSet _ -> failwith "make_symbol"
  | Maybe t -> failwith "make_symbol"

let make_raised_symbol_arg = function
    [_,Imp(_,_,Tensor l)] ->
      Tuple(Xlist.map l (function
        Atom s -> Val s
      | AVar s -> SubstVar s
      | _ -> failwith "make_raised_symbol_arg"))
  | _ -> failwith "make_raised_symbol_arg"

let rec make_raised_symbol = function
  | Tensor l -> failwith "make_raised_symbol"
  | Plus l -> failwith "make_raised_symbol"
  | Imp(s,d,t2) -> if is_raised [d,t2] then make_raised_symbol_arg [d,t2] else make_raised_symbol s
  | One -> failwith "make_raised_symbol"
  | ImpSet(s,l) -> if is_raised l then make_raised_symbol_arg l else make_raised_symbol s
  | WithVar(v,g,e,s) -> make_raised_symbol s
  | Star t -> failwith "make_raised_symbol"
  | Bracket(lf,rf,s) -> make_raised_symbol s
  | BracketSet _ -> failwith "make_raised_symbol"
  | Maybe t -> failwith "make_raised_symbol"

let rec simplify = function
    ImpSet(s,[]),LambdaSet([],t) -> simplify (s,t)
  | ImpSet(s,[d,a]),LambdaSet([v],t) -> let s,t = simplify (s,t) in Imp(s,d,a),Lambda(v,t)
  | ImpSet(s,l),LambdaSet(vl,t) -> let s,t = simplify (s,t) in ImpSet(s,l),LambdaSet(vl,t)
  | WithVar(v,Atom g,e,s),VariantVar(_,t) -> simplify (substitute v (Atom g) s, substitute_substvar v (ENIAM_LCGrules.make_subst e (Atom g)) t)
  | WithVar(v,g,e,s),VariantVar(v2,t) ->
    if count_avar v s = 0 then
      simplify (s, substitute_substvar v (ENIAM_LCGrules.make_subst e g) t)
    else let s,t = simplify (s,t) in WithVar(v,g,e,s),VariantVar(v2,t)
  | Bracket(lf,rf,s),t -> let s,t = simplify (s,t) in Bracket(lf,rf,s),t
  | s,t -> s,t

let make_quant_restriction = function
    [] -> Zero
  | [s] -> Atom s
  | l -> With(Xlist.map l (fun s -> Atom s))
