class BangBang(Controller):
  """bang bang controller."""
  waveform: BreathWaveform = deluca.field(jaxed=False)
  min_action: float = deluca.field(0.0, jaxed=False)
  max_action: float = deluca.field(100.0, jaxed=False)

  def setup(self):
    if self.waveform is None:
      self.waveform = BreathWaveform.create()

  def __call__(self, controller_state, obs):
    pressure, t = obs.predicted_pressure, obs.time
    target = self.waveform.at(t)
    action = jax.lax.cond(pressure < target, lambda x: self.max_action,
                          lambda x: self.min_action, None)
    # update controller_state
    new_dt = jnp.max(
        jnp.array([DEFAULT_DT, t - proper_time(controller_state.time)]))
    new_time = t
    new_steps = controller_state.steps + 1
    controller_state = controller_state.replace(
        time=new_time, steps=new_steps, dt=new_dt)
    return controller_state, action class DeepCnn(Controller):
  """deep cnn controller."""
  params: list = deluca.field(jaxed=True)
  model: nn.module = deluca.field(DeepNetwork, jaxed=False)
  featurizer: jnp.array = deluca.field(jaxed=False)
  H: int = deluca.field(100, jaxed=False)
  input_dim: int = deluca.field(1, jaxed=False)
  history_len: int = deluca.field(10, jaxed=False)
  kernel_size: int = deluca.field(5, jaxed=False)
  clip: float = deluca.field(100.0, jaxed=False)
  normalize: bool = deluca.field(False, jaxed=False)
  u_normalizer: ShiftScaleTransform = deluca.field(jaxed=False)
  p_normalizer: ShiftScaleTransform = deluca.field(jaxed=False)
  model_apply: collections.abc.Callable = deluca.field(jaxed=False)
  use_model_apply_jit: bool = deluca.field(False, jaxed=False)
  use_leaky_clamp: bool = deluca.field(True, jaxed=False)

  def setup(self):
    self.model = DeepNetwork(H=self.H, kernel_size=self.kernel_size)
    if self.params is None:
      self.params = self.model.init(
          jax.random.PRNGKey(0),
          jnp.expand_dims(jnp.ones([self.history_len]), axis=(0, 2)))["params"]
    # linear feature transform:
    # errs -> [average of last h errs, ..., average of last 2 errs, last err]
    # emulates low-pass filter bank
    self.featurizer = jnp.tril(jnp.ones((self.history_len, self.history_len)))
    self.featurizer /= jnp.expand_dims(
        jnp.arange(self.history_len, 0, -1), axis=0)
    # TODO(dsuo): resolve jit load/unload
    if self.use_model_apply_jit:
      self.model_apply = jax.jit(self.model.apply)
    else:
      self.model_apply = self.model.apply

  # TODO(dsuo): Handle dataclass initialization of jax objects
  def init(self, waveform=None):
    if waveform is None:
      waveform = BreathWaveform.create()
    errs = jnp.array([0.0] * self.history_len)
    state = DeepControllerState(errs=errs, waveform=waveform)
    return state

  @jax.jit
  def __call__(self, controller_state, obs):
    state, t = obs.predicted_pressure, obs.time
    errs, waveform = controller_state.errs, controller_state.waveform
    target = waveform.at(t)
    if self.normalize:
      target_normalized = self.p_normalizer(target).squeeze()
      state_normalized = self.p_normalizer(state).squeeze()
      next_errs = jnp.roll(errs, shift=-1)
      next_errs = next_errs.at[-1].set(target_normalized - state_normalized)
    else:
      next_errs = jnp.roll(errs, shift=-1)
      next_errs = next_errs.at[-1].set(target - state)
    controller_state = controller_state.replace(errs=next_errs)
    decay = self.decay(waveform, t)

    def true_func(null_arg):
      trajectory = jnp.expand_dims(next_errs[-self.history_len:], axis=(0, 1))
      input_val = jnp.reshape((trajectory @ self.featurizer),
                              (1, self.history_len, 1))
      u_in = self.model_apply({"params": self.params}, input_val)
      return u_in.squeeze().astype(jnp.float32)

    # changed decay compare from None to float(inf) due to cond requirements
    u_in = jax.lax.cond(
        jnp.isinf(decay), true_func, lambda x: jnp.array(decay), None)
    # Implementing "leaky" clamp to solve the zero gradient problem
    if self.use_leaky_clamp:
      u_in = jax.lax.cond(u_in < 0.0, lambda x: x * 0.01, lambda x: x, u_in)
      u_in = jax.lax.cond(u_in > self.clip, lambda x: self.clip + x * 0.01,
                          lambda x: x, u_in)
    else:
      u_in = jax.lax.clamp(0.0, u_in.astype(jnp.float32), self.clip).squeeze()
    # update controller_state
    new_dt = jnp.max(
        jnp.array([DEFAULT_DT, t - proper_time(controller_state.time)]))
    new_time = t
    new_steps = controller_state.steps + 1
    controller_state = controller_state.replace(
        time=new_time, steps=new_steps, dt=new_dt)
    return controller_state, u_in class DeepMlp(Controller):
  """Deep mlp controller."""
  params: list = deluca.field(jaxed=True)
  model: nn.module = deluca.field(MLP, jaxed=False)
  hidden_dim: int = deluca.field(100, jaxed=False)
  n_layers: int = deluca.field(2, jaxed=False)
  droprate: float = deluca.field(0.0, jaxed=False)
  activation_fn_name: str = deluca.field("leaky_relu", jaxed=False)
  back_history_len: int = deluca.field(10, jaxed=False)
  fwd_history_len: int = deluca.field(2, jaxed=False)
  horizon: int = deluca.field(29, jaxed=False)
  clip: float = deluca.field(100.0, jaxed=False)
  normalize: bool = deluca.field(False, jaxed=False)
  u_normalizer: ShiftScaleTransform = deluca.field(jaxed=False)
  p_normalizer: ShiftScaleTransform = deluca.field(jaxed=False)
  model_apply: collections.abc.Callable = deluca.field(jaxed=False)
  use_model_apply_jit: bool = deluca.field(False, jaxed=False)

  def setup(self):
    if self.activation_fn_name == "leaky_relu":
      activation_fn = nn.leaky_relu
    self.model = MLP(
        hidden_dim=self.hidden_dim,
        out_dim=1,
        n_layers=self.n_layers,
        droprate=self.droprate,
        activation_fn=activation_fn)
    if self.params is None:
      self.params = self.model.init(
          jax.random.PRNGKey(0),
          jnp.ones([self.back_history_len + self.fwd_history_len]))["params"]
    if self.use_model_apply_jit:
      self.model_apply = jax.jit(self.model.apply)
    else:
      self.model_apply = self.model.apply
    # linear feature transform:
    # errs -> [average of last h errs, ..., average of last 2 errs, last err]
    # emulates low-pass filter bank

  # TODO(dsuo): Handle dataclass initialization of jax objects
  def init(self, waveform=None):
    if waveform is None:
      waveform = BreathWaveform.create()
    errs = jnp.array([0.0] * self.back_history_len)
    fwd_targets = jnp.array(
        [waveform.at(t * DEFAULT_DT) for t in range(self.fwd_history_len)])
    state = DeepControllerState(
        errs=errs, fwd_targets=fwd_targets, waveform=waveform)
    return state

  @jax.jit
  def __call__(self, controller_state, obs):
    state, t = obs.predicted_pressure, obs.time
    errs, waveform = controller_state.errs, controller_state.waveform
    fwd_targets = controller_state.fwd_targets
    target = waveform.at(t)
    fwd_t = t + self.fwd_history_len * DEFAULT_DT
    if self.fwd_history_len > 0:
      fwd_target = jax.lax.cond(fwd_t >= self.horizon * DEFAULT_DT,
                                lambda x: fwd_targets[-1],
                                lambda x: waveform.at(fwd_t), None)
    if self.normalize:
      target_normalized = self.p_normalizer(target).squeeze()
      state_normalized = self.p_normalizer(state).squeeze()
      next_errs = jnp.roll(errs, shift=-1)
      next_errs = next_errs.at[-1].set(target_normalized - state_normalized)
      if self.fwd_history_len > 0:
        fwd_target_normalized = self.p_normalizer(fwd_target).squeeze()
        next_fwd_targets = jnp.roll(fwd_targets, shift=-1)
        next_fwd_targets = next_fwd_targets.at[-1].set(fwd_target_normalized)
      else:
        next_fwd_targets = jnp.array([])
    else:
      next_errs = jnp.roll(errs, shift=-1)
      next_errs = next_errs.at[-1].set(target - state)
      if self.fwd_history_len > 0:
        next_fwd_targets = jnp.roll(fwd_targets, shift=-1)
        next_fwd_targets = next_fwd_targets.at[-1].set(fwd_target)
      else:
        next_fwd_targets = jnp.array([])
    controller_state = controller_state.replace(
        errs=next_errs, fwd_targets=next_fwd_targets)
    decay = self.decay(waveform, t)

    def true_func(null_arg):
      trajectory = jnp.hstack([next_errs, next_fwd_targets])
      u_in = self.model_apply({"params": self.params}, trajectory)
      return u_in.squeeze().astype(jnp.float64)

    # changed decay compare from None to float(inf) due to cond requirements
    u_in = jax.lax.cond(
        jnp.isinf(decay), true_func, lambda x: jnp.array(decay), None)
    u_in = jax.lax.clamp(0.0, u_in.astype(jnp.float64), self.clip).squeeze()
    # update controller_state
    new_dt = jnp.max(
        jnp.array([DEFAULT_DT, t - proper_time(controller_state.time)]))
    new_time = t
    new_steps = controller_state.steps + 1
    controller_state = controller_state.replace(
        time=new_time, steps=new_steps, dt=new_dt)
    return controller_state, u_in class PID(Controller):
  """PID controller."""
  model: nn.module = deluca.field(PIDNetwork, jaxed=False)
  params: jnp.array = deluca.field(jaxed=True)  # jnp.array([3.0, 4.0, 0.0]
  RC: float = deluca.field(0.5, jaxed=False)
  dt: float = deluca.field(0.03, jaxed=False)
  model_apply: collections.abc.Callable = deluca.field(jaxed=False)

  def setup(self):
    self.model = PIDNetwork()
    if self.params is None:
      self.params = self.model.init(jax.random.PRNGKey(0), jnp.ones([
          3,
      ]))["params"]
    # TODO(dsuo): Handle dataclass initialization of jax objects
    self.model_apply = jax.jit(self.model.apply)

  def init(self, waveform=None):
    if waveform is None:
      waveform = BreathWaveform.create()
    state = PIDControllerState(waveform=waveform)
    return state

  @jax.jit
  def __call__(self, controller_state, obs):
    pressure, t = obs.predicted_pressure, obs.time
    waveform = controller_state.waveform
    target = waveform.at(t)
    err = jnp.array(target - pressure)
    decay = jnp.array(self.dt / (self.dt + self.RC))
    p, i, d = controller_state.p, controller_state.i, controller_state.d
    next_p = err
    next_i = i + decay * (err - i)
    next_d = d + decay * (err - p - d)
    controller_state = controller_state.replace(p=next_p, i=next_i, d=next_d)
    next_coef = jnp.array([next_p, next_i, next_d])
    u_in = self.model_apply({"params": self.params}, next_coef)
    u_in = jax.lax.clamp(0.0, u_in.astype(jnp.float64), 100.0)

    # update controller_state
    new_dt = jnp.max(
        jnp.array([DEFAULT_DT, t - proper_time(controller_state.time)]))
    new_time = t
    new_steps = controller_state.steps + 1
    controller_state = controller_state.replace(
        time=new_time, steps=new_steps, dt=new_dt)
    return controller_state, u_in class Predestined(Controller):
  """predestined controller."""
  time: float = deluca.field(jaxed=False)
  steps: int = deluca.field(jaxed=False)
  dt: float = deluca.field(jaxed=False)
  u_ins: jnp.array = deluca.field(jaxed=False)

  def __call__(self, state, obs, *args, **kwargs):
    action = jax.lax.dynamic_slice(self.u_ins, (state.steps.astype(int),), (1,))
    time = obs.time
    new_dt = jnp.max(jnp.array([DEFAULT_DT, time - proper_time(state.time)]))
    new_time = time
    new_steps = state.steps + 1
    state = state.replace(time=new_time, steps=new_steps, dt=new_dt)
    return state, action class CompositeController(core.Controller):
  """Controller comprised of a base controller along with a residual controller.

  Only the residual controller is trainable.
  """
  base_controller: core.Controller = deluca.field(jaxed=False)
  resid_controller: core.Controller = deluca.field(jaxed=True)
  use_leaky_clamp: bool = deluca.field(True, jaxed=False)
  clip: float = deluca.field(100.0, jaxed=False)

  def init(self, waveform=None):
    return CompositeState.create(
        self.base_controller.init(waveform),
        self.resid_controller.init(waveform))

  @jax.jit
  def __call__(self, controller_state: CompositeState, obs):
    base_state, u_in_base = self.base_controller(controller_state.base_state,
                                                 obs)
    resid_state, u_in_resid = self.resid_controller(
        controller_state.resid_state, obs)

    u_in = u_in_base + u_in_resid

    # Implementing "leaky" clamp to solve the zero gradient problem
    if self.use_leaky_clamp:
      u_in = jax.lax.cond(u_in.squeeze() < 0.0,
                          lambda x: x * 0.01,
                          lambda x: x, u_in.squeeze())
      u_in = jax.lax.cond(u_in.squeeze() > self.clip,
                          lambda x: self.clip + x * 0.01,
                          lambda x: x, u_in.squeeze())
    else:
      u_in = jax.lax.clamp(0.0, u_in.astype(jnp.float32), self.clip).squeeze()

    return CompositeState.create(base_state, resid_state), u_in class BalloonLung(LungEnv):
  """Lung simulator based on the two-balloon experiment.

     Source: https://en.wikipedia.org/wiki/Two-balloon_experiment
     TODO:
      - time / dt
      - dynamics
      - waveform
      - phase
  """

  leak: bool = deluca.field(False, jaxed=False)
  peep_valve: float = deluca.field(5.0, jaxed=False)
  PC: float = deluca.field(40.0, jaxed=False)
  P0: float = deluca.field(0.0, jaxed=False)
  R: float = deluca.field(15.0, jaxed=False)
  C: float = deluca.field(10.0, jaxed=False)
  dt: float = deluca.field(0.03, jaxed=False)
  min_volume: float = deluca.field(1.5, jaxed=False)
  r0: float = deluca.field(jaxed=False)
  waveform: BreathWaveform = deluca.field(jaxed=False)
  # reward_fn: Callable = deluca.field(None, jaxed=False)
  reset_normalized_peep: float = deluca.field(0.0, jaxed=False)
  flow: float = deluca.field(0, jaxed=False)  # needed for run_controller

  def setup(self):
    if self.waveform is None:
      self.waveform = BreathWaveform.create()

    # dynamics hyperparameters
    # self.time = 0.0

    self.r0 = (3.0 * self.min_volume / (4.0 * jnp.pi))**(1.0 / 3.0)

    # reset states
    # self.reset()

  def reset(self):
    state = SimulatorState(
        steps=0,
        time=0.0,
        volume=self.min_volume,
        predicted_pressure=self.reset_normalized_peep,
        target=self.waveform.at(0.0))
    observation = Observation(
        predicted_pressure=self.reset_normalized_peep, time=0.0)
    return state, observation

  def __call__(self, state, action):
    """call function.

    Args:
      state: (steps, time, volume, predicted_pressure, target)
      action: (u_in, u_out)
    Returns:
      state: next state
      observation: next observation
    """
    volume, pressure = state.volume, state.predicted_pressure
    u_in, u_out = action

    flow = jnp.clip(PropValve(u_in) * self.R, 0.0, 2.0)
    flow -= jax.lax.cond(
        pressure > self.peep_valve,
        lambda x: jnp.clip(Solenoid(u_out), 0.0, 2.0) * 0.05 * pressure,
        lambda x: 0.0,
        flow,
    )

    volume += flow * self.dt
    volume += jax.lax.cond(
        self.leak,
        lambda x: ((self.dt / (5.0 + self.dt) *
                    (self.min_volume - volume))).squeeze(),
        lambda x: 0.0,
        0.0,
    )

    r = (3.0 * volume / (4.0 * jnp.pi))**(1.0 / 3.0)
    pressure = self.P0 + self.PC * (1.0 - (self.r0 / r)**6.0) / (
        self.r0**2.0 * r)
    # pressure = flow * self.R + volume / self.C + self.peep_valve

    state = state.replace(
        steps=state.time + 1,
        time=state.time + self.dt,
        volume=volume.squeeze(),
        predicted_pressure=pressure.squeeze(),
        target=self.waveform.at(state.time + self.dt))

    observation = Observation(
        predicted_pressure=state.predicted_pressure, time=state.time)
    return state, observation class DelayLung(LungEnv):
  """Delay lung."""
  min_volume: float = deluca.field(1.5, jaxed=False)
  R: int = deluca.field(10, jaxed=False)
  C: int = deluca.field(6, jaxed=False)
  delay: int = deluca.field(25, jaxed=False)
  inertia: float = deluca.field(0.995, jaxed=False)
  control_gain: float = deluca.field(0.02, jaxed=False)
  dt: float = deluca.field(0.03, jaxed=False)
  waveform: BreathWaveform = deluca.field(jaxed=False)
  r0: float = deluca.field(jaxed=False)
  flow: float = deluca.field(0.0, jaxed=False)  # Question: initial flow value?

  def setup(self):
    if self.waveform is None:
      self.waveform = BreathWaveform.create()
    self.r0 = (3 * self.min_volume / (4 * jnp.pi))**(1 / 3)

  def reset(self):
    state = SimulatorState(
        steps=0,
        time=0.0,
        volume=self.min_volume,
        predicted_pressure=0.0,
        pipe_pressure=0.0,
        target=self.waveform.at(0.0),
        in_history=jnp.zeros(self.delay),
        out_history=jnp.zeros(self.delay))
    observation = Observation(
        predicted_pressure=0.0, time=0.0)
    return state, observation

  def dynamics(self, state, action):
    """dynamics function.

    Args:
      state: (volume, pressure)
      action: (u_in, u_out)
    Returns:
      state: next state
    """
    flow = state.predicted_pressure / self.R
    volume = state.volume + flow * self.dt
    volume = jnp.maximum(volume, self.min_volume)

    r = (3.0 * volume / (4.0 * jnp.pi))**(1.0 / 3.0)
    lung_pressure = self.C * (1 - (self.r0 / r)**6) / (self.r0**2 * r)

    pipe_impulse = jax.lax.cond(
        state.time < self.delay,
        lambda x: 0.0,
        lambda x: self.control_gain * state.in_history[0],
        None,
    )
    peep = jax.lax.cond(state.time < self.delay, lambda x: 0.0,
                        lambda x: state.out_history[0], None)

    pipe_pressure = self.inertia * state.pipe_pressure + pipe_impulse
    pressure = jnp.maximum(0, pipe_pressure - lung_pressure)

    pipe_pressure = jax.lax.cond(peep, lambda x: x * 0.995, lambda x: x,
                                 pipe_pressure)

    state = state.replace(
        steps=state.time + 1,
        time=state.time + self.dt,
        volume=volume,
        predicted_pressure=pressure,
        pipe_pressure=pipe_pressure,
        target=self.waveform.at(state.time + self.dt))
    return state

  def __call__(self, state, action):
    u_in, u_out = action
    u_in = jax.lax.cond(u_in.squeeze() > 0.0,
                        lambda x: x,
                        lambda x: 0.0,
                        u_in.squeeze())

    next_in_history = jnp.roll(state.in_history, shift=1)
    next_in_history = next_in_history.at[0].set(u_in)
    next_out_history = jnp.roll(state.out_history, shift=1)
    next_out_history = next_out_history.at[0].set(u_out)

    state = state.replace(in_history=next_in_history,
                          out_history=next_out_history)

    next_state = self.dynamics(state, action)
    observation = Observation(
        predicted_pressure=state.predicted_pressure, time=state.time)

    return next_state, observation class LearnedLung(LungEnv):
  """Learned lung."""
  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_normalizer: ShiftScaleTransform = deluca.field(jaxed=False)
  p_normalizer: ShiftScaleTransform = deluca.field(jaxed=False)
  seed: int = deluca.field(0, jaxed=False)
  flow: int = deluca.field(0, jaxed=False)
  transition_threshold: int = deluca.field(0, jaxed=False)

  default_model_parameters: Dict[str, Any] = deluca.field(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_normalized_peep: float = deluca.field(0.0, jaxed=False)

  default_model_name: str = deluca.field("MLP", jaxed=False)
  default_model: nn.module = deluca.field(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 = MLP(
        hidden_dim=self.default_model_parameters["hidden_dim"],
        out_dim=self.default_model_parameters["out_dim"],
        n_layers=self.default_model_parameters["n_layers"],
        droprate=self.default_model_parameters["droprate"],
        activation_fn=self.default_model_parameters["activation_fn"])
    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]
    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 init(self, key):
    return self.reset()


  def reset(self):
    normalized_peep = self.reset_normalized_peep
    state = SimulatorState(
        u_history=jnp.zeros([self.u_history_len]),
        p_history=jnp.hstack(
            [jnp.zeros([self.p_history_len - 1]),
             jnp.array([normalized_peep])]),
        t_in=0,
        steps=0,
        predicted_pressure=normalized_peep,
    )
    state = self.sync_unnormalized_pressure(state)
    obs = LearnedLungObservation(
        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_unnormalized_pressure(self, state):
    normalized_pressure = self.p_normalizer.inverse(
        state.predicted_pressure).squeeze()
    return state.replace(predicted_pressure=normalized_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 = [
        functools.partial(self.ensemble_models[i].apply,
                          {"params": self.params[i]})
        for i in range(self.num_boundary_models + 1)
    ]
    # partial_false_func = functools.partial(
    #     false_func, u_in=u_in, model_idx=model_idx, funcs=funcs)
    partial_false_func = functools.partial(
        false_func,
        u_in=u_in,
        model_idx=model_idx,
        u_normalizer=self.u_normalizer,
        update_history=self.update_history,
        u_history_len=self.u_history_len,
        p_history_len=self.p_history_len,
        u_window=self.u_window,
        p_window=self.p_window,
        sync_unnormalized_pressure=self.sync_unnormalized_pressure,
        num_boundary_models=self.num_boundary_models,
        funcs=funcs)
    partial_true_func = functools.partial(true_func, reset=self.reset)
    state = jax.lax.cond(u_out == 1, partial_true_func, partial_false_func,
                         state)
    state = state.replace(steps=state.steps + 1)
    obs = LearnedLungObservation(
        predicted_pressure=state.predicted_pressure, time=self.time(state))
    return state, obs class SingleCompLung(LungEnv):
  """Single comp lung."""
  R: int = deluca.field(15, jaxed=False)
  C: int = deluca.field(50, jaxed=False)
  min_volume: float = deluca.field(0.2, jaxed=False)
  peep_value: float = deluca.field(5.0, jaxed=False)
  flow: float = deluca.field(0.0, jaxed=False)
  P0: float = deluca.field(10.0, jaxed=False)
  dt: float = deluca.field(0.03, jaxed=False)

  def reset(self):
    vol = self.min_volume
    obs = (vol / self.C) + self.P0

    # state = SimulatorState(0, vol)
    state = SimulatorState(steps=0, volume=vol, time=0.0)

    # print(type(state.steps))
    return state, Observation(obs, 0.0)

  def __call__(self, state, action):
    u_in, _ = action
    # calculate flow from u_in
    flow = jnp.clip(PropValve(u_in), 0, 2)
    # update volume by flow rate
    next_volume = (state.volume + flow * self.dt).squeeze()
    obs = ((flow * self.R) + (state.volume / self.C) +
           self.P0).squeeze()

    state = state.replace(
        steps=state.steps + 1, volume=next_volume, time=state.time + self.dt)

    return state, Observation(predicted_pressure=obs, time=state.time) class BreathDataset:
  """dataset for breaths."""

  @classmethod
  def from_paths(
      cls,
      paths,
      seed=0,
      clip=(0, None),  # used to be (2, -1)
      breath_length=29,
      keys=["train", "test"],
      splits=[0.9, 0.1],
      fit_scaler_key="train",
  ):
    """process data from paths."""
    obj = object.__new__(cls)
    raw_data = []

    for path in paths:
      analyzer = Analyzer(path)
      raw_data += analyzer.get_breaths(clip=clip, breath_length=breath_length)

    # Shuffle and split raw_data
    rng = np.random.default_rng(seed)
    splits = (np.array(splits) / np.sum(splits) *
              len(raw_data)).astype("int")[:-1]
    rng.shuffle(raw_data)
    obj.splits = {
        key: val for key, val in zip(keys, np.split(raw_data, splits))
    }

    # Compute mean, std for u_in and pressure
    obj.u_normalizer = ShiftScaleTransform(
        jnp.array([u for u, p in obj.splits[fit_scaler_key]]))
    print(f"u_in: mean={obj.u_normalizer.mean}, std={obj.u_normalizer.std}")
    obj.p_normalizer = ShiftScaleTransform(
        jnp.array([p for u, p in obj.splits[fit_scaler_key]]))
    print(f"pressure: mean={obj.p_normalizer.mean}, std={obj.p_normalizer.std}")

    obj.data = {}
    for key in keys:
      obj.data[key] = (jnp.array([u for u, p in obj.splits[key]]),
                       jnp.array([p for u, p in obj.splits[key]]))
    return obj def get_initial_pressure(dataset, key):
  return jnp.array([dataset.p_normalizer(p[0]) for u, p in dataset.splits[key]
                   ]).mean().item() @jax.jit
def test_controller(controller, sim, pips, peep):
  """Test controller."""
  # new_controller = controller.replace(use_leaky_clamp=False)
  score = 0.0
  horizon = 29
  for pip in pips:
    waveform = BreathWaveform.create(peep=peep, pip=pip)
    result = run_controller_scan(
        controller,
        T=horizon,
        abort=horizon,
        env=sim,
        waveform=waveform,
        init_controller=True,
    )
    analyzer = Analyzer(result)
    preds = analyzer.pressure  # shape = (29,)
    truth = analyzer.target  # shape = (29,)
    # print('preds.shape: %s', str(preds.shape))
    # print('truth.shape: %s', str(truth.shape))
    score += jnp.abs(preds - truth).mean()
  score = score / len(pips)
  return score
from absl.testing import absltest
from absl.testing import parameterized
from deluca.lung.controllers._bang_bang import BangBang
from deluca.lung.controllers._deep_cnn import DeepCnn
from deluca.lung.controllers._deep_mlp import DeepMlp
from deluca.lung.controllers._pid import PID
from deluca.lung.controllers._predestined import Predestined
from deluca.lung.controllers._residual import CompositeController
from deluca.lung.envs._balloon_lung import BalloonLung
from deluca.lung.envs._delay_lung import DelayLung
from deluca.lung.envs._learned_lung import LearnedLung
from deluca.lung.envs._single_comp_lung import SingleCompLung
from deluca.lung.utils.data.breath_dataset import BreathDataset
from deluca.lung.utils.data.breath_dataset import get_initial_pressure
from deluca.lung.utils.scripts.test_controller import test_controller
import chex
import flax.linen as nn
import jax
import jax.numpy as jnp
# Copyright 2022 The Deluca Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for test_controller."""



class TestControllerParameterizedTest(chex.TestCase):

  def setUp(self):
    super().setUp()
    result = {}
    u_in = jnp.array([1 for i in range(29)] + [0 for i in range(29)] +
                     [1 for i in range(29)] + [0])
    u_out = jnp.array([0 for i in range(29)] + [1 for i in range(29)] +
                      [0 for i in range(29)] + [1])
    pressure = jnp.array([1 for i in range(29)] + [0 for i in range(29)] +
                         [1 for i in range(29)] + [0])
    target = jnp.array([1 for i in range(29)] + [0 for i in range(29)] +
                       [1 for i in range(29)] + [0])
    u_in = jnp.abs(u_in + jax.random.normal(jax.random.PRNGKey(0), (88,)))
    u_out = jnp.abs(u_out + jax.random.normal(jax.random.PRNGKey(1), (88,)))
    pressure = jnp.abs(pressure +
                       jax.random.normal(jax.random.PRNGKey(2), (88,)))
    target = jnp.abs(target + jax.random.normal(jax.random.PRNGKey(3), (88,)))

    timeseries = {
        'timestamp': jnp.array([0.03 * i for i in range(88)]),
        'pressure': pressure,
        'flow': jnp.array([0 for i in range(88)]),
        'target': target,
        'u_in': u_in,
        'u_out': u_out
    }
    result['timeseries'] = timeseries
    paths = [result for i in range(10)]

    self.dataset = BreathDataset.from_paths(paths, clip=(0, None))