open Hsim.Types
open Solvers.Zls

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

type s = Good | Bad

let with_nan = false

type state =
  { s_auto: s; (* active state of the automaton *)
    s_pos : carray;
    s_zin : zarray;
    s_encore : bool }

let sqrt () =
  let zfalse = zmake 1 in
  let fder s y yd =
    yd.{0} <- -1.0;
    match s.s_auto with
    | Good ->
        let o = if with_nan then sqrt y.{0}
                else if y.{0} >= 0.0 then sqrt y.{0} else 0.0 in
        yd.{1} <- o
    | Bad ->
        yd.{1} <- 0.0 in
  let fzero _ y zout = zout.{0} <- -. y.{0} in
  let fout state y =
    let o =
      match state.s_auto with
      | Good -> let o =
                  if with_nan then sqrt y.{0}
                  else if y.{0} >= 0.0 then sqrt y.{0} else 0.0 in
                o
      | Bad -> let o = 42.0 in o in
    of_array [| o; state.s_pos.{0}; state.s_pos.{1} |] in
  let fstep ({ s_auto; s_pos; s_zin; _ } as state) _ =
    match s_auto with
      | Good ->
         let o = if with_nan then sqrt s_pos.{0}
                 else if s_pos.{0} >= 0.0 then sqrt s_pos.{0} else 0.0 in
         let state =
           if s_zin.{0} = 1l then { state with s_auto=Bad; s_encore=true }
           else state in
         let pos = of_array [| state.s_pos.{0}; state.s_pos.{1} |] in
         of_array [| o; state.s_pos.{0}; state.s_pos.{1} |],
         { state with s_zin=zfalse; s_pos=pos }
    | Bad ->
       let o = 42.0 in
       let state = { state with s_encore=false;
                                s_pos=of_array [| s_pos.{0}; 0.0 |] } in
       of_array [| o; state.s_pos.{0}; state.s_pos.{1} |], state in
  let cget { s_pos; _ } = s_pos in
  let cset s l_x = { s with s_pos=l_x } in
  let zset s zin = { s with s_zin=zin } in
  let yd = cmake 2 in
  let zout = cmake 1 in
  let zsize = 1 in
  let csize = 2 in
  let s_init =
    { s_encore = false;
      s_auto = Good;
      s_pos = of_array [| 13.3; 0.0 |];
      s_zin = zmake 1 } in
  let reset _ _ = s_init in
  let jump _ = true in
  HNode { init = (fun _ -> s_init);
    fder = (fun s _ y -> fder s y yd; yd);
    fzer = (fun s _ y -> fzero s y zout; zout);
    fout = (fun s _ y -> fout s y);
    step = (fun s a -> fstep s a);
    horizon = (fun s -> if s.s_encore then 0.0 else max_float);
    cset; cget; zset; zsize; csize; reset; jump }

let errmsg = "Too many arguments to model (needed: 0)"
let init = function
  | [] -> sqrt ()
  | _  -> raise (Invalid_argument errmsg)