open Hsim
open Solvers
open Common
open Types

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 no_assert = ref false
let c_assert  = 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=1)";
  "-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";
  "-noprint", Arg.Set no_print, "\tDo not print output values";
  "-noassert", Arg.Set no_assert, "\tDo not check assertions";
  "-cassert", Arg.Set c_assert, "\tCheck assertions continuously";
]

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 m =
  try match !model with
  | None          -> Format.eprintf "Missing model\n"; exit 2
  | Some "ball"   -> a_of_h @@ Ball.init args
  | Some "vdp"    -> a_of_h @@ Vdp.init args
  | Some "sincos" -> a_of_h @@ Sincos.init args
  | Some "sqrt"   -> a_of_h @@ Sqrt.init args
  | Some "sin1x"  -> a_of_h @@ Sin1x.init args
  | Some "sin1xd" -> a_of_h @@ Sin1x_der.init args
  | Some "aball"  -> Ball_assert.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 then Hsim.Utils.ignore
  else if !speed then Std.Output.print_h
  else Std.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 (fun () -> d_of_dc (z ())) in
    let open Sim.Sim(val st) in
    let sim = if !no_assert then run (fun () -> h_of_a m) s
              else if !c_assert then run_assert_continuous (fun () -> m) s
              else run_assert_sample !sample (fun () -> m) s in 
    Hsim.Utils.run_until_n (output !sample sim ())
  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 sim =
      if !no_assert then
        if !accel then accelerate (fun () -> h_of_a m) s
        else run (fun () -> h_of_a m) (fun () -> d_of_dc (s ()))
      else
        if !c_assert then run_assert_continuous (fun () -> m) (fun () -> d_of_dc (s ()))
        else run_assert_sample !sample (fun () -> m) (fun () -> d_of_dc (s ())) in
    Hsim.Utils.run_until_n (output !sample sim ())

let () = ignore @@ sim !stop !steps ignore