# Path: libml/data.py
# _DATA_CACHE = None
# DATA_DIR = os.environ['ML_DATA']
# FLAGS = flags.FLAGS
# DATASETS = {}
# def record_parse(serialized_example):
# def default_parse(dataset: tf.data.Dataset, parse_fn=record_parse) -> tf.data.Dataset:
# def dataset(filenames: list) -> tf.data.Dataset:
# def memoize(dataset: tf.data.Dataset) -> tf.data.Dataset:
#     def tf_get(index):
#         def get(index):
# def augment_mirror(x):
# def augment_shift(x, w):
# def augment_noise(x, std):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, name, train_labeled, train_unlabeled, test, valid, eval_labeled, eval_unlabeled,
#                  height=32, width=32, colors=3, nclass=10, mean=0, std=1, p_labeled=None, p_unlabeled=None):
#     def creator(cls, name, seed, label, valid, augment, parse_fn=default_parse, do_memoize=True, colors=3,
#                 nclass=10, height=32, width=32, name_suffix=''):
#         def create():
# class DataSet:
# 
# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# def get_config():
# def setup_tf():
# def smart_shape(x):
# def ilog2(x):
# def find_latest_checkpoint(dir, glob_term='model.ckpt-*.meta'):
# def get_latest_global_step(dir):
# def get_latest_global_step_in_subdir(dir):
# def getter_ema(ema, getter, name, *args, **kwargs):
# def model_vars(scope=None):
# def gpu(x):
# def get_available_gpus():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):


import json
import os.path
import shutil
import numpy as np
import tensorflow as tf
from absl import flags
from easydict import EasyDict
from tqdm import trange
from libml import data, utils
    def train_print(self, text):
        self.tmp.print_queue.append(text)

    def _create_initial_files(self):
        for dir in (self.checkpoint_dir, self.arg_dir):
            if not os.path.exists(dir):
                os.makedirs(dir)
        self.save_args()

    def _reset_files(self):
        shutil.rmtree(self.train_dir)
        self._create_initial_files()

    def save_args(self, **extra_params):
        with open(os.path.join(self.arg_dir, 'args.json'), 'w') as f:
            json.dump({**self.params, **extra_params}, f, sort_keys=True, indent=4)

    @classmethod
    def load(cls, train_dir):
        with open(os.path.join(train_dir, 'args/args.json'), 'r') as f:
            params = json.load(f)
        instance = cls(train_dir=train_dir, **params)
        instance.train_dir = train_dir
        return instance

    def experiment_name(self, **kwargs):
        args = [x + str(y) for x, y in sorted(kwargs.items())]
        return '_'.join([self.__class__.__name__] + args)

    def eval_mode(self, ckpt=None):

# Path: libml/data.py
# _DATA_CACHE = None
# DATA_DIR = os.environ['ML_DATA']
# FLAGS = flags.FLAGS
# DATASETS = {}
# def record_parse(serialized_example):
# def default_parse(dataset: tf.data.Dataset, parse_fn=record_parse) -> tf.data.Dataset:
# def dataset(filenames: list) -> tf.data.Dataset:
# def memoize(dataset: tf.data.Dataset) -> tf.data.Dataset:
#     def tf_get(index):
#         def get(index):
# def augment_mirror(x):
# def augment_shift(x, w):
# def augment_noise(x, std):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, name, train_labeled, train_unlabeled, test, valid, eval_labeled, eval_unlabeled,
#                  height=32, width=32, colors=3, nclass=10, mean=0, std=1, p_labeled=None, p_unlabeled=None):
#     def creator(cls, name, seed, label, valid, augment, parse_fn=default_parse, do_memoize=True, colors=3,
#                 nclass=10, height=32, width=32, name_suffix=''):
#         def create():
# class DataSet:
# 
# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# def get_config():
# def setup_tf():
# def smart_shape(x):
# def ilog2(x):
# def find_latest_checkpoint(dir, glob_term='model.ckpt-*.meta'):
# def get_latest_global_step(dir):
# def get_latest_global_step_in_subdir(dir):
# def getter_ema(ema, getter, name, *args, **kwargs):
# def model_vars(scope=None):
# def gpu(x):
# def get_available_gpus():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):


# Path: libml/train.py
import json
import os.path
import shutil
import numpy as np
import tensorflow as tf
from absl import flags
from easydict import EasyDict
from tqdm import trange
from libml import data, utils

    def train_print(self, text):
        self.tmp.print_queue.append(text)

    def _create_initial_files(self):
        for dir in (self.checkpoint_dir, self.arg_dir):
            if not os.path.exists(dir):
                os.makedirs(dir)
        self.save_args()

    def _reset_files(self):
        shutil.rmtree(self.train_dir)
        self._create_initial_files()

    def save_args(self, **extra_params):
        with open(os.path.join(self.arg_dir, 'args.json'), 'w') as f:
            json.dump({**self.params, **extra_params}, f, sort_keys=True, indent=4)

    @classmethod
    def load(cls, train_dir):
        with open(os.path.join(train_dir, 'args/args.json'), 'r') as f:
            params = json.load(f)
        instance = cls(train_dir=train_dir, **params)
        instance.train_dir = train_dir
        return instance

    def experiment_name(self, **kwargs):
        args = [x + str(y) for x, y in sorted(kwargs.items())]
        return '_'.join([self.__class__.__name__] + args)

    def eval_mode(self, ckpt=None):

        self.session = tf.Session(config=utils.get_config())