(* gen_bench.ml — Generate Rocq test files for benchmarking deriving.

   Usage:
     ocaml gen_bench.ml [options]

   Options:
     -types N        number of mutually inductive types (default 1)
     -ctors N        constructors per type (default 3)
     -rec N          recursive arguments per constructor (default 1)
     -nonrec N       nonrecursive arguments per constructor (default 1)
     -nonrec-type T  type of nonrecursive arguments (default nat)
     -seed N         random seed (default 0)
     -o FILE         output file (default stdout)

   The first constructor of each type always takes 0 recursive arguments
   so the type is inhabited. *)

let n_types = ref 1
let n_ctors = ref 3
let n_rec = ref 1
let n_nonrec = ref 1
let nonrec_type = ref "nat"
let seed = ref 0
let output = ref ""

let speclist = [
  "-types",       Arg.Set_int n_types,       "N number of mutually inductive types";
  "-ctors",       Arg.Set_int n_ctors,       "N constructors per type";
  "-rec",         Arg.Set_int n_rec,         "N recursive arguments per constructor";
  "-nonrec",      Arg.Set_int n_nonrec,      "N nonrecursive arguments per constructor";
  "-nonrec-type", Arg.Set_string nonrec_type,"T type of nonrecursive arguments";
  "-seed",        Arg.Set_int seed,          "N random seed";
  "-o",           Arg.Set_string output,     "FILE output file (default stdout)";
]

let usage = "gen_bench.ml [options]"

let () =
  Arg.parse speclist (fun a -> raise (Arg.Bad ("unexpected arg " ^ a))) usage;
  if !n_types < 1 then begin
    prerr_endline "gen_bench: -types must be >= 1"; exit 1
  end;
  Random.init !seed;

  let buf = Buffer.create 4096 in
  let pf fmt = Printf.ksprintf (Buffer.add_string buf) fmt in

  let type_name i = Printf.sprintf "T%d" i in
  let ctor_name i j = Printf.sprintf "C%d_%d" i j in
  let mutual = !n_types > 1 in

  (* Header *)
  pf "(* Auto-generated by bench/gen_bench.ml. Do not edit. *)\n";
  pf "(* types=%d ctors=%d rec=%d nonrec=%d nonrec-type=%s seed=%d *)\n\n"
    !n_types !n_ctors !n_rec !n_nonrec !nonrec_type !seed;
  pf "From HB Require Import structures.\n";
  pf "From mathcomp Require Import\n";
  pf "  ssreflect ssrfun ssrbool ssrnat eqtype seq choice fintype finset order.\n";
  pf "Set SsrOldRewriteGoalsOrder.\n\n";
  pf "From deriving Require Import deriving.\n";
  pf "From deriving Require Import instances.eqtype instances.order.\n\n";
  pf "Set Implicit Arguments.\n";
  pf "Unset Strict Implicit.\n";
  pf "Unset Printing Implicit Defensive.\n\n";

  (* Inductive declaration *)
  if mutual then pf "Unset Elimination Schemes.\n";
  for i = 0 to !n_types - 1 do
    let kw = if i = 0 then "Inductive" else "with" in
    pf "%s %s :=" kw (type_name i);
    if !n_ctors = 0 then ()
    else begin
      pf "\n";
      for j = 0 to !n_ctors - 1 do
        let n_rec_this = if j = 0 then 0 else !n_rec in
        let nonrec_args = List.init !n_nonrec (fun _ -> !nonrec_type) in
        let rec_args =
          List.init n_rec_this (fun _ -> type_name (Random.int !n_types))
        in
        let args = nonrec_args @ rec_args in
        pf "| %s" (ctor_name i j);
        (match args with
         | [] -> ()
         | _  -> pf " of %s" (String.concat " & " args));
        if j < !n_ctors - 1 then pf "\n"
      done
    end;
    if i < !n_types - 1 then pf "\n"
  done;
  pf ".\n";
  if mutual then pf "Set Elimination Schemes.\n";
  pf "\n";

  (* Recursors *)
  let rect_name i = Printf.sprintf "%s_rect" (type_name i) in
  let combined_name = "T_combined_rect" in
  let main_rect =
    if mutual then combined_name else rect_name 0
  in
  if mutual then begin
    for i = 0 to !n_types - 1 do
      let kw = if i = 0 then "Scheme" else "with " in
      pf "%s %s := Induction for %s Sort Type%s\n"
        kw (rect_name i) (type_name i)
        (if i = !n_types - 1 then "." else "")
    done;
    pf "Combined Scheme %s from" combined_name;
    for i = 0 to !n_types - 1 do
      pf "%s %s" (if i = 0 then "" else ",") (rect_name i)
    done;
    pf ".\n\n"
  end;

  let say msg =
    pf "Goal True. idtac \"BENCH: %s\". Abort.\n" msg
  in

  (* indDef + IndType *)
  say "indDef";
  pf "Time Definition T_indDef := [indDef for %s].\n" main_rect;
  for i = 0 to !n_types - 1 do
    let t = type_name i in
    say (Printf.sprintf "IndType %s" t);
    pf "Time Canonical %s_indType := IndType %s T_indDef.\n" t t
  done;
  pf "\n";

  (* nored hasDecEq + simplified hasDecEq for each type *)
  for i = 0 to !n_types - 1 do
    let t = type_name i in
    say (Printf.sprintf "nored hasDecEq %s" t);
    pf "Time Definition %s_hasDecEq_nored : hasDecEq %s :=\n" t t;
    pf "  [derive nored hasDecEq for %s].\n" t;
    say (Printf.sprintf "simpl hasDecEq %s" t);
    pf "Time Definition %s_hasDecEq : hasDecEq %s :=\n" t t;
    pf "  ltac:(simpl_hasDecEq %s %s_hasDecEq_nored).\n" t t;
    pf "HB.instance Definition _ := %s_hasDecEq.\n\n" t
  done;

  (* hasChoice for each type. There is no nored form: the derivation
     piggybacks on tree_of_coq_indK via Choice.copy/pcan_type. *)
  for i = 0 to !n_types - 1 do
    let t = type_name i in
    say (Printf.sprintf "hasChoice %s" t);
    pf "Time Definition %s_hasChoice := [derive hasChoice for %s].\n" t t;
    pf "HB.instance Definition _ := %s_hasChoice.\n\n" t
  done;

  (* isCountable, same shape as hasChoice. *)
  for i = 0 to !n_types - 1 do
    let t = type_name i in
    say (Printf.sprintf "isCountable %s" t);
    pf "Time Definition %s_isCountable := [derive isCountable for %s].\n" t t;
    pf "HB.instance Definition _ := %s_isCountable.\n\n" t
  done;

  (* nored isOrder + simplified isOrder, mirroring hasDecEq. *)
  for i = 0 to !n_types - 1 do
    let t = type_name i in
    say (Printf.sprintf "nored isOrder %s" t);
    pf "Time Definition %s_isOrder_nored\n" t;
    pf "  : Order.isOrder Order.default_display %s :=\n" t;
    pf "  [derive nored isOrder for %s].\n" t;
    say (Printf.sprintf "simpl isOrder %s" t);
    pf "Time Definition %s_isOrder\n" t;
    pf "  : Order.isOrder Order.default_display %s :=\n" t;
    pf "  ltac:(simpl_isOrder %s %s_isOrder_nored).\n" t t;
    pf "HB.instance Definition _ := %s_isOrder.\n\n" t
  done;

  let out = Buffer.contents buf in
  if !output = "" then print_string out
  else begin
    let oc = open_out !output in
    output_string oc out;
    close_out oc
  end
