# Path: init2winit/checkpoint.py
# FLAGS = flags.FLAGS
# def replicate_checkpoint(
#     latest,
#     pytree_keys: Sequence[str],
#     replicate=True):
# def replicate_and_maybe_restore_latest_checkpoint(
#     unreplicated_optimizer_state,
#     unreplicated_params,
#     unreplicated_batch_stats,
#     unreplicated_training_metrics_grabber,
#     train_dir):
# def save_unreplicated_checkpoint_background(
#     train_dir,
#     optimizer_state,
#     params,
#     batch_stats,
#     training_metrics_grabber,
#     global_step,
#     preemption_count,
#     sum_train_cost,
#     max_to_keep=1):
# def _save_checkpoint_background_catch_error(*args, **kwargs):
# def wait_for_checkpoint_save():
# def save_checkpoint_background(*args, **kwargs):
# def save_checkpoint(train_dir,
#                     step,
#                     state,
#                     prefix='ckpt_',
#                     max_to_keep=None):
# def load_checkpoint(
#     checkpoint_path,
#     target=None,
#     prefix='ckpt_'):
# def load_latest_checkpoint(train_dir, target=None, prefix='ckpt_'):


import os
import pathlib
import jax
import optax
import sacrebleu
from init2winit import checkpoint
from tensorflow.io import gfile
  return os.path.join(output_dir, 'bleu_' + eval_split + '_' + str(ckpt_step))


def load_evals(ckpt_dir, ckpt_step, eval_split):
  """Loads results if already available, else return None."""
  ckpt_eval_fpath = get_eval_fpath(ckpt_dir, ckpt_step, eval_split)
  if not gfile.exists(ckpt_eval_fpath):
    return None
  else:
    with gfile.GFile(ckpt_eval_fpath, 'r') as f:
      bleu_score = f.readlines()[-1]
      return float(bleu_score.strip())


def save_evals(ckpt_dir, ckpt_step, eval_split, bleu_score):
  ckpt_eval_fpath = get_eval_fpath(ckpt_dir, ckpt_step, eval_split)
  with gfile.GFile(ckpt_eval_fpath, 'w') as f:
    f.write(str(bleu_score))


def _load_checkpoint(checkpoint_path, params, optimizer_state, batch_stats,
                     replicate=True):
  """Load model (and batch stats) from checkpoint."""
  target = dict(
      params=params,
      optimizer_state=optimizer_state,
      batch_stats=batch_stats,
      global_step=-1,
      preemption_count=0,
      sum_train_cost=0.0)

# Path: init2winit/checkpoint.py
# FLAGS = flags.FLAGS
# def replicate_checkpoint(
#     latest,
#     pytree_keys: Sequence[str],
#     replicate=True):
# def replicate_and_maybe_restore_latest_checkpoint(
#     unreplicated_optimizer_state,
#     unreplicated_params,
#     unreplicated_batch_stats,
#     unreplicated_training_metrics_grabber,
#     train_dir):
# def save_unreplicated_checkpoint_background(
#     train_dir,
#     optimizer_state,
#     params,
#     batch_stats,
#     training_metrics_grabber,
#     global_step,
#     preemption_count,
#     sum_train_cost,
#     max_to_keep=1):
# def _save_checkpoint_background_catch_error(*args, **kwargs):
# def wait_for_checkpoint_save():
# def save_checkpoint_background(*args, **kwargs):
# def save_checkpoint(train_dir,
#                     step,
#                     state,
#                     prefix='ckpt_',
#                     max_to_keep=None):
# def load_checkpoint(
#     checkpoint_path,
#     target=None,
#     prefix='ckpt_'):
# def load_latest_checkpoint(train_dir, target=None, prefix='ckpt_'):


# Path: init2winit/mt_eval/eval_utils.py
import os
import pathlib
import jax
import optax
import sacrebleu
from init2winit import checkpoint
from tensorflow.io import gfile

  return os.path.join(output_dir, 'bleu_' + eval_split + '_' + str(ckpt_step))


def load_evals(ckpt_dir, ckpt_step, eval_split):
  """Loads results if already available, else return None."""
  ckpt_eval_fpath = get_eval_fpath(ckpt_dir, ckpt_step, eval_split)
  if not gfile.exists(ckpt_eval_fpath):
    return None
  else:
    with gfile.GFile(ckpt_eval_fpath, 'r') as f:
      bleu_score = f.readlines()[-1]
      return float(bleu_score.strip())


def save_evals(ckpt_dir, ckpt_step, eval_split, bleu_score):
  ckpt_eval_fpath = get_eval_fpath(ckpt_dir, ckpt_step, eval_split)
  with gfile.GFile(ckpt_eval_fpath, 'w') as f:
    f.write(str(bleu_score))


def _load_checkpoint(checkpoint_path, params, optimizer_state, batch_stats,
                     replicate=True):
  """Load model (and batch stats) from checkpoint."""
  target = dict(
      params=params,
      optimizer_state=optimizer_state,
      batch_stats=batch_stats,
      global_step=-1,
      preemption_count=0,
      sum_train_cost=0.0)

  ckpt = checkpoint.load_checkpoint(