open Hsim.Types
open Solvers.Zls

let of_array a = Bigarray.Array1.of_array Bigarray.Float64 Bigarray.c_layout a

type state = { si : bool; sx : carray }

let csize = 3
let zsize = 1

let fder y yd omega =
  yd.{0} <- omega *. y.{1}; yd.{1} <- -.omega *. y.{0}; yd.{2} <- 1.0; yd
let fout _ _ _ y = of_array [| y.{0}; y.{1}; y.{2} |]
let step { si; sx } sin0 cos0 =
  let sx = if si then of_array [| sin0; cos0; 0.0 |] else sx in
  of_array [| sx.{0}; sx.{1}; sx.{2} |], { sx; si=false }
let cget s     = s.sx
let cset s lx  = { s with sx=lx }
let zset s _   = s
let jump _     = true
let horizon _  = max_float

let sinus_cosinus theta0 omega =
  let sin0 = Float.sin theta0 in
  let cos0 = Float.cos theta0 in
  let yd   = cmake csize in
  let zout = cmake zsize in
  let fder _ _ _ y = fder y yd omega in
  let fzer _ _ _ _ = zout in
  let step s _ _   = step s sin0 cos0 in
  let state      = { sx=of_array [| sin0; cos0; 0.0 |]; si=true } in
  let reset _ _  = state in
  HNode { state; fder; fzer; fout; step; reset; horizon;
          jump; cset; cget; zset; csize; zsize }

let errmsg_invalid = "Invalid arguments to model (needed: [float, float])"
let errmsg_few     = "Too few arguments to model (needed: [float, float])"
let errmsg_many    = "Too many arguments to model (needed: [float, float])"
let init = function
  | [om; t0] ->
      let t0, om = try float_of_string t0, float_of_string om
      with Failure _ -> raise (Invalid_argument errmsg_invalid) in
      sinus_cosinus t0 om
  | [] | [_] -> raise (Invalid_argument errmsg_few)
  | _        -> raise (Invalid_argument errmsg_many)