open Hsim
open Solvers
open Examples
open Common
open Types
open Std.Lift

let sample   = ref 1
let stop     = ref 10.0
let accel    = ref false
let inplace  = ref false
let sundials = ref false
let speed    = ref false
let steps    = ref 1
let model    = ref None
let minstep  = ref None
let maxstep  = ref None
let mintol   = ref None
let maxtol   = ref None
let no_print = ref false
let zelus    = ref false

let gt0i v i  = v := if i <= 0   then 1   else i
let gt0f v f  = v := if f <= 0.0 then 1.0 else f
let opt r s = r := Some s
let modelargs = ref []

let set_model s =
  match !model with
  | None   -> model := Some s
  | Some _ -> modelargs := s :: !modelargs

let opts = [
  "-sample", Arg.Int (gt0i sample), "n \tSample count (default=10)";
  "-stop", Arg.Float (gt0f stop), "n \tStop time (default=10.0)";
  "-debug", Arg.Set Debug.debug, "\tPrint debug information";
  "-accelerate", Arg.Set accel, "\tConcatenate continuous functions";
  "-sundials", Arg.Set sundials, "\tUse sundials (doesn't support -accelerate)";
  "-inplace", Arg.Set inplace, "\tUse imperative solvers";
  "-steps", Arg.Int (gt0i steps), "n \tSplit into [n] steps (default=1)";
  "-speed", Arg.Set speed, "\tLog the step length";
  "-minstep", Arg.String (opt minstep), "\tSet minimum solver step length";
  "-maxstep", Arg.String (opt maxstep), "\tSet maximum solver step length";
  "-mintol",  Arg.String (opt mintol),  "\tSet minimum solver tolerance";
  "-maxtol",  Arg.String (opt maxtol),  "\tSet maximum solver tolerance";
  "-no-print", Arg.Set no_print, "\tDo not print output values";
  "-zelus", Arg.Set zelus, "\tUse the output of the Zélus compiler";
]

let errmsg = "Usage: " ^ Sys.executable_name ^ " [OPTIONS] MODEL\nOptions are:"

let () = try Arg.parse (Arg.align opts) set_model errmsg with _ -> exit 2

let args = List.rev !modelargs

let wrap_zelus (HNode m) =
  let ret = Bigarray.(Array1.create Float64 c_layout 0) in
  let fout s t a y = ignore (m.fout s t a y); ret in
  let step s t () = let _, s = m.step s t () in ret, s in
  HNode { m with fout; step }

let m =
  try
    if !zelus then
      match !model with
      | None            -> Format.eprintf "Missing model\n"; exit 2
      | Some "ballz"    -> wrap_zelus (lift Ballz.main)
      | Some "ballzm"   -> wrap_zelus (lift_hsim Ballz_main.main)
      | Some "sincosz"  -> wrap_zelus (lift Sincosz.f)
      | Some "sincoszm" -> wrap_zelus (lift_hsim Sincosz_main.main)
      (* | Some "count"   -> wrap_zelus (lift Count.count) *)
      | Some s         -> Format.eprintf "Unknown model: %s\n" s; exit 2
    else
      match !model with
      | None          -> Format.eprintf "Missing model\n"; exit 2
      | Some "ball"   -> Ball.init args
      | Some "vdp"    -> Vdp.init args
      | Some "sincos" -> Sincos.init args
      | Some "sqrt"   -> Sqrt.init args
      | Some "sin1x"  -> Sin1x.init args
      | Some "sin1xd" -> Sin1x_der.init args
      | Some s        -> Format.eprintf "Unknown model: %s\n" s; exit 2
  with Invalid_argument s -> Format.eprintf "%s\n" s; exit 2

let st = if !inplace then (module State.InPlaceSimState    : State.SimState)
                     else (module State.FunctionalSimState : State.SimState)

let output =
  if !no_print || !zelus then Hsim.Utils.ignore
  else if !speed then Output.print_h
  else Output.print (* Output.ignore *)

let sim =
  if !sundials then
    let open StatefulSundials in
    let c = if !inplace then InPlace.csolve else Functional.csolve in
    let open StatefulZ in
    let z = if !inplace then InPlace.zsolve else Functional.zsolve in
    let s = Solver.solver c (d_of_dc z) in
    let open Sim.Sim(val st) in
    run_until_n (output !sample (run m s))
  else
    let open StatefulRK45 in
    let c = if !inplace then InPlace.csolve else Functional.csolve in
    let open StatefulZ in
    let z = if !inplace then InPlace.zsolve else Functional.zsolve in
    let s = Solver.solver_c c z in
    let open Sim.Sim(val st) in
    let n = if !accel then accelerate m s else run m (d_of_dc s) in
    run_until_n (output !sample n)

let () = sim !stop !steps ignore