[@@@warning "-27-50-69"]
let todo = assert false

(* Little OCaml reminder *)

type t = { foo : int; bar : int; }

let f () =
  let baz = { foo = 0; bar = 1 } in
  let qux = { baz with foo = 2 } in    (* same as "baz", except field "foo"   *)
  assert (qux = { foo = 2; bar = 1 })








(** Discrete-time node *)
type ('i, 'o, 'r, 's) dnode =
  DNode of {
    state : 's;                        (** current state                      *)
    step  : 's -> 'i -> 's * 'o;       (** step function                      *)
    reset : 's -> 'r -> 's;            (** reset function                     *)
  }


(** Run a discrete node on a list of inputs *)
let drun (DNode n : ('i, 'o, 'r, 's) dnode) (i : 'i list) : 'o list =
  todo






type time =
  float                                (** [≥ 0.0]                            *)


(** Interval-defined functions *)
type 'a dense =
  { h : time;                          (** horizon                            *)
    f : time -> 'a }                   (** [f : [0, h] -> α]                  *)


(** Continuous-time signal *)
type 'a signal =
  'a dense option









(** Initial value problem (IVP) *)
type ('y, 'yder) ivp =
  { y0   : 'y;                         (** initial position                   *)
    fder : time -> 'y -> 'yder;        (** derivative function on [[0.0, h]]  *)
    h    : time; }                     (** maximal horizon                    *) 












(** ODE solver *)
type ('y, 'yder, 's) csolver =
  (time,                               (** requested horizon                  *) 
   'y dense,                           (** solution approximation             *) 
   ('y, 'yder) ivp,                    (** initial value problem              *) 
   's) dnode













(** Zero-crossing problem (ZCP) *)
type ('y, 'zin) zcp =
  { y0   : 'y;                         (** initial position                   *)
    fzer : time -> 'y -> 'zin;         (** zero-crossing function             *)
    h    : time; }                     (** maximal horizon                    *)












(** Zero-crossing solver *)
type ('y, 'zin, 'zout, 's) zsolver =
  ('y dense,                           (** input value                        *)
   time * 'zout,                       (** horizon and zero-crossing events   *)
   ('y, 'zin) zcp,                     (** zero-crossing problem              *)
   's) dnode












(** Full solver (composition of an ODE and zero-crossing solver) *)
type ('y, 'yder, 'zin, 'zout, 's) solver =
  (time,                               (** requested horizon                  *)
   'y dense * 'zout,                   (** output and zero-crossing events    *)
   ('y, 'yder) ivp * ('y, 'zin) zcp,   (** (re)initialization parameters      *)
   's) dnode









(** Compose an ODE solver and a zero-crossing solver *)
let build_solver : ('y, 'yder, 'cs) csolver ->
                   ('y, 'zin, 'zout, 'zs) zsolver ->
                   ('y, 'yder, 'zin, 'zout, 'cs * 'zs) solver
= fun (DNode cs) (DNode zs) ->
  let state = (cs.state, zs.state) in
  let step (cstate, zstate) (h : time) =
    todo in

  
  let reset (cstate, zstate) (ivp, zcp) =
    (cs.reset cstate ivp, zs.reset zstate zcp) in
  DNode { state; step; reset }




(** Hybrid (discrete-time and continuous-time) node *)
type ('i, 'o, 'r, 's, 'y, 'yder, 'zin, 'zout) hnode =
  HNode of {
    state : 's;                        (** current state                      *)
    step  : 's -> 'i -> 's * 'o;       (** discrete step function             *)
    reset : 's -> 'r -> 's;            (** reset function                     *)
    fder  : 's -> 'i -> 'y -> 'yder;   (** derivative function                *)
    fzer  : 's -> 'i -> 'y -> 'zin;    (** zero-crossing function             *)
    fout  : 's -> 'i -> 'y -> 'o;      (** continuous output function         *)
    cget  : 's -> 'y;                  (** continuous state getter            *)
    cset  : 's -> 'y -> 's;            (** continuous state setter            *)
    zset  : 's -> 'zout -> 's;         (** zero-crossing information setter   *)
    jump  : 's -> bool;                (** discrete go-again indicator        *)
  }




(** Simulation mode (either discrete or continuous) *)
type mode = D | C

(** Simulation state *)
type ('i, 'o, 'r, 'y, 'yder, 'zin, 'zout, 'ms, 'ss) state =
  State of {
    solver :                           (** solver state                       *)
      ('y, 'yder, 'zin, 'zout,'ss) solver;
    model :                            (** model state                        *)
      ('i, 'o, 'r, 'ms, 'y, 'yder, 'zin, 'zout) hnode;
    input : 'i signal;                 (** current input                      *)
    time : time;                       (** current time                       *)
    mode : mode;                       (** current step mode                  *)
  }



(** Discrete simulation step *)
let dstep (State ({ model = HNode m; solver = DNode s; _ } as state)) =
  let i = Option.get state.input in
  let (ms, o) = m.step m.state (todo (* current input? *)) in
  let model = HNode { m with state = todo (* ? *) } in
  let state =
    if m.jump ms then State { state with model = todo (* ? *) }
    else if state.time >= i.h then
      State { state with input = todo (* ? *); model; time = todo (* ? *) }
    else
      let y0 = todo (* ? *) and h = i.h -. state.time in
      let ivp = { h; y0; fder = fun t y -> m.fder ms (i.f todo (* ? *)) y } in
      let zcp = { h; y0; fzer = fun t y -> m.fzer ms (i.f todo (* ? *)) y } in
      let solver = DNode { s with state = s.reset s.state (ivp, zcp) } in
      State { state with model; solver; mode = todo (* ? *) } in
  (state, Some { h = 0.; f = fun _ -> o })


  


(** Continuous simulation step *)
let cstep (State ({ model = HNode m; solver = DNode s; _ } as st)) =
  let i = Option.get st.input in
  let (ss, (y, z)) = s.step s.state todo (* ? *) in 
  let ms = m.zset (m.cset m.state (y.f y.h)) z in
  let out = Some { y with f = fun t -> m.fout ms todo (* ? *) (y.f t) } in
  let mode = if m.jump ms || st.time +. y.h >= i.h then D else C in
  let model = HNode { m with state = ms } in
  let solver = DNode { s with state = ss } in
  (State { st with model; solver; mode; time = todo (* ? *) }, out)





(** Simulate a hybrid model with a solver *)
let hsim : ('i, 'o, 'r, 'ms, 'y, 'yder, 'zin, 'zout) hnode ->
           ('y, 'yder, 'zin, 'zout, 'ss) solver ->
           ('i signal, 'o signal, 'r, _) dnode
= fun model solver ->
  let state = State { model; solver; input = None; time = 0.; mode = D } in
  let step (State s as st) input = match (input, s.input, s.mode) with
    | Some _, None,   _ -> dstep (State { s with input; time = 0.; mode = D })
    | None,   Some _, D -> dstep st
    | None,   Some _, C -> cstep st
    | None,   None,   _ -> (st, None)
    | Some _, Some _, _ -> invalid_arg "Not done processing previous input" in
  let reset (State ({ model = HNode m; _ } as s)) r =
    let model = HNode { m with state = m.reset m.state r } in
    State { s with model; input = None; time = 0.; mode = D } in
  DNode { state; step; reset }





(** Run a simulation on a list of inputs *)
let hrun (model : ('i, 'o, 'r, 'ms, 'y, 'yder, 'zin, 'zout) hnode)
         (solver : ('y, 'yder, 'zin, 'zout, 'ss) solver)
         (i : 'i dense list) : 'o dense list
= let sim = hsim model solver and i = List.map Option.some i in
  let rec step os (DNode sim) i =
    let state, o = sim.step sim.state i in
    let sim = DNode { sim with state } in
    if o = None then (sim, List.rev_map Option.get os)
    else step (o :: os) sim None in
  List.fold_left_map (step []) sim i |> snd |> List.flatten