(** Generate C++ from the MIR.

    This module makes extensive use of the Format[0] module via the Fmt[1] API.
    As such, you'll need to understand the "%a" and "@" notation from [0], especially
    the section headed "Formatted pretty-printing." Then, we use functions like
    [pf] and [str] from the Fmt library[1]. On top of that, the "@" pretty-printing
    specifiers all actually correspond 1-to-1 with commands like [open_box] from
    the Format library[0]. The boxing system is best described in this explainer
    pdf [2], particularly Section 3 ("Format basics"). It's worth noting that someone
    was able to make a decent-looking pretty-printer for a subset of Javascript[3] that
    might serve as a good reference. Good luck!

    [0] Format module doc: https://ocaml.org/api/Format.html
    [1] Fmt module doc: https://erratique.ch/software/fmt/doc/Fmt/index.html
    [2] Format Unraveled: https://hal.archives-ouvertes.fr/hal-01503081/file/format-unraveled.pdf
    [3] Javascript pretty-printer https://github.com/Virum/compiler/blob/28e807b842bab5dcf11460c8193dd5b16674951f/JavaScript.ml#L112
*)

open Core_kernel
open Core_kernel.Poly
open Middle
open Fmt
open Expression_gen
open Statement_gen
open Function_gen

let standalone_functions = ref false

(* set by drivers *)
let stanc_args_to_print = ref ""

(** Print name of model function.
  @param prog_name Name of the Stan program.
  @param fname Name of the function.
 *)
let pp_function__ ppf (prog_name, fname) =
  pf ppf
    {|@[<v>static constexpr const char* function__ = "%s_namespace::%s";@,%a@]|}
    prog_name fname pp_unused "function__"

let version = "// Code generated by %%NAME%% %%VERSION%%"
let includes = "#include <stan/model/model_header.hpp>"

(** Validate the dimensions of the C++ object are correct at runtime
 @param ppf A pretty printer
 @param name The name of the object.
 @param st The SizedType of the object.
 *)
let pp_validate_data ppf (name, st) =
  if String.is_suffix ~suffix:"__" name then ()
  else
    let pp_stdvector ppf args =
      let pp_cast ppf x = pf ppf "static_cast<size_t>(%a)" pp_expr x in
      pf ppf "@[<hov 2> std::vector<size_t>{@,%a}@]" (list ~sep:comma pp_cast)
        args in
    pf ppf "@[<hov 4>context__.validate_dims(@,%S,@,%S,@,%S,@,%a);@]@ "
      "data initialization"
      (Mangle.remove_prefix name)
      (stantype_prim_str (SizedType.to_unsized st))
      pp_stdvector (SizedType.get_dims_io st)

let pp_mul ppf () = pf ppf " * "

let get_unconstrained_param_st lst =
  match lst with
  | _, {Program.out_block= Parameters; out_unconstrained_st= st; _} -> (
    match SizedType.get_dims_io st with
    | [] -> Some [Expr.Helpers.loop_bottom]
    | ls -> Some ls )
  | _ -> None

let get_constrained_param_st lst =
  match lst with
  | _, {Program.out_block= Parameters; out_constrained_st= st; _} -> (
    match SizedType.get_dims_io st with
    | [] -> Some [Expr.Helpers.loop_bottom]
    | ls -> Some ls )
  | _ -> None

let pp_num_param ppf (dims : Expr.Typed.t list) =
  match dims with
  | [a] -> pf ppf "@[%a@]@," (list ~sep:pp_mul pp_expr) [a]
  | _ -> pf ppf "@[(%a)@]@," (list ~sep:pp_mul pp_expr) dims

(** Print the constructor of the model class.
 Read in data steps:
   1. context__.validate_dims() to verify the dimensions are correct at runtime.
   1. find vals_%s__ from context__.vals_%s(vident)
   1. keep track of pos__
   1. run checks on resulting vident
*)
let pp_ctor ppf p =
  let params =
    [ "stan::io::var_context& context__"; "unsigned int random_seed__ = 0"
    ; "std::ostream* pstream__ = nullptr" ] in
  pf ppf "%s(@[<hov 0>%a) : model_base_crtp(0) @]" p.Program.prog_name
    (list ~sep:comma string) params ;
  let data_idents = List.map ~f:fst p.input_vars |> String.Set.of_list in
  let pp_stmt_topdecl_size_only ppf (Stmt.Fixed.{pattern; meta} as s) =
    match pattern with
    | Decl {decl_id; decl_type; _} when decl_id <> "pos__" -> (
      match decl_type with
      | Sized st -> (
          Locations.pp_smeta ppf meta ;
          let is_input_data = Set.mem data_idents decl_id in
          match is_input_data with
          | true ->
              pp_validate_data ppf (decl_id, st) ;
              pp_assign_data ppf (decl_id, st, false)
          | false -> pp_assign_data ppf (decl_id, st, true) )
      | Unsized _ -> () )
    | _ -> pp_statement ppf s in
  pp_block ppf
    ( (fun ppf {Program.prog_name; prepare_data; output_vars; _} ->
        pf ppf "int current_statement__ = 0;@ " ;
        pf ppf "using local_scalar_t__ = double ;@ " ;
        pf ppf "boost::ecuyer1988 base_rng__ = @ " ;
        pf ppf "    stan::services::util::create_rng(random_seed__, 0);@ " ;
        pf ppf "%a@ " pp_unused "base_rng__" ;
        pf ppf "%a@ " pp_function__ (prog_name, prog_name) ;
        pf ppf
          "local_scalar_t__ \
           DUMMY_VAR__(std::numeric_limits<double>::quiet_NaN());@ " ;
        pf ppf "%a" pp_unused "DUMMY_VAR__" ;
        pp_located_error ppf
          (pp_block, (list ~sep:cut pp_stmt_topdecl_size_only, prepare_data)) ;
        cut ppf () ;
        let output_params =
          List.filter_map ~f:get_unconstrained_param_st output_vars in
        let pp_plus ppf () = pf ppf " + " in
        let pp_set_params ppf pars = (list ~sep:pp_plus pp_num_param) ppf pars in
        match output_params with
        | [] -> pf ppf "num_params_r__ = 0U;@,"
        | _ ->
            pf ppf "@[<hov 2>num_params_r__ = %a;@]@," pp_set_params
              output_params )
    , p )

let rec top_level_decls Stmt.Fixed.{pattern; _} =
  match pattern with
  | Decl d when d.decl_id <> "pos__" ->
      [(d.decl_id, Type.to_unsized d.decl_type)]
  | SList stmts -> List.concat_map ~f:top_level_decls stmts
  | _ -> []

(** Print the private data members of the model class *)
let pp_model_private ppf {Program.prepare_data; _} =
  let data_decls = List.concat_map ~f:top_level_decls prepare_data in
  (*Filter out Any data that is not an Eigen matrix*)
  let get_eigen_map (name, ut) =
    if UnsizedType.is_eigen_type ut && not (Transform_Mir.is_opencl_var name)
    then true
    else false in
  let eigen_map_decls = (List.filter ~f:get_eigen_map) data_decls in
  pf ppf "%a @ %a"
    (list ~sep:cut pp_data_decl)
    data_decls
    (list ~sep:cut pp_map_decl)
    eigen_map_decls

(** Print the signature and blocks of the model class methods.
  @param ppf A pretty printer
  @param rt The return type.
  @param name The method name.
  @param intro Anything that needs printed before the method body.
  @param outro Anything that needs printed after the method body.
  @param cv_attr Optional parameter to add method attributes.
  @param ppbody (?A pretty printer of the method's body)
 *)
let pp_method ppf rt name params intro ?(outro = nop)
    ?(cv_attr : string list = ["const"]) ppbody =
  pf ppf "@[<v 2>inline %s %s(@[<hov>@,%a@]) %a " rt name
    (list ~sep:comma string) params (list ~sep:cut string) cv_attr ;
  pf ppf "{@,%a@ " intro () ;
  ppbody ppf ;
  outro ppf () ;
  pf ppf "@,} // %s() @,@]" name

(** Print the [get_param_names] method of the model class
  @param ppf A pretty printer.
 *)
let pp_get_param_names ppf {Program.output_vars; _} =
  let add_param = fmt "%S" in
  let extract_name var = Mangle.remove_prefix (fst var) in
  pp_method ppf "void" "get_param_names"
    ["std::vector<std::string>& names__"]
    nop
    (fun ppf ->
      pf ppf "@[<hov 2>names__ = std::vector<std::string>{%a};@]@,"
        (list ~sep:comma add_param)
        (List.map ~f:extract_name output_vars) )
    ~cv_attr:["const"]

(** Print the [get_dims] method of the model class. *)
let pp_get_dims ppf {Program.output_vars; _} =
  let pp_cast ppf cast_dims =
    pf ppf "@[<hov 2>static_cast<size_t>(%a)@]@," pp_expr cast_dims in
  let pp_pack ppf inner_dims =
    pf ppf "std::vector<size_t>{@[<hov>@,%a@]}" (list ~sep:comma pp_cast)
      inner_dims in
  let pp_add_pack ppf dims = pf ppf "%a" pp_pack dims in
  let dim_list =
    List.(
      map ~f:(fun (_, {Program.out_constrained_st= st; _}) -> st) output_vars)
  in
  let pp_output_var ppf dims =
    (list ~sep:comma pp_add_pack) ppf List.(map ~f:SizedType.get_dims_io dims)
  in
  pp_method ppf "void" "get_dims"
    ["std::vector<std::vector<size_t>>& dimss__"]
    nop
    (fun ppf ->
      pf ppf "@[<hov 2>dimss__ = std::vector<std::vector<size_t>>{%a};@]@,"
        pp_output_var dim_list )
    ~cv_attr:["const"]

let pp_method_b ppf rt name params intro ?(outro = nop) ?(cv_attr = ["const"])
    body =
  pp_method ppf rt name params intro
    (fun ppf -> pp_located_error_b ppf body)
    ~outro ~cv_attr

(** Print the write_array method of the model class *)
let pp_write_array ppf {Program.prog_name; generate_quantities; _} =
  pf ppf
    "template <typename RNG, typename VecR, typename VecI, typename VecVar, @ \
     stan::require_vector_like_vt<std::is_floating_point, VecR>* = nullptr, @ \
     stan::require_vector_like_vt<std::is_integral, VecI>* = nullptr, @ \
     stan::require_vector_vt<std::is_floating_point, VecVar>* = nullptr> @ " ;
  let params =
    [ "RNG& base_rng__"; "VecR& params_r__"; "VecI& params_i__"; "VecVar& vars__"
    ; "const bool emit_transformed_parameters__ = true"
    ; "const bool emit_generated_quantities__ = true"
    ; "std::ostream* pstream__ = nullptr" ] in
  let intro ppf () =
    pf ppf "%a@ %a@ %a" (list ~sep:cut string)
      [ "using local_scalar_t__ = double;"
      ; "stan::io::deserializer<local_scalar_t__> in__(params_r__, params_i__);"
      ; "stan::io::serializer<local_scalar_t__> out__(vars__);"
      ; "static constexpr bool propto__ = true;"; "(void) propto__;"
      ; "double lp__ = 0.0;"
      ; "(void) lp__;  // dummy to suppress unused var warning"
      ; "int current_statement__ = 0; "
      ; "stan::math::accumulator<double> lp_accum__;"
      ; "local_scalar_t__ \
         DUMMY_VAR__(std::numeric_limits<double>::quiet_NaN());"
      ; "constexpr bool jacobian__ = false;" ]
      pp_unused "DUMMY_VAR__" pp_function__ (prog_name, "write_array") in
  pp_method_b ppf "void" "write_array_impl" params intro generate_quantities

(** Prints the for loop for [constrained_param_names]
    and [unconstrained_param_names]
  @param index_ids Optional named parameter of a SizedType's dimensions
  @param ppf A pretty printer
  @param dims A list of the dimensions of a SizedType
  @param pp_body Pretty printer for the body of the loop.
 *)
let rec pp_for_loop_iteratee ?(index_ids = []) ppf (iteratee, dims, pp_body) =
  let iter d pp_body =
    let loopvar, gensym_exit = Common.Gensym.enter () in
    pp_for_loop ppf
      ( loopvar
      , Expr.Helpers.loop_bottom
      , d
      , pp_block
      , (pp_body, (iteratee, loopvar :: index_ids)) ) ;
    gensym_exit () in
  match dims with
  | [] -> pp_body ppf (iteratee, index_ids)
  | dim :: dims ->
      iter dim (fun ppf (i, idcs) ->
          pf ppf "@[%a @]" pp_block
            (pp_for_loop_iteratee ~index_ids:idcs, (i, dims, pp_body)) )

let emit_name ppf (name, idcs) =
  let name = Mangle.remove_prefix name in
  let to_string = fmt "std::to_string(%s)" in
  pf ppf "param_names__.emplace_back(std::string() + %a);"
    (list ~sep:(fun ppf () -> pf ppf " + '.' + ") string)
    (str "%S" name :: List.map ~f:(str "%a" to_string) idcs)

let emit_complex_name ppf (name, idcs) =
  let name = Mangle.remove_prefix name in
  let to_string = fmt "std::to_string(%s)" in
  pf ppf "@[param_names__.emplace_back(std::string() + %a);@]@,"
    (list ~sep:(fun ppf () -> pf ppf " + '.' + ") string)
    ((str "%S" name :: List.map ~f:(str "%a" to_string) idcs) @ ["\"real\""]) ;
  pf ppf "param_names__.emplace_back(std::string() + %a);"
    (list ~sep:(fun ppf () -> pf ppf " + '.' + ") string)
    ((str "%S" name :: List.map ~f:(str "%a" to_string) idcs) @ ["\"imag\""])

(** Print the [constrained_param_names] method of the model class. *)
let pp_constrained_param_names ppf {Program.output_vars; _} =
  let params =
    [ "std::vector<std::string>& param_names__"
    ; "bool emit_transformed_parameters__ = true"
    ; "bool emit_generated_quantities__ = true" ] in
  let paramvars, tparamvars, gqvars =
    List.partition3_map
      ~f:(function
        | id, {Program.out_block= Parameters; out_constrained_st= st; _} ->
            `Fst (id, st)
        | id, {out_block= TransformedParameters; out_constrained_st= st; _} ->
            `Snd (id, st)
        | id, {out_block= GeneratedQuantities; out_constrained_st= st; _} ->
            `Trd (id, st) )
      output_vars in
  let pp_param_names ppf (decl_id, st) =
    let gen_name =
      if SizedType.contains_complex st then emit_complex_name else emit_name
    in
    let dims = List.rev (SizedType.get_dims st) in
    pp_for_loop_iteratee ppf (decl_id, dims, gen_name) in
  pp_method ppf "void" "constrained_param_names" params nop
    (fun ppf ->
      (list ~sep:cut pp_param_names) ppf paramvars ;
      pf ppf "@,if (emit_transformed_parameters__) %a@," pp_block
        (list ~sep:cut pp_param_names, tparamvars) ;
      pf ppf "@,if (emit_generated_quantities__) %a@," pp_block
        (list ~sep:cut pp_param_names, gqvars) )
    ~cv_attr:["const"; "final"]

(** Print the [unconstrained_param_names] method of the model class.
  This is just a copy of constrained, I need to figure out which one is wrong
   and fix it eventually. From Bob,

   Off the top of my head, I think the four that change sizes
   from constrained to unconstrained are:

   simplex:  K -> (K - 1)
   covar_matrix:  K^2 -> (K choose 2) + K
   corr_matrix:  K^2 -> (K choose 2)
   cholesky_corr: (K choose 2) + K -> (K choose 2)

   cholesky_cov does not change size (it's (K choose 2) + K).
   Now that our unit vector uses the normal thing, that also doesn't
   change size.  The ordered types and constrained types don't
   change sizes either.
*)
let pp_unconstrained_param_names ppf {Program.output_vars; _} =
  let params =
    [ "std::vector<std::string>& param_names__"
    ; "bool emit_transformed_parameters__ = true"
    ; "bool emit_generated_quantities__ = true" ] in
  let paramvars, tparamvars, gqvars =
    List.partition3_map
      ~f:(function
        | id, {Program.out_block= Parameters; out_unconstrained_st= st; _} ->
            `Fst (id, st)
        | id, {out_block= TransformedParameters; out_unconstrained_st= st; _} ->
            `Snd (id, st)
        | id, {out_block= GeneratedQuantities; out_unconstrained_st= st; _} ->
            `Trd (id, st) )
      output_vars in
  let pp_param_names ppf (decl_id, st) =
    let pp_names =
      if SizedType.contains_complex st then emit_complex_name else emit_name
    in
    pp_for_loop_iteratee ppf
      (decl_id, List.rev (SizedType.get_dims st), pp_names) in
  let cv_attr = ["const"; "final"] in
  pp_method ppf "void" "unconstrained_param_names" params nop
    (fun ppf ->
      (list ~sep:cut pp_param_names) ppf paramvars ;
      pf ppf "@,if (emit_transformed_parameters__) %a@," pp_block
        (list ~sep:cut pp_param_names, tparamvars) ;
      pf ppf "@,if (emit_generated_quantities__) %a@," pp_block
        (list ~sep:cut pp_param_names, gqvars) )
    ~cv_attr

(** Print the [transform_inits] method of the model class *)
let pp_transform_inits_impl ppf {Program.transform_inits; _} =
  pf ppf
    "template <typename VecVar, typename VecI, @ \
     stan::require_vector_t<VecVar>* = nullptr, @ \
     stan::require_vector_like_vt<std::is_integral, VecI>* = nullptr> @ " ;
  let params =
    [ "VecVar& params_r__"; "VecI& params_i__"; "VecVar& vars__"
    ; "std::ostream* pstream__ = nullptr" ] in
  let intro ppf () =
    pf ppf "%a" (list ~sep:cut string)
      [ "using local_scalar_t__ = double;"
      ; "stan::io::deserializer<local_scalar_t__> in__(params_r__, params_i__);"
      ; "stan::io::serializer<local_scalar_t__> out__(vars__);"
      ; "int current_statement__ = 0;"
      ; "local_scalar_t__ \
         DUMMY_VAR__(std::numeric_limits<double>::quiet_NaN());" ] in
  let cv_attr = ["const"] in
  pp_method_b ppf "void" "transform_inits_impl" params intro transform_inits
    ~cv_attr

(** Print the [log_prob] method of the model class *)
let pp_log_prob ppf Program.{prog_name; log_prob; _} =
  pf ppf
    "@ template <bool propto__, bool jacobian__ , typename VecR, typename \
     VecI, @ stan::require_vector_like_t<VecR>* = nullptr, @ \
     stan::require_vector_like_vt<std::is_integral, VecI>* = nullptr> @ " ;
  let params =
    ["VecR& params_r__"; "VecI& params_i__"; "std::ostream* pstream__ = nullptr"]
  in
  let intro ppf () =
    pf ppf "%a@ %a@ %a" (list ~sep:cut string)
      [ "using T__ = stan::scalar_type_t<VecR>;"; "using local_scalar_t__ = T__;"
      ; "T__ lp__(0.0);"; "stan::math::accumulator<T__> lp_accum__;"
      ; "stan::io::deserializer<local_scalar_t__> in__(params_r__, params_i__);"
      ; "int current_statement__ = 0;"
      ; "local_scalar_t__ \
         DUMMY_VAR__(std::numeric_limits<double>::quiet_NaN());" ]
      pp_unused "DUMMY_VAR__" pp_function__ (prog_name, "log_prob") in
  let outro ppf () =
    pf ppf "@ lp_accum__.add(lp__);@ return lp_accum__.sum();" in
  let cv_attr = ["const"] in
  pp_method_b ppf "stan::scalar_type_t<VecR>" "log_prob_impl" params intro
    log_prob ~outro ~cv_attr

(** Print the body of the constrained and unconstrained sizedtype methods
 in the model class
 @param ppf A pretty printer
 @param method_name The name of the method to wrap the body in.
 @param outvars The parameters to gather the sizes for.
 *)
let pp_outvar_metadata ppf (method_name, outvars) =
  let json_str = Cpp_Json.out_var_interpolated_json_str outvars in
  let ppbody ppf = pf ppf "@[<hov 2>return std::string(%s);@]@," json_str in
  pp_method ppf "std::string" method_name [] nop ppbody ~cv_attr:["const"]

(** Print the [get_unconstrained_sizedtypes] method of the model class *)
let pp_unconstrained_types ppf {Program.output_vars; _} =
  let grab_unconstrained (name, {Program.out_unconstrained_st; out_block; _}) =
    (name, out_unconstrained_st, out_block) in
  let outvars = List.map ~f:grab_unconstrained output_vars in
  pp_outvar_metadata ppf ("get_unconstrained_sizedtypes", outvars)

(** Print the [get_constrained_sizedtypes] method of the model class *)
let pp_constrained_types ppf {Program.output_vars; _} =
  let grab_constrained (name, {Program.out_constrained_st; out_block; _}) =
    (name, out_constrained_st, out_block) in
  let outvars = List.map ~f:grab_constrained output_vars in
  pp_outvar_metadata ppf ("get_constrained_sizedtypes", outvars)

(** Print the generic method overloads needed in the model class. *)
let pp_overloads ppf {Program.output_vars; _} =
  (* An expression for the number of individual parameters in a list of output variables *)
  let num_outvars (outvars : Expr.Typed.t Program.outvar list) =
    Expr.Helpers.binop_list
      (List.map
         ~f:(fun outvar ->
           SizedType.num_elems_expr outvar.Program.out_constrained_st )
         outvars )
      Operator.Plus ~default:(Expr.Helpers.int 0) in
  (* The list of output variables that came from a particular block *)
  let block_outvars (block : Program.io_block) =
    List.filter_map output_vars
      ~f:(fun ((_ : string), (outvar : Expr.Typed.t Program.outvar)) ->
        if outvar.out_block = block then Some outvar else None ) in
  let num_gen_quantities = num_outvars (block_outvars GeneratedQuantities) in
  let num_params = num_outvars (block_outvars Parameters) in
  let num_transformed = num_outvars (block_outvars TransformedParameters) in
  pf ppf
    {|
    // Begin method overload boilerplate
    template <typename RNG>
    inline void write_array(RNG& base_rng,
                            Eigen::Matrix<double,Eigen::Dynamic,1>& params_r,
                            Eigen::Matrix<double,Eigen::Dynamic,1>& vars,
                            const bool emit_transformed_parameters = true,
                            const bool emit_generated_quantities = true,
                            std::ostream* pstream = nullptr) const {
      const size_t num_params__ = %a;
      const size_t num_transformed = emit_transformed_parameters * %a;
      const size_t num_gen_quantities = emit_generated_quantities * %a;
      const size_t num_to_write = num_params__ + num_transformed +
        num_gen_quantities;
      std::vector<int> params_i;
      vars = Eigen::Matrix<double, Eigen::Dynamic, 1>::Constant(num_to_write,
        std::numeric_limits<double>::quiet_NaN());
      write_array_impl(base_rng, params_r, params_i, vars,
        emit_transformed_parameters, emit_generated_quantities, pstream);
    }

    template <typename RNG>
    inline void write_array(RNG& base_rng, std::vector<double>& params_r,
                            std::vector<int>& params_i,
                            std::vector<double>& vars,
                            bool emit_transformed_parameters = true,
                            bool emit_generated_quantities = true,
                            std::ostream* pstream = nullptr) const {
      const size_t num_params__ = %a;
      const size_t num_transformed = emit_transformed_parameters * %a;
      const size_t num_gen_quantities = emit_generated_quantities * %a;
      const size_t num_to_write = num_params__ + num_transformed +
        num_gen_quantities;
      vars = std::vector<double>(num_to_write,
        std::numeric_limits<double>::quiet_NaN());
      write_array_impl(base_rng, params_r, params_i, vars,
        emit_transformed_parameters, emit_generated_quantities, pstream);
    }

    template <bool propto__, bool jacobian__, typename T_>
    inline T_ log_prob(Eigen::Matrix<T_,Eigen::Dynamic,1>& params_r,
                       std::ostream* pstream = nullptr) const {
      Eigen::Matrix<int, -1, 1> params_i;
      return log_prob_impl<propto__, jacobian__>(params_r, params_i, pstream);
    }

    template <bool propto__, bool jacobian__, typename T__>
    inline T__ log_prob(std::vector<T__>& params_r,
                        std::vector<int>& params_i,
                        std::ostream* pstream = nullptr) const {
      return log_prob_impl<propto__, jacobian__>(params_r, params_i, pstream);
    }


    inline void transform_inits(const stan::io::var_context& context,
                         Eigen::Matrix<double, Eigen::Dynamic, 1>& params_r,
                         std::ostream* pstream = nullptr) const final {
      std::vector<double> params_r_vec(params_r.size());
      std::vector<int> params_i;
      transform_inits(context, params_i, params_r_vec, pstream);
      params_r = Eigen::Map<Eigen::Matrix<double,Eigen::Dynamic,1>>(
        params_r_vec.data(), params_r_vec.size());
    }
|}
    pp_expr num_params pp_expr num_transformed pp_expr num_gen_quantities
    pp_expr num_params pp_expr num_transformed pp_expr num_gen_quantities

(** Print the [get_constrained_sizedtypes] method of the model class *)
let pp_transform_inits ppf {Program.output_vars; _} =
  let params =
    [ "const stan::io::var_context& context"; "std::vector<int>& params_i"
    ; "std::vector<double>& vars"; "std::ostream* pstream__ = nullptr" ] in
  let list_names ((stri : string), (Program.{out_block; _} : 'a Program.outvar))
      =
    match out_block with Parameters -> Some stri | _ -> None in
  let param_names = List.filter_map ~f:list_names output_vars in
  let list_len = List.length param_names in
  let constrained_params =
    List.filter_map ~f:get_constrained_param_st output_vars in
  let get_names ppf () =
    let add_param = fmt "%S" in
    pf ppf "@[<hov 2> constexpr std::array<const char*, %i> names__{%a};@]@,"
      list_len
      (list ~sep:comma add_param)
      (List.map ~f:Mangle.remove_prefix param_names) in
  let get_constrain_param_size_arr ppf () =
    match constrained_params with
    | [] ->
        pf ppf
          "@[<hov 2> const std::array<Eigen::Index, 0> \
           constrain_param_sizes__{};@]@,"
    | _ ->
        let pp_set_params ppf pars = (list ~sep:comma pp_num_param) ppf pars in
        pf ppf
          "@[<hov 2> const std::array<Eigen::Index, %i> \
           constrain_param_sizes__{%a};@]@,"
          list_len pp_set_params constrained_params in
  let get_constrained_param_size ppf () =
    pf ppf
      "@[<hov 2> const auto num_constrained_params__ = std::accumulate(@,\
      \ constrain_param_sizes__.begin(),@,\
       @ constrain_param_sizes__.end(), 0);@]@," in
  let pp_body ppf =
    pf ppf "%a" (list ~sep:cut string)
      [ " std::vector<double> params_r_flat__(num_constrained_params__);"
      ; " Eigen::Index size_iter__ = 0;"; " Eigen::Index flat_iter__ = 0;"
      ; " for (auto&& param_name__ : names__) {"
      ; "   const auto param_vec__ = context.vals_r(param_name__);"
      ; "   for (Eigen::Index i = 0; i < constrain_param_sizes__[size_iter__]; \
         ++i) {"; "     params_r_flat__[flat_iter__] = param_vec__[i];"
      ; "     ++flat_iter__;"; "   }"; "   ++size_iter__;"; " }"
      ; " vars.resize(num_params_r__);"
      ; " transform_inits_impl(params_r_flat__, params_i, vars, pstream__);" ]
  in
  let cv_attr = ["const"] in
  let intro ppf =
    pf ppf "%a %a %a" get_names () get_constrain_param_size_arr ()
      get_constrained_param_size in
  pp_method ppf "void" "transform_inits" params intro
    (fun ppf -> pp_body ppf)
    ~cv_attr

(** Print the public parts of the model class *)
let pp_model_public ppf p =
  pf ppf "@ %a" pp_ctor p ;
  pf ppf "@ %a" pp_log_prob p ;
  pf ppf "@ %a" pp_write_array p ;
  pf ppf "@ %a" pp_transform_inits_impl p ;
  (* Begin metadata methods *)
  pf ppf "@ %a" pp_get_param_names p ;
  (* Post-data metadata methods *)
  pf ppf "@ %a" pp_get_dims p ;
  pf ppf "@ %a" pp_constrained_param_names p ;
  pf ppf "@ %a" pp_unconstrained_param_names p ;
  pf ppf "@ %a" pp_constrained_types p ;
  pf ppf "@ %a" pp_unconstrained_types p ;
  (* Boilerplate *)
  pf ppf "@ %a" pp_overloads p ;
  pf ppf "@ %a" pp_transform_inits p

let model_prefix = "model_"

(** Print the full model class. *)
let pp_model ppf ({Program.prog_name; _} as p) =
  pf ppf "class %s final : public model_base_crtp<%s> {" prog_name prog_name ;
  pf ppf "@ @[<v 1>@ private:@ @[<v 1> %a@]@ " pp_model_private p ;
  pf ppf "@ public:@ @[<v 1> ~%s() { }" prog_name ;
  pf ppf "@ @ inline std::string model_name() const final { return \"%s\"; }"
    prog_name ;
  pf ppf
    {|

  inline std::vector<std::string> model_compile_info() const noexcept {
    return std::vector<std::string>{"stanc_version = %s", "stancflags = %s"};
  }
  |}
    "%%NAME%%3 %%VERSION%%" !stanc_args_to_print ;
  pf ppf "@ %a@]@]@ };" pp_model_public p

let usings =
  {|
using stan::model::model_base_crtp;
using namespace stan::math;
|}

(** Create the model's namespace. *)
let namespace Program.{prog_name; _} = prog_name ^ "_namespace"

(** Find and register functiors used for map_rect. *)
let pp_register_map_rect_functors ppf p =
  let pp_register_functor ppf (i, f) =
    pf ppf "STAN_REGISTER_MAP_RECT(%d, %s::%s)" i (namespace p) f in
  pf ppf "@ %a"
    (list ~sep:cut pp_register_functor)
    (List.sort ~compare (Hashtbl.to_alist map_rect_calls))

let pp_prog ppf (p : Program.Typed.t) =
  Hashtbl.clear Expression_gen.map_rect_calls ;
  (* First, do some transformations on the MIR itself before we begin printing it.*)
  let p, s = Locations.prepare_prog p in
  pf ppf "@[<v>@ %s@ %s@ namespace %s {@ %s@ %a@ %a@ %a@ }@ @]" version includes
    (namespace p) usings Locations.pp_globals s pp_functions_functors p
    (if !standalone_functions then fun _ _ -> () else pp_model)
    p ;
  if !standalone_functions then
    pf ppf "@[<v>%a@ @]"
      (list ~sep:cut (pp_standalone_fun_def (namespace p)))
      p.functions_block
  else (
    pf ppf "@,using stan_model = %s_namespace::%s;@," p.prog_name p.prog_name ;
    pf ppf
      {|
#ifndef USING_R

// Boilerplate
stan::model::model_base& new_model(
        stan::io::var_context& data_context,
        unsigned int seed,
        std::ostream* msg_stream) {
  stan_model* m = new stan_model(data_context, seed, msg_stream);
  return *m;
}

stan::math::profile_map& get_stan_profile_data() {
  return %s_namespace::profiles__;
}

#endif
|}
      p.prog_name ;
    pf ppf "@[<v>%a@]" pp_register_map_rect_functors p )
