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
     and der x  = x' init x0
     and der x' = .0 init x'0 *)

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

type state =
  { zin : zarray;
    lx  : carray; (* [h';h;x';x] *)
    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 zsize = 1
let csize = 4

let fder y yd =
  if true (* y.{1} >= 0.0 *) then
    begin yd.{0} <- g; yd.{1} <- y.{0}; yd.{2} <- 0.0; yd.{3} <- y.{2} end
  else begin yd.{0} <- 0.0; yd.{1} <- 0.0; yd.{2} <- 0.0; yd.{3} <- 0.0 end;
  yd
let fzer y zo  = zo.{0} <- -. y.{1}; zo
let fout _ _ y = of_array [| y.{1}; y.{0} |]
let jump _     = true
let horizon _  = max_float
let cget s     = s.lx
let cset s lx  = { s with lx }
let zset s zin = { s with zin }
let step ({ zin; lx; _ } as s) zfalse =
  let lx = if zin.{0} = 1l then
    of_array [| -. 0.8 *. lx.{0}; 0.0; lx.{2}; lx.{3} |] else lx in
  of_array [| s.lx.{1}; s.lx.{0} |], { zin=zfalse; lx; i=false }

let bouncing_ball () : (_, _, _, carray, carray, carray, zarray, carray) hrec =
  let yd         = cmake csize in
  let zout       = cmake zsize in
  let zfalse     = zmake 1 in
  let fder _ _ y = fder y yd in
  let fzer _ _ y = fzer y zout in
  let step s _   = step s zfalse in
  let state      = { zin=zfalse; lx=of_array [|y'0;y0;x'0;x0|]; i=true } in
  let reset _ _  = state in
  { state; fder; fzer; fout; step; reset; horizon; jump; cset; cget; zset;
    csize; zsize }

let errmsg = "Too many arguments for the model (needed: 0)"
let init = function
  | [] -> HNode (bouncing_ball ())
  | _  -> raise (Invalid_argument errmsg)