open Common
(** Given a statically typed AST generates a lattice of gradually typed ASTs. *)

let verbose = ref false

let log fmt =
  if !verbose then Format.kasprintf (fun s -> Format.printf "%s" s) fmt
  else Format.ifprintf Format.std_formatter fmt

module Interval : sig
  type t

  val make : int -> int -> t option
  val make_exn : int -> int -> t
  val lower : t -> int
  val upper : t -> int
  val width : t -> int
end = struct
  type t = { lower : int; upper : int }

  let lower interval = interval.lower
  let upper interval = interval.upper
  let width interval = upper interval - lower interval
  let make lower upper = if lower <= upper then Some { lower; upper } else None

  let make_exn lower upper =
    match make lower upper with
    | Some interval -> interval
    | None -> invalid_arg "lower must be less than or equal to upper"
end

(** Generates a list of integer intervals from a given number of intervals and
    interval width.*)
let generate_intervals (interval_num : int) (interval_width : float) :
    Interval.t list =
  if interval_num <= 0 then invalid_arg "interval_num must be greater than 0";
  if interval_width <= 0. then
    invalid_arg "interval_width must be greater than 0.";
  let round = if interval_width > 1. then floor else ceil in
  let rec go interval_num interval_width =
    if interval_num > 1 then
      Interval.make_exn
        (Int.of_float
           (ceil (Float.of_int (interval_num - 1) *. interval_width)))
        (Int.of_float (round @@ (Float.of_int interval_num *. interval_width)))
      :: go (interval_num - 1) interval_width
    else [ Interval.make_exn 0 (Int.of_float (round interval_width)) ]
  in
  go interval_num interval_width

open Syntax.Parse_tree

let ( let* ) x f = Option.bind x f

let sequence_map f lst =
  let rec aux acc = function
    | [] -> Some (List.rev acc)
    | x :: xs ->
        let* v = f x in
        let* rest = aux (v :: acc) xs in
        Some rest
  in
  aux [] lst

let option_of_bool (b : bool) : unit option = if b then Some () else None
let get_fst t = fst t.meta
let get_snd t = snd t.meta
let strip_snd (ty : ('a * 'b) Type.t) : 'a Type.t = Type.map_meta fst ty

module Count = struct
  type t = int * int
  (** Pair composed of the size of the lattice and the number of static types *)

  let default_count = (1, 0)
  let combine_count (p1, s1) (p2, s2) = (p1 * p2, s1 + s2)

  let combine_count_list l f =
    List.fold_left (fun acc x -> combine_count acc (f x)) default_count l

  let increase_count (p, s) = (p + 1, s + 1)

  let rec count_ty (ty : 'a Type.t) : int * int =
    match ty.term with
    | Any -> (1, 1)
    | Int | Bool | Float | Var _ -> increase_count default_count
    | Arrow (tys, ty) ->
        increase_count
        @@ combine_count (count_ty ty) (combine_count_list tys count_ty)
    | Product ts -> increase_count @@ combine_count_list ts count_ty
    | Vec t -> increase_count @@ count_ty t
    | Variant cases ->
        increase_count
        @@ combine_count_list cases (fun (Type.Case (_, tys)) ->
               combine_count_list tys count_ty)

  let rec count_expr (e : 'a Expr.t) : int * int =
    match e.term with
    | Const _ | Var _ -> default_count
    | Lam (args, e) ->
        let _, tys = List.split args in
        combine_count (combine_count_list tys count_ty) (count_expr e)
    | App (e, exprs) ->
        combine_count (count_expr e) (combine_count_list exprs count_expr)
    | Prim1 (_, e) -> count_expr e
    | Prim2 (_, e1, e2) -> combine_count (count_expr e1) (count_expr e2)
    | If (e1, e2, e3) ->
        combine_count (count_expr e1)
          (combine_count (count_expr e2) (count_expr e3))
    | Let (binding, e) -> combine_count (count_binding binding) (count_expr e)
    | Letrec (bindings, e) ->
        combine_count_list bindings count_binding
        |> combine_count (1, 0)
        |> combine_count (count_expr e)
    | Tuple es -> combine_count_list es count_expr
    | Get (e, _) -> count_expr e
    | Vect (e1, e2) -> combine_count (count_expr e1) (count_expr e2)
    | VectGet (e1, e2) -> combine_count (count_expr e1) (count_expr e2)
    | VectSet (e1, e2, e3) ->
        combine_count (count_expr e1)
          (combine_count (count_expr e2) (count_expr e3))
    | VectLength e -> count_expr e
    | Ascribe (e, ty) -> combine_count (count_expr e) (count_ty ty)
    | For (_, e1, e2, e3) ->
        combine_count (count_expr e1)
          (combine_count (count_expr e2) (count_expr e3))
    | Variant (_, es) -> combine_count_list es count_expr
    | Match (e, patterns) ->
        combine_count (count_expr e)
          (combine_count_list patterns (fun pat ->
               count_expr pat.term.Expr.body))
    | Seq es -> combine_count_list es count_expr

  and count_binding (binding : 'a Expr.binding) : int * int =
    match binding.term.ty with
    | None -> count_expr binding.term.e
    | Some ty -> combine_count (count_ty ty) (count_expr binding.term.e)

  let count_program (program : 'a Program.t) : int * int =
    let expr_count = count_expr program.term.expr in
    expr_count
end

let rec annotate_type_with_count (ty : 'a Type.t) : ('a * int) Type.t =
  match ty.term with
  | Any -> annotate Type.Any (ty.meta, 0)
  | Int -> annotate Type.Int (ty.meta, 1)
  | Bool -> annotate Type.Bool (ty.meta, 1)
  | Float -> annotate Type.Float (ty.meta, 1)
  | Arrow (ty_args, ty_ret) ->
      let ty_args' = List.map annotate_type_with_count ty_args in
      let ty_ret' = annotate_type_with_count ty_ret in
      let tys_count =
        List.fold_left (fun acc t -> acc + get_snd t) 0 ty_args'
      in
      annotate
        (Type.Arrow (ty_args', ty_ret'))
        (ty.meta, 1 + get_snd ty_ret' + tys_count)
  | Product ts ->
      let ts' = List.map annotate_type_with_count ts in
      annotate (Type.Product ts')
        (ty.meta, 1 + List.fold_left (fun acc t -> acc + get_snd t) 0 ts')
  | Vec t -> annotate (Type.Vec (annotate_type_with_count t)) (ty.meta, 1)
  | Var type_name -> annotate (Type.Var type_name) (ty.meta, 1)
  | Variant cases ->
      let cases' =
        List.map
          (fun (Type.Case (name, tys)) ->
            Type.Case (name, List.map annotate_type_with_count tys))
          cases
      in
      annotate (Type.Variant cases') (ty.meta, 1)

let local_lattice (program : 'a Program.t) : ('a * int) Type.t list =
  let rec go_expr (e : 'a Expr.t) : ('a * int) Type.t list =
    match e.term with
    | Const _ -> []
    | Var _ -> []
    | Lam (args, e) ->
        let _, tys = List.split args in
        List.map annotate_type_with_count tys @ go_expr e
    | App (e, exprs) -> go_expr e @ List.flatten (List.map go_expr exprs)
    | Prim1 (_, e) -> go_expr e
    | Prim2 (_, e1, e2) -> go_expr e1 @ go_expr e2
    | If (e1, e2, e3) -> go_expr e1 @ go_expr e2 @ go_expr e3
    | Let (binding, e) -> go_binding binding @ go_expr e
    | Letrec (bindings, e) ->
        List.flatten (List.map go_binding bindings) @ go_expr e
    | Tuple es -> List.flatten (List.map go_expr es)
    | Get (e, _) -> go_expr e
    | Vect (e1, e2) -> go_expr e1 @ go_expr e2
    | VectGet (e1, e2) -> go_expr e1 @ go_expr e2
    | VectSet (e1, e2, e3) -> go_expr e1 @ go_expr e2 @ go_expr e3
    | VectLength e -> go_expr e
    | Ascribe (e, ty) -> go_expr e @ [ annotate_type_with_count ty ]
    | For (_, e1, e2, e3) -> go_expr e1 @ go_expr e2 @ go_expr e3
    | Variant (_, es) -> List.flatten (List.map go_expr es)
    | Match (e, patterns) ->
        go_expr e
        @ List.flatten
            (List.map (fun pat -> go_expr pat.term.Expr.body) patterns)
    | Seq es -> List.flatten (List.map go_expr es)
  and go_binding (binding : 'a Expr.binding) : ('a * int) Type.t list =
    match binding.term.ty with
    | None -> go_expr binding.term.e
    | Some ty -> annotate_type_with_count ty :: go_expr binding.term.e
  and go_program (program : 'a Program.t) : ('a * int) Type.t list =
    (* types in decls should not be considered to be potentially deleted. This
       changes the meaning of the program *)
    go_expr program.term.expr
  in
  go_program program

let generate_lattice_info (program : 'a Program.t) :
    ('a * int) Type.t list * int =
  let tys = local_lattice program in
  (tys, List.fold_left (fun acc ty -> acc + get_snd ty) 0 tys)

(* Fisher-Yates shuffle *)
let shuffle (random_state : Random.State.t) (l : 'a list) : 'a list =
  let a = Array.of_list l in
  for i = Array.length a - 1 downto 1 do
    let j = Random.State.int random_state (i + 1) in
    let v = a.(i) in
    a.(i) <- a.(j);
    a.(j) <- v
  done;
  Array.to_list a

(** Generates a list of samples from a list of weighted items.

    Algorithm from: Weighted Random Sampling (2005; Efraimidis, Spirakis)

    Algorithm A Input : A population V of n weighted items Output : A WRS of
    size m 1: For each v_i ∈ V , u_i = random(0, 1) and k_i = u_i^(1/w_i) 2:
    Select the m items with the largest keys k_i as a WRS *)
let weighted_random_sample random_state (list : (int * 'value) list) num_samples
    =
  list
  |> List.map (fun (weight, value) ->
         (weight, value, Random.State.float random_state 1.0))
  |> List.map (fun (weight, value, random) ->
         (weight, value, random ** (1.0 /. Float.of_int weight)))
  |> List.sort (fun (_, _, k1) (_, _, k2) -> compare k2 k1)
  |> List.fold_left
       (fun acc (_, value, _) ->
         if List.length acc < num_samples then value :: acc else acc)
       []
  |> List.rev

let weighted_random_shuffle seed list =
  weighted_random_sample seed list (List.length list)

let reorder_type_list l =
  l |> List.sort (fun (x1, _) (x2, _) -> Int.compare x1 x2) |> List.map snd

let tag_types_with_indices types = List.mapi (fun i ty -> (i, ty)) types
let create_weighted_types_list l = List.map (fun ty -> (ty, get_snd ty)) l

(** generates a random in the range [lo, hi] using the provided random state. *)
let random_int_in_range state lo hi =
  if lo > hi then
    invalid_arg
      (Format.sprintf
         "Invalid range: [%d, %d]. 'lo' must be less than or equal to 'hi'" lo
         hi)
  else lo + Random.State.int state (hi - lo + 1)

let rec sample_less_precise_type :
    Random.State.t -> ('a * int) Type.t -> int -> 'a Type.t option =
 fun random_state static_ty -> sample_type random_state static_ty

and sample_type random_state (ty : ('a * int) Type.t) n : 'a Type.t option =
  let pp_meta fmt (_, n) = Format.fprintf fmt "<%d>" n in
  let invalid_sampling ty n =
    log "sample_type: failed when trying to sample %d nodes from %a@." n
      (Type.pp ~pp_meta) ty;
    None
  in
  match ty with
  | { term = Any; meta = a, _ } ->
      if n = 0 then Option.some @@ annotate Type.Any a
      else invalid_sampling ty n
  | { term = _; meta = _, n' } when n = n' -> Option.some @@ strip_snd ty
  | { term = _; meta = a, _ } when n = 0 -> Option.some @@ annotate Type.Any a
  | { term = Int; meta = a, _ } ->
      if n = 1 then Option.some @@ annotate Type.Int a
      else invalid_sampling ty n
  | { term = Bool; meta = a, _ } ->
      if n = 1 then Option.some @@ annotate Type.Bool a
      else invalid_sampling ty n
  | { term = Float; meta = a, _ } ->
      if n = 1 then Option.some @@ annotate Type.Float a
      else invalid_sampling ty n
  | { term = Arrow (tys, ty); meta = a, type_const_count } ->
      let* ty_sample =
        sample_type_list random_state (n - 1) (type_const_count - 1) (ty :: tys)
      in
      let tys', ty' =
        match ty_sample with
        | [] | [ _ ] -> failwith "should never happen"
        | hd :: tl -> (tl, hd)
      in
      Option.some @@ annotate (Type.Arrow (tys', ty')) a
  | { term = Product tys; meta = a, type_const_count } ->
      let* tys' = sample_type_list random_state n (type_const_count - 1) tys in
      Option.some @@ annotate (Type.Product tys') a
  | { term = Vec ty; meta = a, type_const_count } ->
      let* ty' = sample_type random_state ty (type_const_count - 1) in
      Option.some @@ annotate (Type.Vec ty') a
  | { term = Var x; meta = a, _ } ->
      if n = 1 then Option.some @@ annotate (Type.Var x) a
      else invalid_sampling ty n
  | { term = Variant cases; meta = a, type_const_count } ->
      let* cases' =
        sequence_map
          (fun (Type.Case (name, tys)) ->
            let* tys' =
              sample_type_list random_state n (type_const_count - 1) tys
            in
            Option.some @@ Type.Case (name, tys'))
          cases
      in
      Option.some @@ annotate (Type.Variant cases') a

and sample_type_list random_state target potential types =
  let g
      ((l', (remaining_potential, picked_type_node_count)) :
        (int * 'a Type.t) list * (int * int))
      ((idx, static_type) : int * ('a * int) Type.t) :
      ((int * 'a Type.t) list * (int * int)) option =
    let max_type_node_count = get_snd static_type in
    log "sample_type_list: remaining_potential: %d@." remaining_potential;
    log "sample_type_list: max local type node count: %d@." max_type_node_count;
    let remaining_potential' = remaining_potential - max_type_node_count in
    let lo = max 0 (target - remaining_potential' - picked_type_node_count) in
    let hi = min max_type_node_count (target - picked_type_node_count) in
    let random_type_node_count = random_int_in_range random_state lo hi in
    log
      "sample_type_list: random type node count: %d picked from the interval: \
       [%d, %d]@."
      random_type_node_count lo hi;
    match sample_type random_state static_type random_type_node_count with
    | None -> None
    | Some ty_opt ->
        log
          "sample_type_list: sampled less precise type: %a from the static \
           type %a@."
          (Type.pp ~pp_meta:(fun _ _ -> ()))
          ty_opt
          (Type.pp ~pp_meta:(fun _ _ -> ()))
          static_type;
        Some
          ( (idx, ty_opt) :: l',
            ( remaining_potential',
              picked_type_node_count + random_type_node_count ) )
  in
  let* () = option_of_bool (potential >= target) in
  (* indices are used to reconstruct the list *)
  let type_list_with_indices = List.mapi (fun i ty -> (i, ty)) types in
  (* weights are used when shuffling *)
  let type_list_with_weights =
    List.map (fun ((_, ty) as y) -> (get_snd ty, y)) type_list_with_indices
  in
  (* weighted shuffle *)
  let shuffled_list =
    weighted_random_shuffle random_state type_list_with_weights
  in
  log "sample_type_list: shuffled list: %a@."
    (Format.pp_print_list (fun fmt (i, ty) ->
         Format.fprintf fmt "(%d, %a)" i (Type.pp ~pp_meta:(fun _ _ -> ())) ty))
    shuffled_list;
  shuffled_list
  |> List.fold_left
       (fun acc x ->
         let* l', (remaining_potential, picked_type_node_count) = acc in
         g (l', (remaining_potential, picked_type_node_count)) x)
       (Some ([], (potential, 0)))
  |> Option.map (fun (lst, _) -> reorder_type_list lst)

module Counter : sig
  type t

  val init : unit -> t
  val fresh : t -> int
end = struct
  type t = int ref

  let init () = ref 0

  let fresh counter =
    let n = !counter in
    counter := n + 1;
    n
end

let sample_one random_state
    ((potential, ts) : int * ((int * 'meta) * int) Type.t list) interval :
    (int * 'meta) Type.t list =
  let pp_meta fmt ((id, _), _) = Format.fprintf fmt "<id=%d>" id in
  let f ((l, (remaining_potential, picked_type_node_count)) as x)
      ({ term = _; meta = _, max_type_node_count } as static_type :
        ((int * 'meta) * int) Type.t) =
    let remaining_potential' = remaining_potential - max_type_node_count in
    let lo =
      max 0
        (Interval.lower interval - remaining_potential' - picked_type_node_count)
    in
    let hi =
      min max_type_node_count (Interval.upper interval - picked_type_node_count)
    in
    let random_type_node_count = random_int_in_range random_state lo hi in
    log
      "sample_one:     random type node count: %d picked from the interval: \
       [%d, %d]@."
      random_type_node_count lo hi;
    let ty_opt =
      sample_less_precise_type random_state static_type random_type_node_count
    in
    log
      "sample_one:     sampled less precise type: %a from the static type %a@."
      (Format.pp_print_option
         ~none:(fun fmt () -> Format.fprintf fmt "None")
         (Type.pp ~pp_meta:(fun fmt (id, _) -> Format.fprintf fmt "<id=%d>" id)))
      ty_opt (Type.pp ~pp_meta) static_type;
    if remaining_potential <= 0 then x
    else
      match ty_opt with
      | None -> (l, (remaining_potential', picked_type_node_count))
      | Some ty ->
          ( ty :: l,
            ( remaining_potential',
              picked_type_node_count + random_type_node_count ) )
  in
  log "sample_one:@.";
  log "sample_one:   using the following shuffled list of types: %a@."
    (Format.pp_print_list
       ~pp_sep:(fun fmt _ -> Format.fprintf fmt ",")
       (fun fmt ty ->
         Format.fprintf fmt "%d: %a" (fst (fst ty.meta)) (Type.pp ~pp_meta) ty))
    ts;
  log "sample_one:   target interval: [%d, %d]@." (Interval.lower interval)
    (Interval.upper interval);
  log "sample_one:   potential is: %d@." potential;
  let samples, _ = List.fold_left f ([], (potential, 0)) ts in
  log "sample_one:   result: %a@."
    (Format.pp_print_list
       ~pp_sep:(fun fmt _ -> Format.fprintf fmt ",")
       (fun fmt ty ->
         Format.fprintf fmt "%d: %a" (fst ty.meta)
           (Type.pp ~pp_meta:(fun fmt (id, _) ->
                Format.fprintf fmt "<id=%d>" id))
           ty))
    samples;
  samples

let dynamize_ty (ty : 'a Type.t) : 'a Type.t =
  match ty.term with
  | Any | Int | Bool | Float | Arrow _ | Product _ | Vec _ | Var _ | Variant _
    ->
      { ty with term = Any }

let to_binary_string (replacement_list : bool list) : string =
  replacement_list
  |> List.to_seq
  |> Seq.map (fun b -> if b then '1' else '0')
  |> String.of_seq

let to_hex_string (bools : bool list) : string =
  let hex_chars =
    [|
      '0';
      '1';
      '2';
      '3';
      '4';
      '5';
      '6';
      '7';
      '8';
      '9';
      'A';
      'B';
      'C';
      'D';
      'E';
      'F';
    |]
  in

  let pad l =
    let length = List.length l in
    let extra = length land 0b11 in
    let padding_length = if extra = 0 then 0 else 4 - extra in
    log "to_hex_string: length=%d extra=%d  padding_length=%d\n%!" length extra
      padding_length;
    let padding = List.init padding_length (fun _ -> false) in
    padding @ l
  in

  let rec process acc bit_count value = function
    | [] ->
        if bit_count = 0 then List.rev acc
        else
          let hex_digit = hex_chars.(value) in
          List.rev (hex_digit :: acc)
    | b :: rest ->
        let new_value = (value lsl 1) lor if b then 1 else 0 in
        if bit_count = 3 then
          let hex_digit = hex_chars.(new_value) in
          process (hex_digit :: acc) 0 0 rest
        else process acc (bit_count + 1) new_value rest
  in

  if bools = [] then failwith "Can't hex encode empty list";
  bools |> pad |> process [] 0 0 |> List.to_seq |> String.of_seq

let rec replacement_list (ty : 'a Type.t) (ty' : 'a Type.t) : bool list =
  match (ty.term, ty'.term) with
  | Type.Any, Type.Any -> []
  | Type.Any, _ -> failwith "should be impossible"
  | Type.Int, Type.Any | Type.Bool, Type.Any | Type.Float, Type.Any -> [ false ]
  | Type.Int, Type.Int | Type.Bool, Type.Bool | Type.Float, Type.Float ->
      [ true ]
  | Type.Int, _ | Type.Bool, _ | Type.Float, _ ->
      failwith "should be impossible"
  | Type.Arrow (ty_args, ty_ret), Type.Any ->
      [ false ]
      @ List.flatten (List.map (fun ty -> replacement_list ty ty') ty_args)
      @ replacement_list ty_ret ty'
  | Type.Arrow (ty_args, ty_ret), Type.Arrow (ty_args', ty_ret') ->
      [ true ]
      @ List.flatten (List.map2 replacement_list ty_args ty_args')
      @ replacement_list ty_ret ty_ret'
  | Type.Arrow _, _ -> failwith "should be impossible"
  | Type.Product tys, Type.Any ->
      [ false ]
      @ List.flatten (List.map (fun ty -> replacement_list ty ty') tys)
  | Type.Product tys, Type.Product tys' ->
      [ true ] @ List.flatten (List.map2 replacement_list tys tys')
  | Type.Product _, _ -> failwith "should be impossible"
  | Type.Vec ty, Type.Any -> [ false ] @ replacement_list ty ty'
  | Type.Vec ty, Type.Vec ty' -> [ true ] @ replacement_list ty ty'
  | Type.Vec _, _ -> failwith "should be impossible"
  | Type.Variant cases, Type.Any ->
      [ false ]
      @ List.flatten
          (List.map
             (fun (Type.Case (_, args)) ->
               List.flatten
                 (List.map (fun arg -> replacement_list arg ty') args))
             cases)
  | Type.Variant cases, Type.Variant cases' ->
      [ true ]
      @ List.flatten
          (List.map2
             (fun (Type.Case (_, args)) (Type.Case (_, args')) ->
               List.flatten (List.map2 replacement_list args args'))
             cases cases')
  | Type.Variant _, _ -> failwith "should be impossible"
  | Type.Var _, Type.Any -> [ false ]
  | Type.Var _, Type.Var _ -> [ true ]
  | Type.Var _, _ -> failwith "should be impossible"

let replace_types env program : (int * 'a) Program.t * bool list =
  let list_ref = ref [] in
  let replace_types' (ty : (int * 'a) Type.t) : 'b Type.t =
    match Hashtbl.find_opt env (fst ty.meta) with
    | Some ty' ->
        list_ref := replacement_list ty ty' :: !list_ref;
        ty'
    | None ->
        list_ref := replacement_list ty ty :: !list_ref;
        ty
  in
  let program' = Program.map_ty replace_types' program in
  let replacement_list = List.flatten (List.rev !list_ref) in
  (program', replacement_list)

let sample_many samples_count ast types_info types_nodes_count interval =
  (* TODO: better handling of seed *)
  let seed = 42 in
  let random_state = Random.State.make [| seed |] in
  List.map
    (fun _ ->
      let shuffled_types = shuffle random_state types_info in
      let samples =
        sample_one random_state (types_nodes_count, shuffled_types) interval
      in
      let env =
        samples
        |> List.map (fun ty -> (fst ty.meta, ty))
        |> (fun l ->
        log "sample_many: %a\n%!"
          (Format.pp_print_list
             ~pp_sep:(fun fmt () -> Format.fprintf fmt ", ")
             (fun fmt (id, ty) ->
               Format.fprintf fmt "%d : %a" id
                 (Type.pp ~pp_meta:(fun fmt (id, _) ->
                      Format.fprintf fmt "<id=%d>" id))
                 ty))
          l;
        l)
        |> List.to_seq
        |> Hashtbl.of_seq
      in
      replace_types env ast)
    (List.init samples_count (fun _ -> ()))

let sample_from_bins ?(include_extremes = true) (ast : 'meta Program.t)
    num_samples num_bins =
  (* add unique index to each node in the ast *)
  let counter = Counter.init () in
  let ast' = Program.map_meta (fun meta -> (Counter.fresh counter, meta)) ast in
  log "sample_from_bins: ast:\n%a\n%!"
    (Program.pp ~pp_meta:(fun fmt (id, _) -> Format.fprintf fmt "<id=%d>" id))
    ast';
  let type_info, type_nodes_count = generate_lattice_info ast' in
  let pp_meta fmt ((id, _), _) = Format.fprintf fmt "<id=%d>" id in
  log "sample_from_bins: Types: [%a]\n%!"
    (Format.pp_print_list
       ~pp_sep:(fun fmt () -> Format.fprintf fmt ", ")
       (fun fmt ty ->
         Format.fprintf fmt "%d : %a" (fst (fst ty.meta)) (Type.pp ~pp_meta) ty))
    type_info;
  let interval_width = Float.of_int type_nodes_count /. Float.of_int num_bins in
  let intervals = generate_intervals num_bins interval_width in
  log
    "sample_from_bins: proceeding to sample %d samples from each of the \
     following intervals %a@."
    num_samples
    (Pp.pp_list ~sep:", " (fun fmt interval ->
         Format.fprintf fmt "[%d, %d]" (Interval.lower interval)
           (Interval.upper interval)))
    intervals;
  let sample num_samples interval =
    sample_many num_samples ast' type_info type_nodes_count interval
  in
  let samples = List.map (sample num_samples) intervals in
  let samples =
    if include_extremes then
      let static_interval =
        Interval.make_exn type_nodes_count type_nodes_count
      in
      let dynamic_interval = Interval.make_exn 0 0 in
      [ sample 1 static_interval ] @ samples @ [ sample 1 dynamic_interval ]
    else samples
  in
  samples
