# 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


import copy
import functools
import jax
import jax.numpy as jnp
import optax
from absl import logging
from clu import metric_writers
from deluca.lung.utils.data.breath_dataset import get_shuffled_and_batched_data
from flax.metrics import tensorboard
    logging.info("decay_steps: %s", str(decay_steps))
    cosine_scheduler_fn = optax.cosine_decay_schedule(
        init_value=optimizer_params["learning_rate"], decay_steps=decay_steps)
    optimizer_params["learning_rate"] = cosine_scheduler_fn
    logging.info("optimizer_params: %s", str(optimizer_params))
    optim = optimizer(**optimizer_params)
  optim_state = optim.init(model)

  loop_over_loader_partial = functools.partial(
      loop_over_loader, optim=optim, rollout_fn=rollout, scheduler=scheduler)

  # Tensorboard writer
  if use_tensorboard:
    config = copy.deepcopy(model.default_model_parameters)
    del config["activation_fn"]
    config["activation_fn_name"] = activation_fn_name

    if mode == "train":
      file_name = str(config)
    write_path = tb_dir + file_name
    summary_writer = metric_writers.create_default_writer(
        logdir=write_path, just_logging=jax.process_index() != 0)
    summary_writer = tensorboard.SummaryWriter(write_path)
    summary_writer.write_hparams(dict(config))

  # Main Training Loop
  prng_key = jax.random.PRNGKey(0)
  for epoch in range(epochs + 1):
    if epoch % 10 == 0:
      logging.info("epoch: %s", str(epoch))

# 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/lung/utils/scripts/train_simulator.py
import copy
import functools
import jax
import jax.numpy as jnp
import optax
from absl import logging
from clu import metric_writers
from deluca.lung.utils.data.breath_dataset import get_shuffled_and_batched_data
from flax.metrics import tensorboard

    logging.info("decay_steps: %s", str(decay_steps))
    cosine_scheduler_fn = optax.cosine_decay_schedule(
        init_value=optimizer_params["learning_rate"], decay_steps=decay_steps)
    optimizer_params["learning_rate"] = cosine_scheduler_fn
    logging.info("optimizer_params: %s", str(optimizer_params))
    optim = optimizer(**optimizer_params)
  optim_state = optim.init(model)

  loop_over_loader_partial = functools.partial(
      loop_over_loader, optim=optim, rollout_fn=rollout, scheduler=scheduler)

  # Tensorboard writer
  if use_tensorboard:
    config = copy.deepcopy(model.default_model_parameters)
    del config["activation_fn"]
    config["activation_fn_name"] = activation_fn_name

    if mode == "train":
      file_name = str(config)
    write_path = tb_dir + file_name
    summary_writer = metric_writers.create_default_writer(
        logdir=write_path, just_logging=jax.process_index() != 0)
    summary_writer = tensorboard.SummaryWriter(write_path)
    summary_writer.write_hparams(dict(config))

  # Main Training Loop
  prng_key = jax.random.PRNGKey(0)
  for epoch in range(epochs + 1):
    if epoch % 10 == 0:
      logging.info("epoch: %s", str(epoch))

    X, y, prng_key = get_shuffled_and_batched_data(dataset, batch_size,