open Hsim.Types
open Solvers

(* 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 : Zls.zarray;
    lx : Zls.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 bouncing_ball () =
  let zfalse = Zls.zmake 1 in
  let fder _ (y: Zls.carray) (yd: Zls.carray) =
    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 in
  let fzer _ (y: Zls.carray) (zout: Zls.carray) = zout.{0} <- -. y.{1} in
  let fout _ (y: Zls.carray): Zls.carray = of_array [| y.{1} |] in
  let step ({ zin; lx; _ } as s) _ : Zls.carray * state =
    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} |], { zin=zfalse; lx; i=false } in
  let cget s = s.lx in
  let cset s lx = { s with lx } in
  let zset s zin = { s with zin } in
  let yd    = Zls.cmake 4 in
  let zout  = Zls.cmake 1 in
  let state = { zin=zfalse; lx=of_array [|y'0;y0;x'0;x0|]; i=true } in
  let reset _ _ = state in
  let jump _ = true in
  let zsize = 1 in
  HNode
  { state = state;
    fder = (fun _ _ y -> fder 0.0 y yd; yd);
    fzer = (fun _ _ y -> fzer 0.0 y zout; zout);
    fout = (fun s _ y -> fout s y);
    step;
    horizon = (fun _ -> max_float);
    cset; cget; zset; reset; jump; zsize }

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