# Path: deluca/lung/core.py
# class Controller(deluca.Agent):
#   """Controller."""
# 
#   def init(self, waveform=None):
#     if waveform is None:
#       waveform = BreathWaveform.create()
#     return ControllerState()
# 
#   @property
#   def max(self):
#     return 100
# 
#   @property
#   def min(self):
#     return 0
# 
#   def decay(self, waveform, t):
#     elapsed = waveform.elapsed(t)
# 
#     def false_func():
#       result = jax.lax.cond(
#           elapsed < waveform.xp[waveform.ex_bounds[0]], lambda x: 0.0,
#           lambda x: 5 * (1 - jnp.exp(5 * (waveform.xp[waveform.ex_bounds[
#               0]] - elapsed))).astype(waveform.dtype), None)
#       return result
# 
#     # float(inf) as substitute to None since cond requries same type output
#     result = jax.lax.cond(elapsed < waveform.xp[waveform.in_bounds[1]],
#                           lambda x: float("inf"), lambda x: false_func(), None)
#     return result
# 
# Path: deluca/lung/core.py
# DEFAULT_DT = 0.03
# 
# Path: deluca/lung/core.py
# def proper_time(t):
#   return jax.lax.cond(t == float("inf"), lambda x: 0., lambda x: x, t)


import deluca
import jax
import jax.numpy as jnp
from deluca.lung.core import Controller
from deluca.lung.core import DEFAULT_DT
from deluca.lung.core import proper_time
# 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.

"""predestined controller."""



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

# Path: deluca/lung/core.py
# class Controller(deluca.Agent):
#   """Controller."""
# 
#   def init(self, waveform=None):
#     if waveform is None:
#       waveform = BreathWaveform.create()
#     return ControllerState()
# 
#   @property
#   def max(self):
#     return 100
# 
#   @property
#   def min(self):
#     return 0
# 
#   def decay(self, waveform, t):
#     elapsed = waveform.elapsed(t)
# 
#     def false_func():
#       result = jax.lax.cond(
#           elapsed < waveform.xp[waveform.ex_bounds[0]], lambda x: 0.0,
#           lambda x: 5 * (1 - jnp.exp(5 * (waveform.xp[waveform.ex_bounds[
#               0]] - elapsed))).astype(waveform.dtype), None)
#       return result
# 
#     # float(inf) as substitute to None since cond requries same type output
#     result = jax.lax.cond(elapsed < waveform.xp[waveform.in_bounds[1]],
#                           lambda x: float("inf"), lambda x: false_func(), None)
#     return result
# 
# Path: deluca/lung/core.py
# DEFAULT_DT = 0.03
# 
# Path: deluca/lung/core.py
# def proper_time(t):
#   return jax.lax.cond(t == float("inf"), lambda x: 0., lambda x: x, t)


# Path: deluca/lung/controllers/_predestined.py
import deluca
import jax
import jax.numpy as jnp
from deluca.lung.core import Controller
from deluca.lung.core import DEFAULT_DT
from deluca.lung.core import proper_time

# 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.

"""predestined controller."""



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)]))