# Path: libml/data.py
# DATA_DIR = None
# _DATA_CACHE = None
# SAMPLES_PER_CLASS = [1, 2, 3, 4, 5, 10, 25, 100, 400]
# FLAGS = flags.FLAGS
#     DATA_DIR = FLAGS.data_dir or os.environ['ML_DATA']
# DATASETS = functools.partial(create_datasets, augment_module.augment_function)
# PAIR_DATASETS = functools.partial(create_datasets, augment_module.pair_augment_function)
# MANY_DATASETS = functools.partial(create_datasets, augment_module.many_augment_function)
# QUAD_DATASETS = functools.partial(create_datasets, augment_module.quad_augment_function)
# def _data_setup():
# def record_parse_mnist(serialized_example, image_shape=None):
# def record_parse(serialized_example, image_shape=None):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, data: tf.data.Dataset, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#     def from_files(cls, filenames: list, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#             def fetch_dataset(filename):
#     def empty_data(cls, image_shape, augment_fn: AugmentPair = None):
#         def _get_null_input(_):
#     def __getattr__(self, item):
#         def call_and_update(*args, **kwargs):
#     def parse(self):
#     def numpy_augment(self, *args, **kwargs):
#     def augment(self):
#     def memoize(self):
#         def tf_get(index, image_shape):
#             def get(index):
#     def __init__(self, name, train_labeled: DataSet, train_unlabeled: DataSet, test: DataSet, valid: DataSet,
#                  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=record_parse, do_memoize=False,
#                 nclass=10, colors=3, height=32, width=32):
#         def create():
# def create_datasets(augment_fn):
# class DataSet:
# class DataSets:
# 
# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
# def get_config():
# def setup_main():
# 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 get_gpu():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# def interleave(x, batch):
# def de_interleave(x, batch):
# def combine_dicts(*args):
# 
# Path: libml/utils.py
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
#         super(EasyDict, self).__init__(*args, **kwargs)
#         self.__dict__ = self


import functools
import json
import os.path
import shutil
import numpy as np
import tensorflow as tf
from absl import flags
from tqdm import trange, tqdm
from libml import data, utils
from libml.utils import EasyDict
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Training loop, checkpoint saving and loading, evaluation code."""



FLAGS = flags.FLAGS
flags.DEFINE_string('train_dir', './experiments',
                    'Folder where to save training data.')
flags.DEFINE_float('lr', 0.0001, 'Learning rate.')
flags.DEFINE_integer('batch', 64, 'Batch size.')
flags.DEFINE_integer('train_kimg', 1 << 14, 'Training duration in kibi-samples.')
flags.DEFINE_integer('report_kimg', 64, 'Report summary period in kibi-samples.')
flags.DEFINE_integer('save_kimg', 64, 'Save checkpoint period in kibi-samples.')
flags.DEFINE_integer('keep_ckpt', 50, 'Number of checkpoints to keep.')
flags.DEFINE_string('eval_ckpt', '', 'Checkpoint to evaluate. If provided, do not do training, just do eval.')
flags.DEFINE_string('rerun', '', 'A string to identify a run if running multiple ones with same parameters.')


class Model:
    def __init__(self, train_dir: str, dataset: data.DataSets, **kwargs):
        self.train_dir = os.path.join(train_dir, FLAGS.rerun, self.experiment_name(**kwargs))

# Path: libml/data.py
# DATA_DIR = None
# _DATA_CACHE = None
# SAMPLES_PER_CLASS = [1, 2, 3, 4, 5, 10, 25, 100, 400]
# FLAGS = flags.FLAGS
#     DATA_DIR = FLAGS.data_dir or os.environ['ML_DATA']
# DATASETS = functools.partial(create_datasets, augment_module.augment_function)
# PAIR_DATASETS = functools.partial(create_datasets, augment_module.pair_augment_function)
# MANY_DATASETS = functools.partial(create_datasets, augment_module.many_augment_function)
# QUAD_DATASETS = functools.partial(create_datasets, augment_module.quad_augment_function)
# def _data_setup():
# def record_parse_mnist(serialized_example, image_shape=None):
# def record_parse(serialized_example, image_shape=None):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, data: tf.data.Dataset, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#     def from_files(cls, filenames: list, augment_fn: AugmentPair, parse_fn=record_parse, image_shape=None):
#             def fetch_dataset(filename):
#     def empty_data(cls, image_shape, augment_fn: AugmentPair = None):
#         def _get_null_input(_):
#     def __getattr__(self, item):
#         def call_and_update(*args, **kwargs):
#     def parse(self):
#     def numpy_augment(self, *args, **kwargs):
#     def augment(self):
#     def memoize(self):
#         def tf_get(index, image_shape):
#             def get(index):
#     def __init__(self, name, train_labeled: DataSet, train_unlabeled: DataSet, test: DataSet, valid: DataSet,
#                  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=record_parse, do_memoize=False,
#                 nclass=10, colors=3, height=32, width=32):
#         def create():
# def create_datasets(augment_fn):
# class DataSet:
# class DataSets:
# 
# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
# def get_config():
# def setup_main():
# 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 get_gpu():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# def interleave(x, batch):
# def de_interleave(x, batch):
# def combine_dicts(*args):
# 
# Path: libml/utils.py
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
#         super(EasyDict, self).__init__(*args, **kwargs)
#         self.__dict__ = self


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

# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Training loop, checkpoint saving and loading, evaluation code."""



FLAGS = flags.FLAGS
flags.DEFINE_string('train_dir', './experiments',
                    'Folder where to save training data.')
flags.DEFINE_float('lr', 0.0001, 'Learning rate.')
flags.DEFINE_integer('batch', 64, 'Batch size.')
flags.DEFINE_integer('train_kimg', 1 << 14, 'Training duration in kibi-samples.')
flags.DEFINE_integer('report_kimg', 64, 'Report summary period in kibi-samples.')
flags.DEFINE_integer('save_kimg', 64, 'Save checkpoint period in kibi-samples.')
flags.DEFINE_integer('keep_ckpt', 50, 'Number of checkpoints to keep.')
flags.DEFINE_string('eval_ckpt', '', 'Checkpoint to evaluate. If provided, do not do training, just do eval.')
flags.DEFINE_string('rerun', '', 'A string to identify a run if running multiple ones with same parameters.')


class Model:
    def __init__(self, train_dir: str, dataset: data.DataSets, **kwargs):
        self.train_dir = os.path.join(train_dir, FLAGS.rerun, self.experiment_name(**kwargs))

        self.params = EasyDict(kwargs)