open Hsim.Types
open Hsim.Solver
open Zls

module Functional =
  struct
    type ('state, 'vec) state = { state: 'state; vec: 'vec }

    let csolve : (carray, carray) csolver_c =
      let open Odexx.Ode45 in

      let state =
        let v = Zls.cmake 0 in
        let state = initialize (fun _ _ _ -> ()) (vec v) in
        set_stop_time state 1.0; { state; vec=v } in

      let reset { fder; init; stop; _ } _ =
        let fder t cvec dvec = Zls.blit (fder t cvec) dvec in
        let state = initialize fder (vec init) in
        set_stop_time state stop;
        { state; vec = init } in

      let step ({ state ; vec=v } as s) h =
        let y_nv = vec v in
        let h = step state h y_nv in
        let state = copy state in
        let dky t = get_dky state y_nv t 0; unvec y_nv in
        (h, dky), s in

      let copy { state; vec } = { state; vec } in

      DNodeC { state; step; reset; copy }
  end

module InPlace =
  struct
    type ('state, 'vec) state = { mutable state: 'state; mutable vec : 'vec }

    let csolve : (carray, carray) csolver_c =
      let open Odexx.Ode45 in

      let state =
        let v = Zls.cmake 0 in
        let state = initialize (fun _ _ _ -> ()) (vec v) in
        set_stop_time state 1.0;
        { state; vec=v } in

      let reset { fder; init; stop; _ } s =
        let fder t cvec dvec =
          let dvec' = fder t cvec in Zls.blit dvec' dvec in
        let state = initialize fder (vec init) in
        set_stop_time state stop; s.state <- state ; s.vec <- init; s in

      let step ({ state; vec=v } as s) h =
        let y_nv = vec v in
        let h = step state h y_nv in
        let get_dky t = get_dky state y_nv t 0; unvec y_nv in
        (h, get_dky), s in

      let copy { state; vec } =
        { state = copy state; vec = Zls.copy vec } in

      DNodeC { state; reset; step; copy }
  end