(*
 *  ENIAMsubsyntax: MWE, abbreviation and sentence detecion for Polish
 *  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 Xstd
open ENIAMtokenizerTypes

let to_string (paths,last) =
  String.concat "\n" (Xlist.map paths (fun t -> ENIAMtokens.string_of_tokens 0 (Token t)))
  ^ Printf.sprintf "\nlast=%d" last

let to_string_indexed (paths,last) =
  String.concat "\n" (Xlist.map paths (fun (i,j,t) ->
    Printf.sprintf "%2d %2d %s" i j (ENIAMtokens.string_of_tokens 0 (Token t))))
  ^ Printf.sprintf "\nlast=%d" last

let compare_token_record p r =
  let v = compare p.beg r.beg in
  if v <> 0 then v else
  let v = compare p.next r.next in
  if v <> 0 then v else
  compare p r

let sort (paths,last) =
  Xlist.sort paths compare_token_record, last

let rec uniq_rec rev = function
    [] -> List.rev rev
  | [p] -> List.rev (p :: rev)
  | p :: r :: l -> if p = r then uniq_rec rev (r :: l) else uniq_rec (p :: rev) (r :: l)

let uniq (paths,last) =
  uniq_rec [] paths, last

let rec translate_into_paths_rec paths = function
    Token t -> t :: paths
  | Seq l -> Xlist.fold l paths translate_into_paths_rec
  | Variant l -> Xlist.fold l paths translate_into_paths_rec

let translate_into_paths tokens =
  let paths = Xlist.fold tokens [] (fun paths token ->
    translate_into_paths_rec paths token) in
  let last = if paths = [] then 0 else (List.hd paths).next in
  let paths = sort (paths,last) in
  let paths = uniq paths in
  paths

(**********************************************************************************)

let excluded_interps = StringSet.of_list [
  "praet:sg:f:ter:perf";
  "praet:sg:f:ter:imperf.perf";
  "praet:sg:f:ter:imperf";
  "praet:pl:m2.m3.f.n1.n2.p2.p3:ter:perf";
  "praet:pl:m2.m3.f.n1.n2.p2.p3:ter:imperf.perf";
  "praet:pl:m2.m3.f.n1.n2.p2.p3:ter:imperf";
  "praet:pl:m1.p1:ter:perf";
  "praet:pl:m1.p1:ter:imperf.perf";
  "praet:pl:m1.p1:ter:imperf";
  "praet:sg:m1.m2.m3:ter:perf";
  "praet:sg:m1.m2.m3:ter:imperf.perf";
  "praet:sg:m1.m2.m3:ter:imperf";
  "praet:sg:n1.n2:ter:perf";
  "praet:sg:n1.n2:ter:imperf.perf";
  "praet:sg:n1.n2:ter:imperf";
  "cond:sg:m1.m2.m3:ter:perf";
  "cond:sg:m1.m2.m3:ter:imperf.perf";
  "cond:sg:m1.m2.m3:ter:imperf";
  "cond:sg:m1.m2.m3:sec:perf";
  "cond:sg:m1.m2.m3:sec:imperf.perf";
  "cond:sg:m1.m2.m3:sec:imperf";
  "cond:sg:m1.m2.m3:pri:perf";
  "cond:sg:m1.m2.m3:pri:imperf.perf";
  "cond:sg:m1.m2.m3:pri:imperf";
  "cond:sg:m1.m2.m3:perf";
  "cond:sg:m1.m2.m3:imperf.perf";
  "cond:sg:m1.m2.m3:imperf";
  "cond:sg:f:ter:perf";
  "cond:sg:f:ter:imperf.perf";
  "cond:sg:f:ter:imperf";
  "cond:sg:f:sec:perf";
  "cond:sg:f:sec:imperf.perf";
  "cond:sg:f:sec:imperf";
  "cond:sg:f:pri:perf";
  "cond:sg:f:pri:imperf.perf";
  "cond:sg:f:pri:imperf";
  "cond:sg:f:perf";
  "cond:sg:f:imperf.perf";
  "cond:sg:f:imperf";
  "cond:sg:n1.n2:ter:perf";
  "cond:sg:n1.n2:ter:imperf.perf";
  "cond:sg:n1.n2:ter:imperf";
  "cond:sg:n1.n2:sec:perf";
  "cond:sg:n1.n2:sec:imperf.perf";
  "cond:sg:n1.n2:sec:imperf";
  "cond:sg:n1.n2:pri:perf";
  "cond:sg:n1.n2:pri:imperf.perf";
  "cond:sg:n1.n2:pri:imperf";
  "cond:sg:n1.n2:perf";
  "cond:sg:n1.n2:imperf.perf";
  "cond:sg:n1.n2:imperf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:ter:perf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:ter:imperf.perf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:ter:imperf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:sec:perf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:sec:imperf.perf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:sec:imperf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:pri:perf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:pri:imperf.perf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:pri:imperf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:perf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:imperf.perf";
  "cond:pl:m2.m3.f.n1.n2.p2.p3:imperf";
  "cond:pl:m1.p1:ter:perf";
  "cond:pl:m1.p1:ter:imperf.perf";
  "cond:pl:m1.p1:ter:imperf";
  "cond:pl:m1.p1:sec:perf";
  "cond:pl:m1.p1:sec:imperf.perf";
  "cond:pl:m1.p1:sec:imperf";
  "cond:pl:m1.p1:pri:perf";
  "cond:pl:m1.p1:pri:imperf.perf";
  "cond:pl:m1.p1:pri:imperf";
  "cond:pl:m1.p1:perf";
  "cond:pl:m1.p1:imperf.perf";
  "cond:pl:m1.p1:imperf";
  "winien:sg:n1.n2:ter:imperf";
  "winien:sg:n1.n2:sec:imperf";
  "winien:sg:n1.n2:pri:imperf";
(*   "winien:sg:n1.n2:imperf"; *)
  "winien:sg:m1.m2.m3:ter:imperf";
  "winien:sg:m1.m2.m3:sec:imperf";
  "winien:sg:m1.m2.m3:pri:imperf";
(*   "winien:sg:m1.m2.m3:imperf"; *)
  "winien:sg:f:ter:imperf";
  "winien:sg:f:sec:imperf";
  "winien:sg:f:pri:imperf";
(*   "winien:sg:f:imperf"; *)
  "winien:pl:m2.m3.f.n1.n2.p2.p3:ter:imperf";
  "winien:pl:m2.m3.f.n1.n2.p2.p3:sec:imperf";
  "winien:pl:m2.m3.f.n1.n2.p2.p3:pri:imperf";
(*   "winien:pl:m2.m3.f.n1.n2.p2.p3:imperf"; *)
  "winien:pl:m1.p1:ter:imperf";
  "winien:pl:m1.p1:sec:imperf";
  "winien:pl:m1.p1:pri:imperf";
(*   "winien:pl:m1.p1:imperf"; *)
  ]

let transformed_interps = Xlist.fold [
  "praet:sg:f:perf:agl","praet:sg:f:perf";
  "praet:sg:f:imperf.perf:agl","praet:sg:f:imperf.perf";
  "praet:sg:f:imperf:agl","praet:sg:f:imperf";
  "praet:pl:m2.m3.f.n1.n2.p2.p3:perf:agl","praet:pl:m2.m3.f.n1.n2.p2.p3:perf";
  "praet:pl:m2.m3.f.n1.n2.p2.p3:imperf.perf:agl","praet:pl:m2.m3.f.n1.n2.p2.p3:imperf.perf";
  "praet:pl:m2.m3.f.n1.n2.p2.p3:imperf:agl","praet:pl:m2.m3.f.n1.n2.p2.p3:imperf";
  "praet:pl:m1.p1:perf:agl","praet:pl:m1.p1:perf";
  "praet:pl:m1.p1:imperf.perf:agl","praet:pl:m1.p1:imperf.perf";
  "praet:pl:m1.p1:imperf:agl","praet:pl:m1.p1:imperf";
  "praet:sg:m1.m2.m3:perf:agl","praet:sg:m1.m2.m3:perf";
  "praet:sg:m1.m2.m3:imperf.perf:agl","praet:sg:m1.m2.m3:imperf.perf";
  "praet:sg:m1.m2.m3:imperf:agl","praet:sg:m1.m2.m3:imperf";
  "praet:sg:n1.n2:perf:agl","praet:sg:n1.n2:perf";
  "praet:sg:n1.n2:imperf.perf:agl","praet:sg:n1.n2:imperf.perf";
  "praet:sg:n1.n2:imperf:agl","praet:sg:n1.n2:imperf";
  "praet:sg:f:perf:nagl","praet:sg:f:perf";
  "praet:sg:f:imperf.perf:nagl","praet:sg:f:imperf.perf";
  "praet:sg:f:imperf:nagl","praet:sg:f:imperf";
  "praet:pl:m2.m3.f.n1.n2.p2.p3:perf:nagl","praet:pl:m2.m3.f.n1.n2.p2.p3:perf";
  "praet:pl:m2.m3.f.n1.n2.p2.p3:imperf.perf:nagl","praet:pl:m2.m3.f.n1.n2.p2.p3:imperf.perf";
  "praet:pl:m2.m3.f.n1.n2.p2.p3:imperf:nagl","praet:pl:m2.m3.f.n1.n2.p2.p3:imperf";
  "praet:pl:m1.p1:perf:nagl","praet:pl:m1.p1:perf";
  "praet:pl:m1.p1:imperf.perf:nagl","praet:pl:m1.p1:imperf.perf";
  "praet:pl:m1.p1:imperf:nagl","praet:pl:m1.p1:imperf";
  "praet:sg:m1.m2.m3:perf:nagl","praet:sg:m1.m2.m3:perf";
  "praet:sg:m1.m2.m3:imperf.perf:nagl","praet:sg:m1.m2.m3:imperf.perf";
  "praet:sg:m1.m2.m3:imperf:nagl","praet:sg:m1.m2.m3:imperf";
  "praet:sg:n1.n2:perf:nagl","praet:sg:n1.n2:perf";
  "praet:sg:n1.n2:imperf.perf:nagl","praet:sg:n1.n2:imperf.perf";
  "praet:sg:n1.n2:imperf:nagl","praet:sg:n1.n2:imperf";
  ] StringMap.empty (fun map (k,v) -> StringMap.add map k v)

let merge_lemmata l =
  let map = Xlist.fold l StringMap.empty (fun map (lemma,interp,quantity,attrs) ->
    let interp = if interp = "num:comp" then "numc" else interp in
    if StringSet.mem excluded_interps interp then map else
    let interp = try StringMap.find transformed_interps interp with Not_found -> interp in
    let s = lemma ^ "#" ^ String.concat "|" attrs in
    StringMap.add_inc map s (lemma,quantity,[interp],attrs) (fun (_,q,l,_) -> lemma,q+quantity,interp :: l,attrs)) in
  let map = StringMap.map map (fun (lemma,quantity,interps,attrs) ->
    lemma,
    Xlist.fold interps StringMap.empty (fun map interp ->
      Xlist.fold (ENIAMtokens.parse_postags interp) map (fun map (pos,tags) ->
        StringMap.add_inc map pos [tags] (fun l -> tags :: l))),
    max 1 (quantity / Xlist.size interps),
    attrs) in
  StringMap.fold map [] (fun l _ (lemma,map,quantity,attrs) ->
    StringMap.fold map l (fun l cat interp ->
      (lemma,cat,interp,quantity,attrs) :: l))

let uppercase lemma cl ll =
  let n = String.length lemma in
  let nll = String.length ll in
  cl ^ String.sub lemma nll (n - nll)

let quant_mod quantity =
  log10 (float quantity)

let lemmatize_token = function
  | {token=AllSmall s} as t ->
       t :: (Xlist.map (merge_lemmata (ENIAMinflexion.get_interpretations s)) (fun (lemma,cat,interp,quantity,attrs) ->
         {t with token=Lemma(lemma,cat,interp); weight=t.weight+.(quant_mod quantity); attrs=attrs @ t.attrs}))
  | {token=SmallLetter s} as t ->
       t :: (Xlist.map (merge_lemmata (ENIAMinflexion.get_interpretations s)) (fun (lemma,cat,interp,quantity,attrs) ->
         {t with token=Lemma(lemma,cat,interp); weight=t.weight+.(quant_mod quantity); attrs=attrs @ t.attrs}))
  | {token=FirstCap(s,lower,cl,ll)} as t ->
       let l = Xlist.map (merge_lemmata (ENIAMinflexion.get_interpretations s)) (fun (lemma,cat,interp,quantity,attrs) ->
           {t with token=Lemma(lemma,cat,interp); weight=t.weight+.(quant_mod quantity); attrs="notvalidated proper" :: attrs @ t.attrs}) in
       let b = Xlist.fold l false (fun b t -> if Xlist.mem t.attrs "lemma not validated" || Xlist.mem t.attrs "token not found" then b else true) in
       if b then t :: l else
       t :: (Xlist.map (merge_lemmata (ENIAMinflexion.get_interpretations lower)) (fun (lemma,cat,interp,quantity,attrs) ->
           {t with token=Lemma(uppercase lemma cl ll,cat,interp); weight=t.weight+.(quant_mod quantity); attrs="notvalidated proper" :: "lemmatized as lowercase" :: attrs @ t.attrs}))
  | {token=AllCap(s,_,_)} as t ->
       t :: (Xlist.map (merge_lemmata (ENIAMinflexion.get_interpretations s)) (fun (lemma,cat,interp,quantity,attrs) ->
         {t with token=Lemma(lemma,cat,interp); weight=t.weight+.(quant_mod quantity); attrs="notvalidated proper" :: attrs @ t.attrs}))
  | {token=CapLetter(s,_)} as t ->
       t :: (Xlist.map (merge_lemmata (ENIAMinflexion.get_interpretations s)) (fun (lemma,cat,interp,quantity,attrs) ->
         {t with token=Lemma(lemma,cat,interp); weight=t.weight+.(quant_mod quantity); attrs="notvalidated proper" :: attrs @ t.attrs}))
  | {token=SomeCap s} as t ->
       t :: (Xlist.map (merge_lemmata (ENIAMinflexion.get_interpretations s)) (fun (lemma,cat,interp,quantity,attrs) ->
         {t with token=Lemma(lemma,cat,interp); weight=t.weight+.(quant_mod quantity); attrs="notvalidated proper" :: attrs @ t.attrs}))
  | t -> [t]

let rec lemmatize_rec rev = function
    [] -> List.rev rev
  | t :: l -> lemmatize_rec (lemmatize_token t @ rev) l

let lemmatize (paths,last) =
  List.rev (Xlist.fold (lemmatize_rec [] paths) [] (fun paths t ->
    match t.token with
      Lemma _ -> if (Xlist.mem t.attrs "lemma not validated" || Xlist.mem t.attrs "token not found") &&
        (Xlist.mem t.attrs "required validated lemmatization") then paths
        else t(*{t with attrs=List.remove "required validated lemmatization" t.attrs}*) :: paths
    | _ -> t :: paths)), last


