# Path: studio/fs_tracker.py
# STUDIOML_EXPERIMENT = 'STUDIOML_EXPERIMENT'
# STUDIOML_HOME = 'STUDIOML_HOME'
# STUDIOML_ARTIFACT_MAPPING = 'STUDIOML_ARTIFACT_MAPPING'
# def get_experiment_key():
# def get_studio_home():
# def setup_experiment(env, experiment, clean=True):
# def get_artifact(tag):
# def get_artifacts():
# def get_model_directory(experiment_name=None):
# def get_artifact_cache(tag, experiment_name=None):
# def get_blob_cache(blobkey):
# def get_model_directory(experiment_name=None):
# def _get_artifact_mapping_path(experiment_name=None):
# def _get_experiment_key(experiment):
# def _setup_model_directory(experiment_name, clean=False):
# def get_queue_directory():
# def get_tensorboard_dir(experiment_name=None):


import sys
import tensorflow as tf
from studio import fs_tracker
from keras.layers import Dense, Flatten, Conv2D, BatchNormalization
from keras.models import Sequential
from keras.datasets import cifar10
from keras.utils import to_categorical
from keras import optimizers
from keras.callbacks import ModelCheckpoint, TensorBoard
from keras import backend as backend
model.add(Conv2D(32, (3, 3), activation='relu'))
model.add(BatchNormalization())
model.add(Conv2D(64, (5, 5), activation='relu'))
model.add(Conv2D(64, (3, 3), activation='relu'))
model.add(BatchNormalization())
model.add(Conv2D(128, (3, 3), activation='relu'))
model.add(Conv2D(128, (3, 3), activation='relu'))
model.add(BatchNormalization())
model.add(Flatten())
model.add(Dense(10, activation='softmax'))
model.summary()


batch_size = 128
no_epochs = int(sys.argv[1]) if len(sys.argv) > 1 else 40
lr = 0.01

print('learning rate = {}'.format(lr))
print('batch size = {}'.format(batch_size))
print('no_epochs = {}'.format(no_epochs))

model.compile(loss='categorical_crossentropy', optimizer=optimizers.SGD(lr=lr),
              metrics=['accuracy'])


checkpointer = ModelCheckpoint(
    fs_tracker.get_model_directory() +
    '/checkpoint.{epoch:02d}-{val_loss:.2f}.hdf')



# Path: studio/fs_tracker.py
# STUDIOML_EXPERIMENT = 'STUDIOML_EXPERIMENT'
# STUDIOML_HOME = 'STUDIOML_HOME'
# STUDIOML_ARTIFACT_MAPPING = 'STUDIOML_ARTIFACT_MAPPING'
# def get_experiment_key():
# def get_studio_home():
# def setup_experiment(env, experiment, clean=True):
# def get_artifact(tag):
# def get_artifacts():
# def get_model_directory(experiment_name=None):
# def get_artifact_cache(tag, experiment_name=None):
# def get_blob_cache(blobkey):
# def get_model_directory(experiment_name=None):
# def _get_artifact_mapping_path(experiment_name=None):
# def _get_experiment_key(experiment):
# def _setup_model_directory(experiment_name, clean=False):
# def get_queue_directory():
# def get_tensorboard_dir(experiment_name=None):


# Path: examples/keras/train_cifar10.py
import sys
import tensorflow as tf
from studio import fs_tracker
from keras.layers import Dense, Flatten, Conv2D, BatchNormalization
from keras.models import Sequential
from keras.datasets import cifar10
from keras.utils import to_categorical
from keras import optimizers
from keras.callbacks import ModelCheckpoint, TensorBoard
from keras import backend as backend

model.add(Conv2D(32, (3, 3), activation='relu'))
model.add(BatchNormalization())
model.add(Conv2D(64, (5, 5), activation='relu'))
model.add(Conv2D(64, (3, 3), activation='relu'))
model.add(BatchNormalization())
model.add(Conv2D(128, (3, 3), activation='relu'))
model.add(Conv2D(128, (3, 3), activation='relu'))
model.add(BatchNormalization())
model.add(Flatten())
model.add(Dense(10, activation='softmax'))
model.summary()


batch_size = 128
no_epochs = int(sys.argv[1]) if len(sys.argv) > 1 else 40
lr = 0.01

print('learning rate = {}'.format(lr))
print('batch size = {}'.format(batch_size))
print('no_epochs = {}'.format(no_epochs))

model.compile(loss='categorical_crossentropy', optimizer=optimizers.SGD(lr=lr),
              metrics=['accuracy'])


checkpointer = ModelCheckpoint(
    fs_tracker.get_model_directory() +
    '/checkpoint.{epoch:02d}-{val_loss:.2f}.hdf')



tbcallback = TensorBoard(log_dir=fs_tracker.get_tensorboard_dir(),