import datetime
import os
import time


import numpy as np
import tensorflow as tf
from tensorflow.keras.applications.resnet50 import ResNet50,preprocess_input
from tensorflow.keras.callbacks import TensorBoard
from tensorflow.keras.layers import GlobalAveragePooling2D,Dense
from tensorflow.keras.models import Model, Sequential
from skimage.transform import resize
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.python.keras.optimizers import Adam

def random_crop(img, random_crop_size):
    # Note: image_data_format is 'channel_last'
    assert img.shape[2] == 3
    height, width = img.shape[0], img.shape[1]
    dy, dx = random_crop_size
    x = np.random.randint(0, width - dx + 1)
    y = np.random.randint(0, height - dy + 1)
    return img[y:(y+dy), x:(x+dx), :]


def crop_generator(batches, crop_length):
    """
    Take as input a Keras ImageGen (Iterator) and generate random
    crops from the image batches generated by the original iterator.
    """

    while True:
        batch_x, batch_y = next(batches)
        batch_crops = np.zeros((batch_x.shape[0], crop_length, crop_length, 3))
        for i in range(batch_x.shape[0]):
            batch_crops[i] = random_crop(batch_x[i], (crop_length, crop_length))
        yield (batch_crops, batch_y)

DATA_DIR = "../../annotations_landmarks_clean_struct"
HEIGHT = 250
WIDTH = 250
BATCH_SIZE = 16
NUM_CLASSES=585

datagen =  ImageDataGenerator(preprocessing_function=preprocess_input,validation_split=0.2,rotation_range=40,
width_shift_range=0.2,
height_shift_range=0.2,
shear_range=0.2,
zoom_range=0.2,
horizontal_flip=True,)

train_generator = datagen.flow_from_directory(DATA_DIR, 
                                                    target_size=(HEIGHT, WIDTH), 
						    interpolation='bicubic',
                                                    batch_size=BATCH_SIZE,
						    class_mode='categorical',
                                                    shuffle=True,
						    subset='training',
                                                    color_mode='rgb')

valid_generator = datagen.flow_from_directory(DATA_DIR,
                                                    target_size=(HEIGHT, WIDTH),
						    interpolation='bicubic',
                                                    batch_size=BATCH_SIZE,
						    class_mode='categorical',
                                                    shuffle=False,
						    subset='validation',
                                                    color_mode='rgb')

train_crops = crop_generator(train_generator, 224)
valid_crops = crop_generator(valid_generator, 224)

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

def timestamp():
  return datetime.datetime.fromtimestamp(time.time()).strftime('%y%m%d-%H%M%S')


def download():
  """ Download Cifar10 and Resnet50. """
  ResNet50(weights='imagenet', include_top=False)

def extract_train():
  """ Extract and train in the same model. """
  resnet50 = ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
  
  fc = GlobalAveragePooling2D()(resnet50.output)
  fc = Dense(NUM_CLASSES, activation='softmax', name='fc')(fc)
  model = Model(resnet50.input, fc)

  for layer in resnet50.layers:
    layer.trainable = False

  print('RestNet50-FC-1024-585')
  model.summary(line_length=100)
  model.compile(
    optimizer=Adam(lr=2e-5),
    loss='categorical_crossentropy',
    metrics=['accuracy']
  )

  tensorboard = TensorBoard(
    log_dir='./graphs/transfer_learning' + timestamp(), histogram_freq=0, write_graph=True, write_images=False)

  model.fit_generator(train_crops,steps_per_epoch=train_generator.samples//BATCH_SIZE, epochs=30,validation_data = valid_crops,validation_steps = valid_generator.samples//BATCH_SIZE, callbacks=[tensorboard])

  for layer in resnet50.layers:
    layer.trainable = True

  model.compile(
    optimizer=Adam(lr=1e-5),
    loss='categorical_crossentropy',
    metrics=['accuracy']
  )


  model.fit_generator(train_crops,steps_per_epoch=train_generator.samples//BATCH_SIZE, epochs=30,validation_data = valid_crops,validation_steps = valid_generator.samples//BATCH_SIZE, callbacks=[tensorboard])
  model.save('finetunned.h5')
  #evaluation = model.evaluate(X_tst, Y_tst)
  #print(evaluation)


def main(args):
  err_msg = 'Unknown function, options: download, extract_train'
  if len(args) > 1:
    func_name = args[1]
    if func_name == 'download':
      download()
    elif func_name == 'extract_train':
      extract_train()
    else:
      print(err_msg)
  else:
    print(err_msg)
  return 0


if __name__ == '__main__':
  import sys
  sys.exit(main(sys.argv))
