# Path: ehr_prediction_modeling/types.py
# class ForwardFunctionReturns(
#     collections.namedtuple(
#         "ForwardFunctionReturns",
#         ["model_output", "hidden_states", "inputs", "activations"],
#     )):
# class EmbeddingType(object):
# class EmbeddingEncoderType(object):
# class ActivationFunction(object):
# class RegularizationType(object):
# class TaskType(object):
# class BestModelMetric(object):
# class TaskNames(object):
# class TaskTypes(object):
# class TaskLossType(object):
# class RNNCellType(object):
# class FeatureTypes(object):
# class Optimizers(object):
# class LearningRateScheduling(object):
# class EmbeddingCombinationMethod(object):
# class SNRBlockConnType(object):
# class SubNettoSubNetConnType(object):
# class SNRInputCombinationType(object):
# class LossCombinationType(object):
# class ModelTypes(object):
# class TaskLayerTypes(object):
# class ModelMode(object):
#   LOOKUP = "Lookup"
#   DEEP_EMBEDDING = "DeepEmbedding"
#   RESIDUAL = "residual"
#   FC = "fc"
#   SNR = "snr"
#   TANH = "tanh"
#   RELU = "relu"
#   LRELU = "lrelu"
#   SWISH = "swish"
#   ELU = "elu"
#   SLU = "selu"
#   ELISH = "elish"
#   HARD_ELISH = "hard_elish"
#   SIGMOID = "sigmoid"
#   HARD_SIGMOID = "hard_sigmoid"
#   TANH_PEN = "tanh_pen"
#   NONE = "None"
#   L1 = "L1"
#   L2 = "L2"
#   BINARY_CLASSIFICATION = "BinaryClassification"
#   REGRESSION = "Regression"
#   PRAUC = "prauc"
#   ROCAUC = "rocauc"
#   L1 = "l1"
#   ADVERSE_OUTCOME_RISK = "AdverseOutcomeRisk"
#   LAB_REGRESSION = "Labs"
#   BINARIZED_LOS = "BinarizedLengthOfStay"
#   REGRESSION_LOS = "RegressionLengthOfStay"
#   MORTALITY = "MortalityRisk"
#   READMISSION = "ReadmissionRisk"
#   ADVERSE_OUTCOME_RISK = "ao_risk"
#   LAB_REGRESSION = "labs_regression"
#   LOS = "length_of_stay"
#   MORTALITY = "mortality"
#   READMISSION = "readmission"
#   CE = "CE"
#   MULTI_CE = "MULTI_CE"
#   BRIER = "Brier"
#   L1 = "L1"
#   L2 = "L2"
#   SRU = "SRU"
#   LSTM = "LSTM"
#   UGRNN = "UGRNN"
#   PRESENCE_SEQ = "pres_s"  # Sequential presence features
#   NUMERIC_SEQ = "num_s"  # Sequential numerical features
#   CATEGORY_COUNTS_SEQ = "count_s"  # Sequential category counts features
# FEAT_TO_NAME = {
#     "pres_s": "presence",
#     "num_s": "numeric",
#     "count_s": "category_counts",
# }
#   SGD = "SGD"
#   ADAM = "Adam"
#   RMSPROP = "RMSprop"
#   FIXED = "fixed"
#   EXPONENTIAL_DECAY = "exponential_decay"
#   CONCATENATE = "concatenate"
#   SUM_ALL = "sum_all"
#   SUM_BY_SUFFIX = "sum_by_suffix"
#   COMBINE_SNR_OUT = "combine_snr_out"
#   FC = "fc"
#   HIGHWAY = "highway"
#   RESIDUAL = "residual"
#   NONE = "None"
#   BOOL = "bool"
#   SCALAR_WEIGHT = "scalar_weight"
#   CONCATENATE = "concatenate"
#   SUM_ALL = "sum"
#   SUM_ALL = "SUM_ALL"
#   UNCERTAINTY_WEIGHTED = "UNCERTAINTY_WEIGHTED"
#   RNN = "RNN"
#   SNRNN = "SNRNN"
#   NONE = "none"
#   MLP = "MLP"
#   SNRMLP = "SNRMLP"
#   TRAIN = "train"
#   EVAL = "eval"
#   PREDICT = "predict"
#   def is_train(cls, mode: str) -> bool:
# 
# Path: ehr_prediction_modeling/models/reset_core.py
# class ResetCore(snt.RNNCore):
#   def __init__(self, core: RNNCore):
#   def _build(self, input_should_reset: Tuple[tf.Tensor, tf.Tensor],
#              state: tf.Tensor, **kwargs) -> Tuple[tf.Tensor, tf.Tensor]:
#   def wrapped_core(self) -> RNNCore:
#   def state_size(self) -> tf.TensorShape:
#   def output_size(self) -> tf.TensorShape:
#   def trainable_weights(self) -> List[tf.Tensor]:
#   def initial_state(self, *args, **kwargs) -> List[tf.Tensor]:


import functools
import typing
import sonnet as snt
import tensorflow.compat.v1 as tf
from ehr_prediction_modeling import types
from ehr_prediction_modeling.models import reset_core
from tensorflow import contrib as tf_contrib
  from ehr_prediction_modeling import configdict
              tf_contrib.rnn.UGRNNCell,
              num_units=config.ndim_lstm,
              activation=config.activation_fn),
  }

  if config.cell_type not in cell_construction_mapping:
    raise ValueError(f"Invalid cell_type: {config.cell_type}")

  cell = cell_construction_mapping[config.cell_type]()

  if config.cell_type in cells_not_allowing_wrappers:
    # Certain cells do not support wrappers.
    return cell

  if config.use_highway_connections:
    cell = tf_contrib.rnn.HighwayWrapper(cell)
  return cell


def init_deep_cell(
    cell_config: "configdict.ConfigDict") -> reset_core.ResetCore:
  """Initializes and wraps a cell in a ResetCore.

  Args:
    cell_config: Cell ConfigDict.

  Returns:
    An RNNCore object.
  """
  cell = _map_cell_type(cell_config)

# Path: ehr_prediction_modeling/types.py
# class ForwardFunctionReturns(
#     collections.namedtuple(
#         "ForwardFunctionReturns",
#         ["model_output", "hidden_states", "inputs", "activations"],
#     )):
# class EmbeddingType(object):
# class EmbeddingEncoderType(object):
# class ActivationFunction(object):
# class RegularizationType(object):
# class TaskType(object):
# class BestModelMetric(object):
# class TaskNames(object):
# class TaskTypes(object):
# class TaskLossType(object):
# class RNNCellType(object):
# class FeatureTypes(object):
# class Optimizers(object):
# class LearningRateScheduling(object):
# class EmbeddingCombinationMethod(object):
# class SNRBlockConnType(object):
# class SubNettoSubNetConnType(object):
# class SNRInputCombinationType(object):
# class LossCombinationType(object):
# class ModelTypes(object):
# class TaskLayerTypes(object):
# class ModelMode(object):
#   LOOKUP = "Lookup"
#   DEEP_EMBEDDING = "DeepEmbedding"
#   RESIDUAL = "residual"
#   FC = "fc"
#   SNR = "snr"
#   TANH = "tanh"
#   RELU = "relu"
#   LRELU = "lrelu"
#   SWISH = "swish"
#   ELU = "elu"
#   SLU = "selu"
#   ELISH = "elish"
#   HARD_ELISH = "hard_elish"
#   SIGMOID = "sigmoid"
#   HARD_SIGMOID = "hard_sigmoid"
#   TANH_PEN = "tanh_pen"
#   NONE = "None"
#   L1 = "L1"
#   L2 = "L2"
#   BINARY_CLASSIFICATION = "BinaryClassification"
#   REGRESSION = "Regression"
#   PRAUC = "prauc"
#   ROCAUC = "rocauc"
#   L1 = "l1"
#   ADVERSE_OUTCOME_RISK = "AdverseOutcomeRisk"
#   LAB_REGRESSION = "Labs"
#   BINARIZED_LOS = "BinarizedLengthOfStay"
#   REGRESSION_LOS = "RegressionLengthOfStay"
#   MORTALITY = "MortalityRisk"
#   READMISSION = "ReadmissionRisk"
#   ADVERSE_OUTCOME_RISK = "ao_risk"
#   LAB_REGRESSION = "labs_regression"
#   LOS = "length_of_stay"
#   MORTALITY = "mortality"
#   READMISSION = "readmission"
#   CE = "CE"
#   MULTI_CE = "MULTI_CE"
#   BRIER = "Brier"
#   L1 = "L1"
#   L2 = "L2"
#   SRU = "SRU"
#   LSTM = "LSTM"
#   UGRNN = "UGRNN"
#   PRESENCE_SEQ = "pres_s"  # Sequential presence features
#   NUMERIC_SEQ = "num_s"  # Sequential numerical features
#   CATEGORY_COUNTS_SEQ = "count_s"  # Sequential category counts features
# FEAT_TO_NAME = {
#     "pres_s": "presence",
#     "num_s": "numeric",
#     "count_s": "category_counts",
# }
#   SGD = "SGD"
#   ADAM = "Adam"
#   RMSPROP = "RMSprop"
#   FIXED = "fixed"
#   EXPONENTIAL_DECAY = "exponential_decay"
#   CONCATENATE = "concatenate"
#   SUM_ALL = "sum_all"
#   SUM_BY_SUFFIX = "sum_by_suffix"
#   COMBINE_SNR_OUT = "combine_snr_out"
#   FC = "fc"
#   HIGHWAY = "highway"
#   RESIDUAL = "residual"
#   NONE = "None"
#   BOOL = "bool"
#   SCALAR_WEIGHT = "scalar_weight"
#   CONCATENATE = "concatenate"
#   SUM_ALL = "sum"
#   SUM_ALL = "SUM_ALL"
#   UNCERTAINTY_WEIGHTED = "UNCERTAINTY_WEIGHTED"
#   RNN = "RNN"
#   SNRNN = "SNRNN"
#   NONE = "none"
#   MLP = "MLP"
#   SNRMLP = "SNRMLP"
#   TRAIN = "train"
#   EVAL = "eval"
#   PREDICT = "predict"
#   def is_train(cls, mode: str) -> bool:
# 
# Path: ehr_prediction_modeling/models/reset_core.py
# class ResetCore(snt.RNNCore):
#   def __init__(self, core: RNNCore):
#   def _build(self, input_should_reset: Tuple[tf.Tensor, tf.Tensor],
#              state: tf.Tensor, **kwargs) -> Tuple[tf.Tensor, tf.Tensor]:
#   def wrapped_core(self) -> RNNCore:
#   def state_size(self) -> tf.TensorShape:
#   def output_size(self) -> tf.TensorShape:
#   def trainable_weights(self) -> List[tf.Tensor]:
#   def initial_state(self, *args, **kwargs) -> List[tf.Tensor]:


# Path: ehr_prediction_modeling/models/cell_factory.py
import functools
import typing
import sonnet as snt
import tensorflow.compat.v1 as tf
from ehr_prediction_modeling import types
from ehr_prediction_modeling.models import reset_core
from tensorflow import contrib as tf_contrib
  from ehr_prediction_modeling import configdict

              tf_contrib.rnn.UGRNNCell,
              num_units=config.ndim_lstm,
              activation=config.activation_fn),
  }

  if config.cell_type not in cell_construction_mapping:
    raise ValueError(f"Invalid cell_type: {config.cell_type}")

  cell = cell_construction_mapping[config.cell_type]()

  if config.cell_type in cells_not_allowing_wrappers:
    # Certain cells do not support wrappers.
    return cell

  if config.use_highway_connections:
    cell = tf_contrib.rnn.HighwayWrapper(cell)
  return cell


def init_deep_cell(
    cell_config: "configdict.ConfigDict") -> reset_core.ResetCore:
  """Initializes and wraps a cell in a ResetCore.

  Args:
    cell_config: Cell ConfigDict.

  Returns:
    An RNNCore object.
  """
  cell = _map_cell_type(cell_config)

  return reset_core.ResetCore(cell)