pqTree.ml

open! DataStruct 
open! Combi 

type 'elt t =
 | Leaf of 'elt
 | P of 'elt t list 
 | Q of 'elt t list 


let rec enumerate_permutation = 
  let open Enumeration in 
  function
  | Leaf elt -> Seq.return [elt]
  | P children -> 
     children
     |> List.map enumerate_permutation
     |> permutation_list_enumerator
     |> Seq.map List.concat 
  | Q children -> 
     children
     |> List.map enumerate_permutation
     |> list_enumerator >>= fun lists ->
     let list = List.concat lists in 
     List.to_seq [list; List.rev list]


let frontier tree =
  let rec go elements = function
    | Leaf e -> 
       e::elements 
    | P children | Q children ->
       List.fold_left go elements (List.rev children) 
  in 
  go [] tree 


let rec length = function 
  | Leaf _ -> 1
  | P children | Q children -> 
     List.fold_left (+) 0 (List.map length children)


let rec count_permutations = 
  let open Big_int_Z in 
  function
  | Leaf _ -> big_int_of_int 1
  | P children -> 
     children 
     |> List.map count_permutations 
     |> List.fold_left mult_big_int (Combinatorics.factorial (List.length children))
  | Q children -> 
     children 
     |> List.map count_permutations 
     |> List.fold_left mult_big_int (big_int_of_int 2)


let rec shuffle = function
  | Leaf elt -> Leaf elt
  | P children -> 
     children 
     |> List.map shuffle
     |> Combinatorics.random_permutation
     |> fun c -> P c 
  | Q children -> 
     children 
     |> List.map shuffle
     |> (if Random.bool () then List.rev else (fun x -> x))
     |> fun c -> Q c 


let sample_permutation pqtree =
  pqtree 
  |> shuffle
  |> frontier 
 

let sample_interval pqtree = 
  pqtree
  |> sample_permutation 
  |> Combinatorics.random_sublist



let rec is_well_formed = function 
    | Leaf _ -> true
    | P children ->
       List.length children > 1 
       && List.for_all is_well_formed children 
    | Q children -> 
       List.length children > 2 
       && List.for_all is_well_formed children 


module Skeleton = 
  struct 

    type skeleton = 
      | Hole 
      | HoledP of skeleton list 
      | HoledQ of skeleton list

    let to_pqtree = 
      let rec insert (accu, mini) skeleton = 
        let (child, next) = go mini skeleton in
        (child::accu, next)
      and go mini = function
        | Hole -> (Leaf mini, mini + 1)
        | HoledP skeletons -> 
           let (children,value) = 
             List.fold_left insert ([],mini) skeletons
           in 
           ( P (List.rev children), value) 
        | HoledQ skeletons -> 
           let (children,value) = 
             List.fold_left insert ([],mini) skeletons
           in 
           (Q (List.rev children), value) 
      in 
      fun skeleton -> fst (go 0 skeleton)

  end


module type EnumerateSig = 
sig
  type elt 
  include FeatCore.EnumSig.ENUM
  val enumeration : elt t enum 
  val compatible_orders : 'elt t -> 'elt list IFSeq.seq
end

module MakeEnumerate =
  functor (RandomM : FeatCore.RandomSig.S) ->
  struct 

    module PQTIFSeq = FeatCore.IFSeq.Make(Feat.Num)(RandomM)
    module PQTEnum = FeatCore.Enum.Make(PQTIFSeq)


    open Skeleton 

    type elt = int 

    let enum_skeletons =
      Fix.Memoize.Int.fix 
        (fun quintuple n ->
          let open PQTIFSeq in 
          let (>>=) l f = MoreList.flat_map f l in 
          let enum i = let (fst,_,_,_) = quintuple i in fst in 
          let enum_list i = let (_,snd,_,_) = quintuple i in snd in 
          let enum_list1 i = let (_,_,trd,_) = quintuple i in trd in 
          let enum_list2 i = let (_,_,_,frt) = quintuple i in frt in 
          let incr list_j =
            bigsum 
              begin 
                MoreList.range 1 (n-1) >>= fun i -> 
                [ map (fun (hd,tl) -> hd::tl) (enum i ** list_j (n-i))]
              end
          in 
          let enum_list3 = incr enum_list2 in 
          let enum_list2 = incr enum_list1 in 
          let enum_n = 
            match n with 
            | 0 -> empty 
            | 1 -> one Hole 
            | _ -> 
               bigsum 
                 [ map (fun l -> HoledP l) enum_list2;
                   map (fun l -> HoledQ l) enum_list3
                 ]
          in
          let enum_list1 = 
            (map (fun t -> [t]) enum_n) ++ incr enum_list
          in 
          let enum_list =
            match n with 
            | 0 -> one [] 
            | _ -> enum_list1  
          in 
          ( enum_n,
            enum_list,
            enum_list1,
            enum_list2
          )
        )

      let enum_skeleton n = (fun (fst,_,_,_) -> fst)  (enum_skeletons n)


      include PQTEnum
      let enumeration = 
        PQTEnum.map to_pqtree enum_skeleton


      let extensions list = 
        let rec go previous = function 
        | [] -> []
        | this::next -> (this, List.rev_append previous next) :: go (this::previous) next
        in 
        go [] list


      let rec enum_permutations = 
        let open PQTIFSeq in 
        function 
        | [] -> singleton [] 
        | elements ->
            exists 
              (extensions elements)
              (fun (head,tail) -> map (List.cons head) (enum_permutations tail))
           
            
      let rec big_product = 
        let open PQTIFSeq in 
        function 
        | [] -> singleton []
        | head::tail -> 
            product head (big_product tail)
            |> map (fun (sample_head, samples_tail) -> sample_head :: samples_tail)

      let apply_permutation (sigma : int list) (elements : 'a list) : 'a list = 
        List.map (List.nth elements) sigma

      let rec compatible_orders pqtree =
        let open PQTIFSeq in 
        match pqtree with 
        | Leaf x -> singleton [x]
        | P children -> 
            let permutations = enum_permutations (List.mapi (fun i _ -> i) children) in
            let orders = big_product (List.map compatible_orders children) in 
            product permutations orders
            |> map (fun (permutation, orders) -> apply_permutation permutation orders |> List.concat)
        | Q children -> 
            let left_to_right_seq = 
              children 
              |> List.map compatible_orders
              |> big_product 
              |> map List.concat 
            in 
            sum (left_to_right_seq) (left_to_right_seq |> map List.rev)
                        
  end 


module Enum = MakeEnumerate(Stdlib.Random) 

let sample size = 
  match Enum.sample 1 Enum.enumeration size (size+1) Seq.empty () with 
  | Seq.Cons (pqtree,_) -> pqtree 
  | _ -> invalid_arg (Format.sprintf "PqTree.sample %d" size)


let sample_compatible_order pqtree = 
  let sequence = Enum.compatible_orders pqtree in 
  match Seq.uncons (Enum.PQTIFSeq.sample 1 sequence Seq.empty) with 
  | Some (sample,_) -> sample 
  | None -> assert false



let rec shrink_a_child children = 
  MoreList.extend children 
  |> List.to_seq 
  |> Seq.flat_map (function 
         | (P grandchildren as child, pre, post)
         | (Q grandchildren as child, pre, post ) -> 
            let reset c = List.flatten [List.rev pre; c; post] in
            shrink child 
            |> Seq.map (fun new_child -> reset [new_child])
            |> Seq.cons (reset grandchildren)
         | (Leaf _,_,_) -> 
            Seq.empty
       )

and shrink : ('a t -> 'a t Seq.t) = function
    | Leaf _ -> Seq.empty
    | P children -> 
       children
       |> shrink_a_child
       |> Seq.map (fun new_children -> P new_children)
    | Q children -> 
       children
       |> shrink_a_child
       |> Seq.map (fun new_children -> Q new_children)
       

module type CanonicalSig = 
  sig 
    type elt 
    val canonical : elt t -> elt t 
    val compare : elt t -> elt t -> int 
    val compare' : elt t -> elt t -> int 
    type nonrec t = elt t
  end 


module MakeCanonical = 
  functor (C : Map.OrderedType) -> 
  struct 
    
    type elt = C.t
    type nonrec t = elt t

    let canonical tree = 
      let rec go = function 
        | Leaf e -> (e, Leaf e) 
        | Q children -> 
           let canonicals = List.map go children in 
           let first = fst (List.hd canonicals) in 
           let last = List.fold_left (fun _ (e,_) -> e) first canonicals in 
           if C.compare first last <= 0 then 
             (first, Q (List.map snd canonicals))
           else 
             (last, Q (List.map snd canonicals |> List.rev))
        | P children -> 
           let canonicals = 
             children
             |> List.map go 
             |> List.sort (fun (min1,_) (min2,_) -> C.compare min1 min2)
           in 
           (fst (List.hd canonicals), P (List.map snd canonicals))     
      in 
      snd (go tree)

    let rec compare' tree1 tree2 =
      match tree1, tree2 with 
      | Leaf a, Leaf b -> C.compare a b
      | Leaf _, _ -> -1 
      | _, Leaf _ -> 1
      | P children1, P children2
      | Q children1, Q children2 -> 
         List.compare compare' children1 children2 
      | P _, _ -> -1 
      | _, P _ -> 1 

    let compare tree1 tree2 = compare' (canonical tree1) (canonical tree2)

  end
  



module type EltMap = 
sig
  type key 
  type 'value t

  val empty : 'value t
  val add : key -> 'value -> 'value t -> 'value t
  val find_opt : key -> 'value t -> 'value option 
  val equal : ('value -> 'value -> bool) -> 'value t -> 'value t -> bool
end


module type CheckPermutationSig = 
  sig
    type elt 
    val contains_permutation : elt t -> elt list -> bool
    type checker
    val checker : elt t -> checker
    val checks : checker -> elt list -> bool 
  end 


module MakeCheckPermutation = 
  functor (M : EltMap) ->
  struct 

    type elt = M.key 

    let (>>=) = Option.bind 

    let leaf_reader elt = 
      ( [elt],
        function
        | hd::tail when elt = hd -> Some tail 
        | _ -> None
      )


    let p_reader children_readers =
      let child_map = 
        List.fold_left
          (fun map (elements, reader) ->
            List.fold_left (fun map element -> M.add element reader map) map elements
          )
          M.empty
        children_readers
      in
      let read_child map = function 
        | first::_ as permutation ->
           M.find_opt first map >>= fun reader ->
           reader permutation
        | _ -> None
      in
      let read_children map permutation = 
        List.fold_left 
          (fun permutation _ -> permutation >>= read_child map) 
          (Some permutation)
          children_readers
      in            
      ( List.concat (List.map fst children_readers),
        read_children child_map 
      )

    let q_reader children_readers = 
      let left = fst (List.hd children_readers) in 
      let right = List.fold_left (fun _ l -> fst l) left children_readers in 
      let read_children readers permutation = List.fold_left Option.bind permutation readers in 
      let readers = List.map snd children_readers in 
      let sredaer = List.rev readers in 
      let map =
            List.fold_left 
            (fun map (element,readers) -> M.add element readers map)
            M.empty
            (List.rev_append 
               (List.map (fun e -> (e,readers)) left)
               (List.map (fun e -> (e,sredaer)) right)
            )
      in
      let reader = function
        | first::_ as permutation -> 
           M.find_opt first map >>= fun childrens -> read_children childrens (Some permutation)
        | _ -> None 
      in        
      ( List.rev_append left right,
        reader
      )

    type checker = M.key list * (M.key list -> M.key list option)

    let rec reader = function 
      | Leaf e -> 
         leaf_reader e
      | P children ->
         p_reader (List.map reader children) 
      | Q children ->
         q_reader (List.map reader children)
      
    let checker pqtree = (frontier pqtree, snd (reader pqtree))

    let are_permutation list1 list2 = 
      let count map elt = 
        M.find_opt elt map 
        |> Option.value ~default:0
        |> (+) 1 
        |> fun value -> M.add elt value map
      in 
      let count_map list = List.fold_left count M.empty list in 
      List.length list1 = List.length list2 
      && M.equal (=) (count_map list1) (count_map list2)
 
    let checks (frontier,reader) permutation = 
      are_permutation frontier permutation 
      && reader permutation = Some []

    let contains_permutation pqtree = checks (checker pqtree)

  end



module type PrintableType = sig 
  type t
  val print : Format.formatter -> t -> unit
end

module MakePrinter =
  functor (Prt : PrintableType) -> 
  struct 

    type nonrec t = Prt.t t 


    let rec print formatter = function
        | Leaf elt -> Format.fprintf formatter "@[<h>Leaf %a@]" Prt.print elt 
        | P children -> 
           Format.fprintf formatter "@[<h>P@,[@[<hov 2>%a@]]@]"
             (Format.pp_print_list ~pp_sep:(fun fmt () -> Format.fprintf fmt ";@,") print) 
             children 
        | Q children -> 
           Format.fprintf formatter "@[<h>Q@,[@[<hov 2>%a@]]@]"
             (Format.pp_print_list ~pp_sep:(fun fmt () -> Format.fprintf fmt ";@,") print) 
             children 

  end