From db56cb2e400986a9879f5aa48971e9a6b10cfc15 Mon Sep 17 00:00:00 2001
From: "Julian D. Otalvaro" <juliandavid347@gmail.com>
Date: Thu, 23 Nov 2023 17:00:30 +0000
Subject: [PATCH] Lag re-implemented

---
 examples/bimodal_ke/main.rs        | 11 +++++++----
 src/routines/simulation/predict.rs |  4 ++--
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/examples/bimodal_ke/main.rs b/examples/bimodal_ke/main.rs
index 8a5626b27..7ea4acdb0 100644
--- a/examples/bimodal_ke/main.rs
+++ b/examples/bimodal_ke/main.rs
@@ -52,14 +52,17 @@ struct Ode {}
 impl<'a> Predict<'a> for Ode {
     type Model = Model;
     type State = State;
-    fn initial_system(&self, params: &Vec<f64>, scenario: Scenario) -> Self::Model {
+    fn initial_system(&self, params: &Vec<f64>, scenario: Scenario) -> (Self::Model, Scenario) {
         let params = HashMap::from([("ke".to_string(), params[0]), ("v".to_string(), params[1])]);
-        Model {
+        (Model {
             params,
-            _scenario: scenario,
+            _scenario: scenario.clone(),//TODO remove
             infusions: vec![],
             cov: None,
-        }
+        },
+        scenario.reorder_with_lag(vec![(0.0, 1)])
+    )
+
     }
     fn get_output(&self, x: &Self::State, system: &Self::Model, outeq: usize) -> f64 {
         let v = system.get_param("v");
diff --git a/src/routines/simulation/predict.rs b/src/routines/simulation/predict.rs
index 2567982d9..63d39d0a7 100644
--- a/src/routines/simulation/predict.rs
+++ b/src/routines/simulation/predict.rs
@@ -34,7 +34,7 @@ impl Model {
 pub trait Predict<'a> {
     type Model: 'a + Clone;
     type State;
-    fn initial_system(&self, params: &Vec<f64>, scenario: Scenario) -> Self::Model;
+    fn initial_system(&self, params: &Vec<f64>, scenario: Scenario) -> (Self::Model, Scenario);
     fn initial_state(&self) -> Self::State;
     fn add_covs(&self, system: Self::Model, cov: Option<HashMap<String, CovLine>>) -> Self::Model;
     fn add_infusion(&self, system: Self::Model, infusion: Infusion) -> Self::Model;
@@ -65,7 +65,7 @@ where
         Self { ode }
     }
     pub fn pred(&self, scenario: Scenario, params: Vec<f64>) -> Vec<f64> {
-        let system = self.ode.initial_system(&params, scenario.clone());
+        let (system, scenario) = self.ode.initial_system(&params, scenario.clone());
         let mut yout = vec![];
         let mut x = self.ode.initial_state();
         let mut index: usize = 0;