# Path: deluca/lung/utils/data/breath_dataset.py
# 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
# 
# Path: deluca/lung/utils/data/breath_dataset.py
# def get_initial_pressure(dataset, key):
#   return jnp.array([dataset.p_normalizer(p[0]) for u, p in dataset.splits[key]
#                    ]).mean().item()
# 
# Path: deluca/lung/utils/data/breath_dataset.py
# def get_shuffled_and_batched_data(dataset, batch_size, key, prng_key):
#   """function to shuffle and batch data."""
#   x, y = dataset.data[key]
#   x = jax.random.permutation(prng_key, x)
#   y = jax.random.permutation(prng_key, y)
#   prng_key, _ = jax.random.split(prng_key)
#   num_batches = x.shape[0] // batch_size
#   trunc_len = num_batches * batch_size
#   trunc_x = x[:trunc_len]
#   trunc_y = y[:trunc_len]
#   batched_x = jnp.reshape(trunc_x, (num_batches, batch_size, trunc_x.shape[1]))
#   batched_y = jnp.reshape(trunc_y, (num_batches, batch_size, trunc_y.shape[1]))
#   return batched_x, batched_y, prng_key


from absl.testing import absltest
from deluca.lung.utils.data.breath_dataset import BreathDataset
from deluca.lung.utils.data.breath_dataset import get_initial_pressure
from deluca.lung.utils.data.breath_dataset import get_shuffled_and_batched_data
import chex
import jax
import jax.numpy as jnp
  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])
    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))

  def test_get_shuffled_and_batched_data(self):
    prng_key = jax.random.PRNGKey(0)
    batch_size = 2

# Path: deluca/lung/utils/data/breath_dataset.py
# 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
# 
# Path: deluca/lung/utils/data/breath_dataset.py
# def get_initial_pressure(dataset, key):
#   return jnp.array([dataset.p_normalizer(p[0]) for u, p in dataset.splits[key]
#                    ]).mean().item()
# 
# Path: deluca/lung/utils/data/breath_dataset.py
# def get_shuffled_and_batched_data(dataset, batch_size, key, prng_key):
#   """function to shuffle and batch data."""
#   x, y = dataset.data[key]
#   x = jax.random.permutation(prng_key, x)
#   y = jax.random.permutation(prng_key, y)
#   prng_key, _ = jax.random.split(prng_key)
#   num_batches = x.shape[0] // batch_size
#   trunc_len = num_batches * batch_size
#   trunc_x = x[:trunc_len]
#   trunc_y = y[:trunc_len]
#   batched_x = jnp.reshape(trunc_x, (num_batches, batch_size, trunc_x.shape[1]))
#   batched_y = jnp.reshape(trunc_y, (num_batches, batch_size, trunc_y.shape[1]))
#   return batched_x, batched_y, prng_key


# Path: deluca/tests/lung/utils/data/breath_dataset_test.py
from absl.testing import absltest
from deluca.lung.utils.data.breath_dataset import BreathDataset
from deluca.lung.utils.data.breath_dataset import get_initial_pressure
from deluca.lung.utils.data.breath_dataset import get_shuffled_and_batched_data
import chex
import jax
import jax.numpy as jnp

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

  def test_get_shuffled_and_batched_data(self):
    prng_key = jax.random.PRNGKey(0)
    batch_size = 2

    x, y, prng_key = get_shuffled_and_batched_data(