# Path: deluca/lung/core.py
# class LungEnv(deluca.Env):
#   """Lung environment."""
# 
#   def time(self, state):
#     return self.dt * state.steps
# 
#   @property
#   def dt(self):
#     return DEFAULT_DT
# 
#   @property
#   def physical(self):
#     return False
# 
#   def should_abort(self):
#     return False
# 
#   def wait(self, duration):
#     pass
# 
#   def cleanup(self):
#     pass


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

"""Single comp lung."""


class Observation(deluca.Obj):
  predicted_pressure: float = 0.0
  time: float = 0.0


class SimulatorState(deluca.Obj):
  steps: int = 0
  volume: float = 0.0
  time: float = 0.0


def PropValve(x):
  y = 3.0 * x
  flow_new = 1.0 * (jnp.tanh(0.03 * (y - 130)) + 1.0)
  flow_new = jnp.clip(flow_new, 0.0, 1.72)
  return flow_new



# Path: deluca/lung/core.py
# class LungEnv(deluca.Env):
#   """Lung environment."""
# 
#   def time(self, state):
#     return self.dt * state.steps
# 
#   @property
#   def dt(self):
#     return DEFAULT_DT
# 
#   @property
#   def physical(self):
#     return False
# 
#   def should_abort(self):
#     return False
# 
#   def wait(self, duration):
#     pass
# 
#   def cleanup(self):
#     pass


# Path: deluca/lung/envs/_single_comp_lung.py
import deluca.core
import jax.numpy as jnp
from deluca.lung.core import LungEnv

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

"""Single comp lung."""


class Observation(deluca.Obj):
  predicted_pressure: float = 0.0
  time: float = 0.0


class SimulatorState(deluca.Obj):
  steps: int = 0
  volume: float = 0.0
  time: float = 0.0


def PropValve(x):
  y = 3.0 * x
  flow_new = 1.0 * (jnp.tanh(0.03 * (y - 130)) + 1.0)
  flow_new = jnp.clip(flow_new, 0.0, 1.72)
  return flow_new



class SingleCompLung(LungEnv):