[@@@warning "-27-50-69"]

(* Little OCaml reminder: *)
type _s = A | B of int | C of float * string (* sum types *) 
type _t = { a : int; b : int; }              (* product types *)

let _f () =
  let x = { a = 0; b = 1 } in
  let y = { x with a = 2 } in    (* same as "x", except field "a"   *)
  assert (y = { a = 2; b = 1 })

(* Everything is immutable, except explicitly declared record fields! *)
type _q = { c : int (* immutable *); mutable d : int; }

(* Types can be parameterized by other types: *)
type 'a _llist = Nil | Cons of { v : 'a; mutable next : 'a _llist }




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


(** Run a discrete node on a list of inputs *)
let drun (DNode n : ('i, 'o, 'r) dnode) (i : 'i list) : 'o list =
  snd (List.fold_left_map n.step n.state i)






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                *)
    h    : time; }                     (** maximal horizon                    *) 












(** ODE solver *)
type ('y, 'yder) csolver =
  (time,                               (** requested horizon                  *)
   'y dense,                           (** solution approximation             *)
   ('y, 'yder) ivp)                    (** initial value problem              *)
  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) zsolver =
  ('y dense,                           (** input value                        *)
   time * 'zout,                       (** horizon and zero-crossing events   *)
   ('y, 'zin) zcp)                     (** zero-crossing problem              *)
  dnode












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









(** Compose an ODE solver and a zero-crossing solver. *)
let compose_solvers : ('y, 'yder) csolver ->
                      ('y, 'zin, 'zout) zsolver ->
                      ('y, 'yder, 'zin, 'zout) solver
= fun (DNode csolver) (DNode zsolver) ->
    let state = (csolver.state, zsolver.state) in
    let step (cstate, zstate) h =
      let cstate, y      = csolver.step cstate h in
      let zstate, (h, z) = zsolver.step zstate y in
      (cstate, zstate), ({ y with h }, z) in
    let reset (cstate, zstate) (ivp, zcp) =
      (csolver.reset cstate ivp, zsolver.reset zstate zcp) in
    DNode { state; step; reset }




(** Hybrid (discrete-time and continuous-time) node *)
type ('i, 'o, 'r, 'y, 'yder, 'zin, 'zout) hnode =
  HNode : {
    state : 's;                        (** current state                      *)
    step  : 's -> time -> 'i -> 's * 'o; (** discrete step function           *)
    reset : 's -> 'r -> 's;            (** reset function                     *)
    fder  : 's -> time -> 'i -> 'y -> 'yder; (** derivative function          *)
    fzer  : 's -> time -> 'i -> 'y -> 'zin; (** zero-crossing function        *)
    fout  : 's -> time -> '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 function         *)
  } -> ('i, 'o, 'r, 'y, 'yder, 'zin, 'zout) hnode





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

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




(** 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 state.time (i.f state.time) in
  let model = HNode { m with state = ms } in
  let state = if m.jump ms then
    State { state with model }
  else if state.time >= i.h then
    State { state with input = None; model; time = 0. }
  else
    let y0 = m.cget ms and h = i.h -. state.time and ofs = (+.) state.time in
    let ivp = { h; y0; fder = fun t y -> m.fder ms (ofs t) (i.f (ofs t)) y } in
    let zcp = { h; y0; fzer = fun t y -> m.fzer ms (ofs t) (i.f (ofs t)) y } in
    let solver = DNode { s with state = s.reset s.state (ivp, zcp) } in
    let input = Some { h; f = fun t -> i.f (ofs t) } in
    State { model; solver; mode = C; time = 0.; input } in
  state, Some { h = 0.; f = fun _ -> o }




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




(** Simulate a hybrid model with a solver *)
let hsim : ('i, 'o, 'r, 'y, 'yder, 'zin, 'zout) hnode ->
           ('y, 'yder, 'zin, 'zout) 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, 'y, 'yder, 'zin, 'zout) hnode)
         (solver : ('y, 'yder, 'zin, 'zout) 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