# 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/utils/data/transform.py
# class ShiftScaleTransform:
#   """Data normalizer."""
# 
#   # vectors is an array of vectors
#   def __init__(self, vectors):
#     vectors_concat = jnp.concatenate(vectors)
#     self.mean = jnp.mean(vectors_concat)
#     self.std = jnp.std(vectors_concat)
#     print(self.mean, self.std)
# 
#   def _transform(self, x, mean, std):
#     return (x - mean) / std
# 
#   def _inverse_transform(self, x, mean, std):
#     return (x * std) + mean
# 
#   def __call__(self, vector):
#     return self._transform(vector, self.mean, self.std)
# 
#   def inverse(self, vector):
#     return self._inverse_transform(vector, self.mean, self.std)
# 
# Path: deluca/lung/utils/nn/mlp.py
# class MLP(nn.Module):
#   """multilayered perceptron."""
#   hidden_dim: int = 10
#   out_dim: int = 1
#   n_layers: int = 2
#   droprate: float = 0.0
#   activation_fn: Callable = nn.relu
# 
#   @nn.compact
#   def __call__(self, x):
#     for i in range(self.n_layers - 1):
#       x = nn.Dense(
#           features=self.hidden_dim, use_bias=True, name=f"MLP_fc{i}")(
#               x)
#       x = nn.Dropout(
#           rate=self.droprate, deterministic=False)(
#               x, rng=jax.random.PRNGKey(0))
#       x = self.activation_fn(x)
#     x = nn.Dense(features=self.out_dim, use_bias=True, name=f"MLP_fc{i + 1}")(x)
#     return x.squeeze()  # squeeze for consistent shape w/ boundary model output
# 
# Path: deluca/lung/utils/nn/shallow_boundary_model.py
# class ShallowBoundaryModel(nn.Module):
#   """Shallow boundary module."""
#   out_dim: int = 1
#   hidden_dim: int = 100
#   model_num: int = 0
# 
#   @nn.compact
#   def __call__(self, x):
#     # need to flatten extra dimensions required by CNN and LSTM
#     x = x.squeeze()
#     x = nn.Dense(
#         features=self.hidden_dim,
#         use_bias=False,
#         name=f"shallow_fc{1}_model" + str(self.model_num),
#     )(
#         x)
#     x = nn.tanh(x)
#     x = nn.Dense(
#         features=self.out_dim,
#         use_bias=True,
#         name=f"shallow_fc{2}_model" + str(self.model_num))(
#             x)
#     return x.squeeze()  # squeeze for consistent shape w/ boundary model output


import functools
import deluca.core
import flax.linen as nn
import jax
import jax.numpy as jnp
from typing import Any, Dict
from deluca.lung.core import LungEnv
from deluca.lung.utils.data.transform import ShiftScaleTransform
from deluca.lung.utils.nn import MLP
from deluca.lung.utils.nn import ShallowBoundaryModel
  default_features = jnp.hstack(
      [state.u_history[-u_window:], state.p_history[-p_window:]])
  default_pad_len = (u_history_len + p_history_len) - (u_window + p_window)
  default_features = jnp.hstack(
      [jnp.zeros((default_pad_len,)), default_features])
  features = jax.lax.cond(
      model_idx == num_boundary_models,
      lambda x: default_features,
      lambda x: boundary_features,
      None,
  )
  normalized_pressure = jax.lax.switch(model_idx, funcs, features)
  normalized_pressure = normalized_pressure.astype(jnp.float64)
  state = state.replace(predicted_pressure=normalized_pressure)
  new_p_history = update_history(state.p_history, normalized_pressure)
  state = state.replace(p_history=new_p_history
                       )  # just for inference. This is undone during training
  state = sync_unnormalized_pressure(state)  # unscale predicted pressure
  return state


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)

# 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/utils/data/transform.py
# class ShiftScaleTransform:
#   """Data normalizer."""
# 
#   # vectors is an array of vectors
#   def __init__(self, vectors):
#     vectors_concat = jnp.concatenate(vectors)
#     self.mean = jnp.mean(vectors_concat)
#     self.std = jnp.std(vectors_concat)
#     print(self.mean, self.std)
# 
#   def _transform(self, x, mean, std):
#     return (x - mean) / std
# 
#   def _inverse_transform(self, x, mean, std):
#     return (x * std) + mean
# 
#   def __call__(self, vector):
#     return self._transform(vector, self.mean, self.std)
# 
#   def inverse(self, vector):
#     return self._inverse_transform(vector, self.mean, self.std)
# 
# Path: deluca/lung/utils/nn/mlp.py
# class MLP(nn.Module):
#   """multilayered perceptron."""
#   hidden_dim: int = 10
#   out_dim: int = 1
#   n_layers: int = 2
#   droprate: float = 0.0
#   activation_fn: Callable = nn.relu
# 
#   @nn.compact
#   def __call__(self, x):
#     for i in range(self.n_layers - 1):
#       x = nn.Dense(
#           features=self.hidden_dim, use_bias=True, name=f"MLP_fc{i}")(
#               x)
#       x = nn.Dropout(
#           rate=self.droprate, deterministic=False)(
#               x, rng=jax.random.PRNGKey(0))
#       x = self.activation_fn(x)
#     x = nn.Dense(features=self.out_dim, use_bias=True, name=f"MLP_fc{i + 1}")(x)
#     return x.squeeze()  # squeeze for consistent shape w/ boundary model output
# 
# Path: deluca/lung/utils/nn/shallow_boundary_model.py
# class ShallowBoundaryModel(nn.Module):
#   """Shallow boundary module."""
#   out_dim: int = 1
#   hidden_dim: int = 100
#   model_num: int = 0
# 
#   @nn.compact
#   def __call__(self, x):
#     # need to flatten extra dimensions required by CNN and LSTM
#     x = x.squeeze()
#     x = nn.Dense(
#         features=self.hidden_dim,
#         use_bias=False,
#         name=f"shallow_fc{1}_model" + str(self.model_num),
#     )(
#         x)
#     x = nn.tanh(x)
#     x = nn.Dense(
#         features=self.out_dim,
#         use_bias=True,
#         name=f"shallow_fc{2}_model" + str(self.model_num))(
#             x)
#     return x.squeeze()  # squeeze for consistent shape w/ boundary model output


# Path: deluca/lung/envs/_learned_lung.py
import functools
import deluca.core
import flax.linen as nn
import jax
import jax.numpy as jnp
from typing import Any, Dict
from deluca.lung.core import LungEnv
from deluca.lung.utils.data.transform import ShiftScaleTransform
from deluca.lung.utils.nn import MLP
from deluca.lung.utils.nn import ShallowBoundaryModel

  default_features = jnp.hstack(
      [state.u_history[-u_window:], state.p_history[-p_window:]])
  default_pad_len = (u_history_len + p_history_len) - (u_window + p_window)
  default_features = jnp.hstack(
      [jnp.zeros((default_pad_len,)), default_features])
  features = jax.lax.cond(
      model_idx == num_boundary_models,
      lambda x: default_features,
      lambda x: boundary_features,
      None,
  )
  normalized_pressure = jax.lax.switch(model_idx, funcs, features)
  normalized_pressure = normalized_pressure.astype(jnp.float64)
  state = state.replace(predicted_pressure=normalized_pressure)
  new_p_history = update_history(state.p_history, normalized_pressure)
  state = state.replace(p_history=new_p_history
                       )  # just for inference. This is undone during training
  state = sync_unnormalized_pressure(state)  # unscale predicted pressure
  return state


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)