(* Copyright (C) 2022  Petter A. Urkedal <paurkedal@gmail.com>
 *
 * This library is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or (at your
 * option) any later version, with the LGPL-3.0 Linking Exception.
 *
 * This library is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
 * License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * and the LGPL-3.0 Linking Exception along with this library.  If not, see
 * <http://www.gnu.org/licenses/> and <https://spdx.org>, respectively.
 *)

open Bechamel
open Bechamel.Toolkit

let connect_uri = Uri.of_string "postgresql://?use_single_row_mode=true"

let fetch_many_request =
  let open Caqti_request.Infix in
  let open Caqti_type.Std in
  unit -->* tup4 int int int int @:-
  {|
    WITH tmp (i) AS (VALUES (0), (1), (2), (3), (4), (5), (6), (7), (8), (9))
    SELECT a.i, b.i, c.i, d.i + e.i
    FROM tmp a, tmp b, tmp c, tmp d, tmp e
  |}

let allocate () = Caqti_blocking.(connect connect_uri |> or_fail)
let free (module C : Caqti_blocking.CONNECTION) = C.disconnect ()

let test_fetch_many =
  Staged.stage @@ fun (module C : Caqti_blocking.CONNECTION) ->
  let aux (_, _, _, _) = succ in
  let _count = C.fold fetch_many_request aux () 0 |> Caqti_blocking.or_fail in
  ()

let test =
  let open Test in
  make_with_resource uniq ~name:"fetch-many" ~allocate ~free test_fetch_many

(* Benchmark Function *)

let benchmark () =
  let ols =
    Analyze.ols ~bootstrap:0 ~r_square:true ~predictors:Measure.[|run|]
  in
  let instances =
    Instance.[minor_allocated; major_allocated; monotonic_clock]
  in
  let cfg =
    (* TODO: Why does this segfault with ~kde:(Some 1000)? *)
    Benchmark.cfg ~limit:2000 ~quota:(Time.second 1.0) ()
  in
  let raw_results = Benchmark.all cfg instances test in
  let results =
    List.map (fun instance -> Analyze.all ols instance raw_results) instances
  in
  let results = Analyze.merge ols instances results in
  (results, raw_results)

(* TTY Boilerplate *)

let () =
  List.iter
    (fun v -> Bechamel_notty.Unit.add v (Measure.unit v))
    Instance.[minor_allocated; major_allocated; monotonic_clock]

let img (window, results) =
  Bechamel_notty.Multiple.image_of_ols_results ~rect:window
    ~predictor:Measure.run results

open Notty_unix

let () =
  let window =
    match winsize Unix.stdout with
    | Some (w, h) -> {Bechamel_notty.w; h}
    | None -> {Bechamel_notty.w = 80; h = 1}
  in
  let results, _ = benchmark () in
  img (window, results) |> eol |> output_image