# Path: fedjax/core/serialization.py
# def save_state(state, path):
# def load_state(path):
# def _ndarray_to_bytes(arr):
# def _dtype_from_name(name):
# def _ndarray_from_bytes(data):
# def _bytes_ndarray_to_bytes(x):
# def _object_ndarray_from_bytes(data):
# def _msgpack_ext_pack(x):
# def _msgpack_ext_unpack(code, data):
# def msgpack_serialize(pytree):
# def msgpack_deserialize(encoded_pytree):
# class _MsgpackExtType(enum.IntEnum):
# 
# Path: fedjax/core/util.py
# def safe_div(a, b):
#   def __getattr__(self, _):
# def import_tf():
# class FakeTensorFlowModule:


import os.path
import re
from typing import Any, List, Optional, Tuple
from fedjax.core import serialization
from fedjax.core import util
  pattern = base_path + r'[0-9]{8}$'
  checkpoint_paths = []
  for path in tf.io.gfile.glob(base_path + '*'):
    if re.match(pattern, path):
      checkpoint_paths.append(path)

  def sort_key(path):
    return int(path.split(base_path)[-1])

  return sorted(checkpoint_paths, key=sort_key)


def load_latest_checkpoint(root_dir: str) -> Optional[Tuple[Any, int]]:
  """Loads latest checkpoint and round number."""
  base_path = os.path.join(root_dir, _CHECKPOINT_PREFIX)
  all_checkpoint_paths = _get_checkpoint_paths(base_path)
  if all_checkpoint_paths:
    latest_checkpoint_path = all_checkpoint_paths[-1]
    latest_round_num = int(latest_checkpoint_path.split(base_path)[-1])
    latest_state = serialization.load_state(latest_checkpoint_path)
    return latest_state, latest_round_num


def save_checkpoint(root_dir: str,
                    state: Any,
                    round_num: int = 0,
                    keep: int = 1):
  """Saves checkpoint and cleans up old checkpoints."""
  base_path = os.path.join(root_dir, _CHECKPOINT_PREFIX)
  checkpoint_path = f'{base_path}{round_num:08d}'

# Path: fedjax/core/serialization.py
# def save_state(state, path):
# def load_state(path):
# def _ndarray_to_bytes(arr):
# def _dtype_from_name(name):
# def _ndarray_from_bytes(data):
# def _bytes_ndarray_to_bytes(x):
# def _object_ndarray_from_bytes(data):
# def _msgpack_ext_pack(x):
# def _msgpack_ext_unpack(code, data):
# def msgpack_serialize(pytree):
# def msgpack_deserialize(encoded_pytree):
# class _MsgpackExtType(enum.IntEnum):
# 
# Path: fedjax/core/util.py
# def safe_div(a, b):
#   def __getattr__(self, _):
# def import_tf():
# class FakeTensorFlowModule:


# Path: fedjax/training/checkpoint.py
import os.path
import re
from typing import Any, List, Optional, Tuple
from fedjax.core import serialization
from fedjax.core import util

  pattern = base_path + r'[0-9]{8}$'
  checkpoint_paths = []
  for path in tf.io.gfile.glob(base_path + '*'):
    if re.match(pattern, path):
      checkpoint_paths.append(path)

  def sort_key(path):
    return int(path.split(base_path)[-1])

  return sorted(checkpoint_paths, key=sort_key)


def load_latest_checkpoint(root_dir: str) -> Optional[Tuple[Any, int]]:
  """Loads latest checkpoint and round number."""
  base_path = os.path.join(root_dir, _CHECKPOINT_PREFIX)
  all_checkpoint_paths = _get_checkpoint_paths(base_path)
  if all_checkpoint_paths:
    latest_checkpoint_path = all_checkpoint_paths[-1]
    latest_round_num = int(latest_checkpoint_path.split(base_path)[-1])
    latest_state = serialization.load_state(latest_checkpoint_path)
    return latest_state, latest_round_num


def save_checkpoint(root_dir: str,
                    state: Any,
                    round_num: int = 0,
                    keep: int = 1):
  """Saves checkpoint and cleans up old checkpoints."""
  base_path = os.path.join(root_dir, _CHECKPOINT_PREFIX)
  checkpoint_path = f'{base_path}{round_num:08d}'

  serialization.save_state(state, checkpoint_path)