import tensorflow as tf
from deep_heatmaps_model_primary_valid import DeepHeatmapsModel
import os

# data_dir ='/mnt/External1/Yarden/deep_face_heatmaps/data/conventional_landmark_detection_dataset/'
data_dir = '/Users/arik/Dropbox/a_mac_thesis/face_heatmap_networks/conventional_landmark_detection_dataset/'
pre_train_path = 'saved_models/0.01/model/deep_heatmaps-50000'
output_dir = os.getcwd()

flags = tf.app.flags

flags.DEFINE_string('mode', 'TRAIN', "'TRAIN' or 'TEST'")

# define paths
flags.DEFINE_string('save_model_path', 'model', "directory for saving the model")
flags.DEFINE_string('save_sample_path', 'sample', "directory for saving the sampled images")
flags.DEFINE_string('save_log_path', 'logs', "directory for saving the log file")
flags.DEFINE_string('img_path', data_dir, "data directory")
flags.DEFINE_string('test_model_path', 'model/deep_heatmaps-5', "saved model to test")
flags.DEFINE_string('test_data','full', 'test set to use full/common/challenging/test/art')

# pretrain parameters
flags.DEFINE_string('pre_train_path', pre_train_path, 'pretrained model path')
flags.DEFINE_bool('load_pretrain', False, "load pretrained weight?")

flags.DEFINE_integer('image_size', 256, "image size")
flags.DEFINE_integer('c_dim', 3, "color channels")
flags.DEFINE_integer('num_landmarks', 68, "number of face landmarks")

# optimization parameters
flags.DEFINE_integer('train_iter', 100000, 'maximum training iterations')
flags.DEFINE_integer('batch_size', 10, "batch_size")
flags.DEFINE_float('learning_rate', 1e-6, "initial learning rate")
flags.DEFINE_float('momentum', 0.95, 'optimizer momentum')
flags.DEFINE_integer('step', 100000, 'step for lr decay')
flags.DEFINE_float('gamma', 0.1, 'exponential base for lr decay')

# augmentation parameters
flags.DEFINE_bool('augment_basic', True,"use basic augmentation?")
flags.DEFINE_bool('augment_texture', False,"use artistic texture augmentation?")
flags.DEFINE_bool('augment_geom', False,"use artistic geometric augmentation?")
flags.DEFINE_integer('basic_start', 0,  'min epoch to start basic augmentation')
flags.DEFINE_float('p_texture', 0., 'initial probability of artistic texture augmentation')
flags.DEFINE_float('p_geom', 0., 'initial probability of artistic geometric augmentation')
flags.DEFINE_integer('artistic_step', 10, 'increase probability of artistic augmentation every X epochs')
flags.DEFINE_integer('artistic_start', 0, 'min epoch to start artistic augmentation')

# directory of test
flags.DEFINE_string('output_dir', output_dir, "directory for saving test")

FLAGS = flags.FLAGS

if not os.path.exists(FLAGS.output_dir):
    os.mkdir(FLAGS.output_dir)


def main(_):

    save_model_path = os.path.join(FLAGS.output_dir, FLAGS.save_model_path)
    save_sample_path = os.path.join(FLAGS.output_dir, FLAGS.save_sample_path)
    save_log_path = os.path.join(FLAGS.output_dir, FLAGS.save_log_path)

    # create directories if not exist
    if not os.path.exists(save_model_path):
        os.mkdir(save_model_path)
    if not os.path.exists(save_sample_path):
        os.mkdir(save_sample_path)
    if not os.path.exists(save_log_path):
        os.mkdir(save_log_path)

    model = DeepHeatmapsModel(mode=FLAGS.mode, train_iter=FLAGS.train_iter, learning_rate=FLAGS.learning_rate,
                              momentum=FLAGS.momentum, step=FLAGS.step, gamma=FLAGS.gamma, batch_size=FLAGS.batch_size,
                              image_size=FLAGS.image_size, c_dim=FLAGS.c_dim, num_landmarks=FLAGS.num_landmarks,
                              augment_basic=FLAGS.augment_basic, basic_start=FLAGS.basic_start,
                              augment_texture=FLAGS.augment_texture, p_texture=FLAGS.p_texture,
                              augment_geom=FLAGS.augment_geom, p_geom=FLAGS.p_geom,
                              artistic_start=FLAGS.artistic_start, artistic_step=FLAGS.artistic_step,
                              img_path=FLAGS.img_path, save_log_path=save_log_path,
                              save_sample_path=save_sample_path, save_model_path=save_model_path,
                              test_data=FLAGS.test_data, test_model_path=FLAGS.test_model_path,
                              load_pretrain=FLAGS.load_pretrain, pre_train_path=FLAGS.pre_train_path)

    if FLAGS.mode == 'TRAIN':
        model.train()
    else:
        model.eval()

if __name__ == '__main__':
    tf.app.run()