text: from functools import partial
  
  import jax
  import jax.numpy as jnp
  import flax.linen as fnn
  import optax
  
  import deluca.core
  from deluca.lung.core import LungEnv
  from deluca.lung.utils.data.transform import ShiftScaleTransform
  
  from deluca.lung.utils.sim.nn import SNN
  from deluca.lung.utils.sim.nn import ShallowBoundaryModel
  
  class StitchedSimObservation(deluca.Obj):
      predicted_pressure: float = 0.0
      time: float = 0.0
  
  class SimulatorState(deluca.Obj):
      u_history: jnp.ndarray
      p_history: jnp.ndarray
      t_in: int = 0
      steps: int = 0
      predicted_pressure: float = 0.0
  
  
  class StitchedSim(LungEnv):
      params: list = deluca.field(jaxed=True)
  
      init_rng: jnp.array = deluca.field(jaxed=False)
      u_window: int = deluca.field(5, jaxed=False)
      p_window: int = deluca.field(3, jaxed=False)
      u_history_len: int = deluca.field(5, jaxed=False)
      p_history_len: int = deluca.field(5, jaxed=False)
      u_scaler: ShiftScaleTransform = deluca.field(jaxed=False)
      p_scaler: ShiftScaleTransform = deluca.field(jaxed=False)
      to_round: bool = deluca.field(False, jaxed=False)
      seed: int = deluca.field(0, jaxed=False)
      flow: int = deluca.field(0, jaxed=False)
  
      default_out_dim: int = deluca.field(1, jaxed=False)
      default_hidden_dim: int = deluca.field(100, jaxed=False)
      default_n_layers: int = deluca.field(4, jaxed=False)
      default_dropout_prob: float = deluca.field(0.0, jaxed=False)
  
      num_boundary_models: int = deluca.field(5, jaxed=False)
      boundary_out_dim: int = deluca.field(1, jaxed=False)
      boundary_hidden_dim: int = deluca.field(100, jaxed=False)
      reset_scaled_peep: float = deluca.field(0.0, jaxed=False)
  
      default_model: fnn.module = deluca.field(SNN, jaxed=False)
      boundary_models: list = deluca.field(default_factory=list, jaxed=False)
      ensemble_models: list = deluca.field(default_factory=list, jaxed=False)
  
      def setup(self):
          self.u_history_len = max(self.u_window, self.num_boundary_models)
          self.p_history_len = max(self.p_window, self.num_boundary_models)
  
          self.default_model = SNN(
              out_dim=self.default_out_dim,
              hidden_dim=self.default_hidden_dim,
              n_layers=self.default_n_layers,
          )
          default_params = self.default_model.init(
              jax.random.PRNGKey(0), jnp.ones([self.u_history_len + self.p_history_len])
          )["params"]
  
          self.boundary_models = [
              ShallowBoundaryModel(
                  out_dim=self.boundary_out_dim, hidden_dim=self.boundary_hidden_dim, model_num=i + 1
              )
              for i in range(self.num_boundary_models)
          ]
          boundary_params = [
              self.boundary_models[i].init(
                  jax.random.PRNGKey(0), jnp.ones([self.u_history_len + self.p_history_len])
              )["params"]
              for i in range(self.num_boundary_models)
          ]
  
          self.ensemble_models = self.boundary_models + [self.default_model]
          # print("ENSEMBLE MODELS:")
          # print(self.ensemble_models)
          # without ensemble i.e. flattened params
          if self.params is None:
              self.params = boundary_params + [default_params]
          # print("TREE MAP")
          # print(jax.tree_map(lambda x: x.shape, self.params))
  
      def reset(self):
          scaled_peep = self.reset_scaled_peep
          state = SimulatorState(
              u_history=jnp.zeros([self.u_history_len]),
              p_history=jnp.hstack([jnp.zeros([self.p_history_len]), jnp.array([scaled_peep])]),
              t_in=0,
              steps=0,
              predicted_pressure=scaled_peep,
          )
          state = self.sync_unscaled_pressure(state)
          obs = StitchedSimObservation(predicted_pressure=state.predicted_pressure, time=self.time(state))
          return state, obs
  
      def update_history(self, history, value):
          history = jnp.roll(history, shift=-1)
          history = history.at[-1].set(value)
          return history
  
      def sync_unscaled_pressure(self, state):
          scaled_pressure = self.p_scaler.inverse(state.predicted_pressure).squeeze()
          return state.replace(predicted_pressure=scaled_pressure)
  
      def __call__(self, state, action):
          u_in, u_out = action
          model_idx = jnp.min(jnp.array([state.t_in, self.num_boundary_models]))
          funcs = [partial(self.ensemble_models[i].apply, {"params": self.params[i]}) 
                   for i in range(self.num_boundary_models+1)]
          def true_func(state):
              def call_if_t_in_positive(state):
                  new_state, _ = self.reset()
                  return new_state
  
              state = jax.lax.cond(state.t_in > 0, call_if_t_in_positive, lambda x: x, state)
              return state
  
          def false_func(state, u_in, model_idx):
              if self.to_round:
                  u_in_scaled = self.u_scaler(jnp.round(u_in)).squeeze()
              else:
                  u_in_scaled = self.u_scaler(u_in).squeeze()
              state = state.replace(
                  u_history=self.update_history(state.u_history, u_in_scaled), t_in=state.t_in + 1
              )
              boundary_features = jnp.hstack(
                  [state.u_history[-self.u_history_len :], state.p_history[-self.p_history_len :]]
              )
              default_features = jnp.hstack(
                  [state.u_history[-self.u_window :], state.p_history[-self.p_window :]]
              )
              default_pad_len = (self.u_history_len + self.p_history_len) - (
                  self.u_window + self.p_window
              )
              default_features = jnp.hstack([jnp.zeros((default_pad_len,)), default_features])
              features = jax.lax.cond(
                  model_idx == self.num_boundary_models,
                  lambda x: default_features,
                  lambda x: boundary_features,
                  None,
              )
  
              scaled_pressure = jax.lax.switch(model_idx, funcs, features)[0]
  
              state = state.replace(predicted_pressure=scaled_pressure.squeeze())
              new_p_history = self.update_history(state.p_history, scaled_pressure.squeeze())
              state = state.replace(
                  p_history=new_p_history
              )  # just for inference. This is undone during training
              state = self.sync_unscaled_pressure(state)  # unscale predicted pressure
              return state
  
          partial_false_func = partial(false_func, u_in=u_in, model_idx=model_idx)
          state = jax.lax.cond(u_out == 1, true_func, partial_false_func, state)
          obs = StitchedSimObservation(predicted_pressure=state.predicted_pressure, time=self.time(state))
          return state.replace(steps=state.steps + 1), obs
  
  
  def get_X_y_for_next_epoch_tf(loader, batch_size):
      # loader is a tensorflow.data.Dataset
      loader = loader.shuffle(buffer_size=len(loader))
      loader = loader.batch(batch_size=batch_size, drop_remainder=True)
      loader_list = list(loader)
      unzipped_loader = list(zip(*loader_list))
      X = jnp.array([jnp.array(z.numpy()) for z in unzipped_loader[0]])
      y = jnp.array([jnp.array(z.numpy()) for z in unzipped_loader[1]])
      return X, y
  
  
  def rollout(model, data, num_boundary_models):
      # data.shape == (2, N)
      start_idx = 0
      end_idx = len(data[0])
      state, _ = model.reset()
      new_u_history = jnp.zeros((num_boundary_models,))
      new_p_history = jnp.zeros((num_boundary_models,))
      state = state.replace(u_history=new_u_history, p_history=new_p_history)
      state_loss_init = (state, 0.0)
  
      def predict_and_update_state(i, state_loss):
          state, loss = state_loss
          u_in, pressure = data[0, i], data[1, i]  # unscaled u_in and pressure
          new_p_history = jnp.append(state.p_history, model.p_scaler(pressure).squeeze())
          state = state.replace(p_history=new_p_history)
          next_state, _ = model(state=state, action=(u_in, 0))
          # drop predicted_pressure so that u_history.shape = u_history_len+1, p_history.shape = p_history_len+1
          truncated_p_history = next_state.p_history[:-1]
          next_state = next_state.replace(p_history=truncated_p_history)
          pred = model.p_scaler(next_state.predicted_pressure)  # scaled gradient step
          return (
              next_state,
              loss + jnp.abs(model.p_scaler(data[1, i + 1]) - pred),
          )  # scaled gradient step
  
      (state, total_loss) = jax.lax.fori_loop(
          start_idx, end_idx, predict_and_update_state, state_loss_init
      )
      return total_loss / (end_idx - start_idx)
  
  
  def loop_over_loader(model_optimState, X_Y, optim, rollout):
      """
      rollout has signature (model, data) -> loss where data.shape = (2, N)
      X_batch.shape = Y_batch.shape = (num_batches, batch_size, N=29)
      """
      # print('=================== ENTER loop_over_loader ======================')
      X_batch, y_batch = X_Y
      model, optim_state = model_optimState
      value, grad = jax.value_and_grad(map_rollout_over_batch)(model, (X_batch, y_batch), rollout)
      updates, optim_state = optim.update(grad, optim_state, model)
      model = optax.apply_updates(model, updates)
      return (model, optim_state), None
  
  
  def map_rollout_over_batch(model, data, rollout):
      """
      rollout has signature (model, data) -> loss where data.shape = (2, N)
      data.shape = ((batch_size, N), (batch_size, N))
      """
      rollout_partial = lambda xs: partial(rollout, model=model)(data=xs)
      data_zipped = jnp.array(list(zip(data[0], data[1])))  # (batch_size, 2, N)
      losses = jax.vmap(rollout_partial)(data_zipped)
      return jnp.array(losses).mean()
  
  
  # TODO: add scheduler and scheduler_params
  def stitched_sim_train(
      munger,
      model,
      num_boundary_models,
      # 0 to num_boundary_models-1 are boundary models, and num_boundary_models is #
      # default_model
      train_key="train",
      test_key="test",
      batch_size=512,
      epochs=500,
      optimizer=optax.adam,
      optimizer_params={"learning_rate": 1e-3},
      # scheduler=torch.optim.lr_scheduler.ReduceLROnPlateau,
      # scheduler_params={"factor": 0.9, "patience": 15},
      loss_fn=lambda x, y: (jnp.abs(x - y)).mean(),
      # loss_fn_params={},
      print_loss=100,
  ):
      # Unpack dataloader
      loader = munger.get_dataloader_tf(
          key=train_key, shuffle=True, scale_data=False, batch_size=batch_size
      )
  
      # evaluate on these at end of epoch
      X_train, y_train = munger.get_scaled_data_jax(train_key, scale_data=False)
      X_test, y_test = munger.get_scaled_data_jax(test_key, scale_data=False)
  
      # set up model and optimizer
      optim = optimizer(**optimizer_params)
      optim_state = optim.init(model)
  
      rollout_partial = partial(rollout, num_boundary_models=num_boundary_models)
      loop_over_loader_partial = partial(loop_over_loader, optim=optim, rollout=rollout_partial)
      for epoch in range(epochs + 1):
          if epoch % 25 == 0:
              print("epoch:" + str(epoch))
          X, y = get_X_y_for_next_epoch_tf(loader, batch_size)
  
          (model, optim_state), _ = jax.lax.scan(
              loop_over_loader_partial, (model, optim_state), (X, y)
          )
  
          if epoch % print_loss == 0:
              # expensive end-of-epoch eval, just for intuition
              print("X_train.shape:" + str(X_train.shape))
              print("y_train.shape:" + str(y_train.shape))
              train_loss = map_rollout_over_batch(model, (X_train, y_train), rollout_partial)
              # cross-validation
              print("X_test.shape:" + str(X_test.shape))
              print("y_test.shape:" + str(y_test.shape))
              test_loss = map_rollout_over_batch(model, (X_test, y_test), rollout_partial)
              print(f"Epoch {epoch:2d}: train={train_loss.item():.5f}, test={test_loss.item():.5f}")
              print("-----------------------------------")
      print("finished looping over epochs")
      return model
  
  