open Hsim.Types
open Solvers.Zls

(* let hybrid bouncing () = (x, y) where
     rec der y  = y'   init y0
     and der y' = -. g init y'0 reset z -> -0.8 *. last y'
     and z = up (-. y)
     and assert (up (-. (y +. epsilon))) *)

let of_array (a : time array) : carray =
  Bigarray.(Array1.of_array Float64 c_layout) a

type state =
  { zin : zarray;
    lx  : carray; (* [| y'; y |] *)
    i   : bool; }

let g = -9.81
let y0 = 50.0
let y'0 = 0.0
let x0 = 0.0
let x'0 = 1.0

let ball ()
: (state, unit, unit, carray, carray, carray, zarray, carray) hrec
= let zsize = 1 in
  let csize = 2 in
  let yd     = cmake csize in
  let zout   = cmake zsize in
  let zfalse = zmake 1 in
  let state = { zin=zfalse; lx=of_array [| y'0; y0 |]; i=true } in
  let fder _ _ () y = yd.{0} <- g; yd.{1} <- y.{0}; yd in
  let fzer _ _ () y = zout.{0} <- -. y.{1}; zout in
  let fout _ _ _ y = y in
  let step s _ () =
    let lx =
      if s.zin.{0} = 1l then of_array [| -. 0.8 *. s.lx.{0}; s.lx.{1} |]
      else s.lx in
    s.lx, { zin=zfalse; lx; i=false } in
  let reset () _ = state in
  let horizon _ = max_float in
  let jump _ = true in
  let cset s lx = { s with lx } in
  let cget s = s.lx in
  let zset s zin = { s with zin } in
  { state; fder; fzer; fout; step; reset; horizon; jump; cset; cget; zset;
    csize; zsize }

type astate = { zin_a: zarray }

let aball epsilon
: (astate, unit, state, bool, carray, carray, zarray, carray) hrec
= let zsize = 1 in
  let csize = 0 in
  let zin_a = zmake zsize in
  let yd   = cmake csize in
  let zout = cmake zsize in
  let state = { zin_a } in
  let fder _ _ _ _ = yd in
  let fzer _ _ st _ = zout.{0} <- -. (st.lx.{1} +. epsilon); zout in
  let fout _ _ st _ = st.lx.{1} +. epsilon >= 0.0 in
  let step { zin_a } _ st =
    zin_a.{0} <> 1l && st.lx.{1} +. epsilon >= 0.0, { zin_a } in
  let reset _ _ = state in
  let horizon _ = max_float in
  let jump _ = true in
  let cset s _ = s in
  let cget s = yd in
  let zset _ zin_a = { zin_a } in
  { state; fder; fzer; fout; step; reset; horizon; jump; cset; cget; zset; csize; zsize }

let errmsg_invalid = "Invalid arguments to model (needed: [float])"
let errmsg_few     = "Too few arguments to model (needed: [float])"
let errmsg_many    = "Too many arguments to model (needed: [float])"
let init = function
  | [eps] ->
      let eps = try float_of_string eps
      with Failure _ -> raise (Invalid_argument errmsg_invalid) in
      let a = HNodeA { body=aball eps; assertions=[] } in
      HNodeA { body=ball (); assertions=[a] }
  | [] -> raise (Invalid_argument errmsg_few)
  | _  -> raise (Invalid_argument errmsg_many)