(** The 20 amino-acids and their various physico-chemical classes. *)

module type PARAM =
  sig
    val gen_class_size : int  (** The maximum size of amino-acids classes that can be generated by [gen]. *)
  end

module Make (Param : PARAM) =
  struct
    include Logic.Default

    type t = string LSet.t

    let all_aa = LSet.of_list ["A";"C";"D";"E";"F";"G";"H";"I";"K";"L";"M";"N";"P";"Q";"R";"S";"T";"V";"W";"Y";"X"]

    let valid_aa s = List.mem s all_aa

    let class_of_string s =
      let res = ref (LSet.empty ()) in
      for i = String.length s - 1 downto 0 do
        res := LSet.add (String.sub s i 1) !res
      done;
      !res

    let acidic = class_of_string "DE"
    let basic = class_of_string "KR"
    let ion_polar = class_of_string "DEKRH"
    let non_ion_polar = class_of_string "DEKRHNQST"
    let ion_hbond_non_polar = class_of_string "CY"
    let ion = class_of_string "DEKRHCY"
    let hbond_aromatic = class_of_string "YW"
    let hbond_non_polar = class_of_string "CYW"
    let hbond = class_of_string "DEKRHNQSTCY"
    let sulfur = class_of_string "CM"
    let aromatic = class_of_string "YWF"
    let aliph_non_polar = class_of_string "AILPV"
    let non_polar = class_of_string "CYWMFAILPV"
    let aliph = class_of_string "AILPVG"

    let all_classes =
      List.map LSet.singleton all_aa @
      [acidic; basic; ion_polar; non_ion_polar; ion_hbond_non_polar; ion; hbond_aromatic; hbond_non_polar;
       hbond; sulfur; aromatic; aliph_non_polar; non_polar; aliph]

    open Token

    let parse = parser
      | [<'Term t>] ->
          if LSet.mem t all_aa
          then LSet.singleton t
          else raise (Stream.Error ("Syntax error: invalid amino-acid: " ^ t))
      | [<'LeftBra;
          'Term t ?? "Syntax error: set of amino-acids expected after '['";
          'RightBra ?? "Syntax error: missing ']' after: [" ^ t
         >] ->
          let cl = class_of_string t in
          if LSet.contains all_aa cl
          then cl
          else raise (Stream.Error ("Syntax error: invalid amino-acid in: [" ^ t ^ "]"))
      | [<'LeftAcc;
          'Term t ?? "Syntax error: set of amino-acids expected after '{'";
          'RightAcc ?? "Syntax error: missing '}' after: {" ^ t
         >] ->
          let cl = class_of_string t in
          if LSet.contains all_aa cl
          then LSet.subtract all_aa cl
          else raise (Stream.Error ("Syntax error: invalid amino-acid in: {" ^ t ^ "}"))

    let parse_compact = parser
      | [<'Term s>] ->
          let res = ref [] in
          let l = String.length s in
          for i = l - 1 downto 0 do
            res := Term (String.sub s i 1) :: !res
          done;
          !res

    let print = function
      | [] -> assert false
      | [aa] -> [Term aa]
      | cl ->
         if LSet.cardinal cl <= 10
         then [LeftBra; Term (String.concat "" cl); RightBra]
         else [LeftAcc; Term (String.concat "" (LSet.subtract all_aa cl)); RightAcc]

    let top () = all_aa

    let entails f g = LSet.contains g f

    let conj f g = LSet.inter f g

    let add f g = conj f g

    let sub f g =
       let h = LSet.subtract f g in
       if LSet.is_empty h then raise Not_found else h

    let features f = []

    let gen f g hs =
      try
        let x = List.find (fun c -> LSet.contains c (LSet.union f g)) all_classes in
        if LSet.cardinal x > Param.gen_class_size
         or List.exists (fun h -> entails h x) hs
        then LSet.empty ()
        else LSet.singleton x
      with Not_found -> LSet.empty ()

  end
