import tensorflow as tf

import numpy as np

import itertools

import sys

import os

from tensorflow.keras import layers

from tensorflow.keras import backend as K

# for saving models and csvlogger

import datetime

print("TF version: ", tf.version.VERSION)

tf.keras.backend.clear_session()  # For easy reset of notebook state.

#  config = tf.ConfigProto()

#  config.gpu_options.allow_growth = True  # dynamically grow the memory used on the GPU

#  sess = tf.Session(config=config)

#  K.set_session(sess)  # set this TensorFlow session as the default session for Keras

# dimensions of our images.

img_width, img_height = 150, 150

# IMAGE_SIZE    = (100, 100)

# CROP_LENGTH   = 84

if len(sys.argv) != 3:

    print('Error: pls provide train path and validation path')

    exit(0)

train_data_dir = sys.argv[1]

validation_data_dir = sys.argv[2]

nb_classes = 4

nb_train_samples = 0

nb_validation_samples = 0

nb_sample_per_class = []

nb_val_sample_per_class = []

folders = ['1', '2', '3', '4']

for folder in folders:

    num_tr = len(os.listdir(os.path.join(train_data_dir, folder)))

    nb_train_samples += num_tr

    nb_sample_per_class.append(num_tr)

for folder in folders:

    num_val = len(os.listdir(os.path.join(validation_data_dir, folder)))

    nb_validation_samples += num_val

    nb_val_sample_per_class.append(num_val)

# data_folder = train_data_dir.split(os.sep)

# data_folder = [e for e in data_folder if e != '']

# data_folder = data_folder[-2]

# print(data_folder)

print("\nnb_train_samples: ", nb_train_samples)

print("\nnb_validation_samples: ", nb_validation_samples)

print("\nnb_sample_per_class: ", nb_sample_per_class)

print("\nnb_val_sample_per_class: ", nb_val_sample_per_class)

print("--")

epochs = 100

batch_size = 128

from functools import partial, update_wrapper

def wrapped_partial(func, *args, **kwargs):

    partial_func = partial(func, *args, **kwargs)

    update_wrapper(partial_func, func)

    return partial_func

def w_categorical_crossentropy(y_true, y_pred, weights):

    nb_cl = len(weights)

    final_mask = K.zeros_like(y_pred[:, 0])

    y_pred_max = K.max(y_pred, axis=1)

    y_pred_max = K.expand_dims(y_pred_max, 1)

    y_pred_max_mat = K.equal(y_pred, y_pred_max)

    for c_p, c_t in itertools.product(range(nb_cl), range(nb_cl)):

        final_mask += (K.cast(weights[c_t, c_p],K.floatx()) * K.cast(y_pred_max_mat[:, c_p] ,K.floatx())* K.cast(y_true[:, c_t],K.floatx()))

    return K.categorical_crossentropy(y_pred, y_true) * final_mask

w_array = np.ones((4,4))

w_array[0, 1] = 1

w_array[0, 2] = 1

w_array[0, 3] = 1

w_array[1, 0] = float(nb_sample_per_class[0])/float(nb_sample_per_class[1])

w_array[1, 2] = float(nb_sample_per_class[0])/float(nb_sample_per_class[1])

w_array[1, 3] = float(nb_sample_per_class[0])/float(nb_sample_per_class[1])

w_array[2, 0] = float(nb_sample_per_class[0])/float(nb_sample_per_class[2])

w_array[2, 1] = float(nb_sample_per_class[0])/float(nb_sample_per_class[2])

w_array[2, 3] = float(nb_sample_per_class[0])/float(nb_sample_per_class[2])

w_array[3, 0] = float(nb_sample_per_class[0])/float(nb_sample_per_class[3])

w_array[3, 1] = float(nb_sample_per_class[0])/float(nb_sample_per_class[3])

w_array[3, 2] = float(nb_sample_per_class[0])/float(nb_sample_per_class[3])

ncce = partial(w_categorical_crossentropy, weights=w_array)

ncce.__name__ = 'w_categorical_crossentropy'

# sgd = optimizers.SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)

#  """ # RSNA version custom small model

model = tf.keras.Sequential()

model.add(layers.Conv2D(32, (3, 3), input_shape=(img_height, img_width, 3), name='conv1'))

model.add(layers.Activation('relu'))

model.add(layers.BatchNormalization())

model.add(layers.MaxPooling2D(pool_size=(2, 2)))

model.add(layers.Dropout(0.2))

model.add(layers.Conv2D(64, (3, 3)))

model.add(layers.Activation('relu'))

model.add(layers.BatchNormalization())

model.add(layers.MaxPooling2D(pool_size=(2, 2)))

model.add(layers.Conv2D(64, (3, 3)))

model.add(layers.Activation('relu'))

model.add(layers.BatchNormalization())

model.add(layers.MaxPooling2D(pool_size=(2, 2)))

model.add(layers.Conv2D(64, (3, 3)))

model.add(layers.Activation('relu'))

model.add(layers.BatchNormalization())

model.add(layers.MaxPooling2D(pool_size=(2, 2)))

model.add(layers.Conv2D(32, (3, 3), padding='same'))

model.add(layers.Activation('relu'))

model.add(layers.BatchNormalization())

model.add(layers.MaxPooling2D(pool_size=(2, 2)))

model.add(layers.Conv2D(32, (3, 3), padding='same'))

model.add(layers.Activation('relu'))

model.add(layers.BatchNormalization())

model.add(layers.MaxPooling2D(pool_size=(2, 2)))

model.add(layers.Dropout(0.2))

model.add(layers.Flatten())

model.add(layers.Dense(512))

model.add(layers.Activation('relu'))

model.add(layers.Dropout(0.2))

model.add(layers.Dense(nb_classes, name='output'))

model.add(layers.Activation('softmax'))

#  """

# model.compile(loss=ncce,

#               optimizer=sgd,

#               metrics=['accuracy'])

# base_model = InceptionV3(include_top=False, weights='imagenet', input_shape=(img_height, img_width, 3))

# x = base_model.output

# x = GlobalAveragePooling2D()(x)

# # let's add a fully-connected layer

# x = Dense(1024, activation='relu')(x)

# # and a logistic layer -- let's say we have 200 classes

# predictions = Dense(4, activation='softmax')(x)

# model = Model(inputs=base_model.input, outputs=predictions)

model.compile(loss=ncce,

              optimizer=tf.keras.optimizers.Adam(),

              metrics=['accuracy'])

# Display the model's architecture

model.summary()

# # random crop patch from:

# # https://mc.ai/extending-keras-imagedatagenerator-to-support-random-cropping/

# 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)

# this is the augmentation configuration we will use for training

train_datagen = tf.keras.preprocessing.image.ImageDataGenerator(

    rescale=1. / 255,

    shear_range=0.1,

    zoom_range=0.1,

    #  horizontal_flip=True,

    rotation_range=5,

    #  width_shift_range=0.01,

    #  height_shift_range=0.01,

    #  brightness_range=[0.2, 1.0],

    )

# this is the augmentation configuration we will use for testing:

# only rescaling

test_datagen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1. / 255)

# flow_from_directory will print:

# Found xxx images belonging to xxx classes

train_generator = train_datagen.flow_from_directory(

    train_data_dir,

    target_size=(img_height, img_width),

    batch_size=batch_size,

    class_mode='categorical')

validation_generator = test_datagen.flow_from_directory(

    validation_data_dir,

    target_size=(img_height, img_width),

    batch_size=batch_size,

    class_mode='categorical')

# iterator that returns image_batch, label_batch pairs

for image_batch, label_batch in train_generator:

    print("Image batch shape: ", image_batch.shape)

    print("Label batch shape: ", label_batch.shape)

    break

for image_batch, label_batch in validation_generator:

    print("Image batch shape: ", image_batch.shape)

    print("Label batch shape: ", label_batch.shape)

    break

# # crop from 100 to CROP_LENGTH = 84

# train_crops = crop_generator(train_batches, CROP_LENGTH)

########################

# callbacks setting

########################

# for modelcheckpoint:

# if True, then only the model's weights will

# be saved (model.save_weights(filepath)),

# else the full model is saved (model.save(filepath)).

# for ModelCheckpoint and model saves

# Saving everything into a single archive in the TensorFlow SavedModel

# format (or in the older Keras H5 format). This is the standard practice.

formatted_time = datetime.datetime.now().strftime("%m%d-%H%M")

save_model_dir = "Axial_center_resnetscale150V1_150x150bat128_6LDropout_Date{}".format(formatted_time)

if not os.path.exists(save_model_dir):

    print(save_model_dir, " will be created")

    os.makedirs(save_model_dir)

# store the model json as:

store_model_json_name = "axial-center-Date{}.json".format(formatted_time)

# store model

model_json = model.to_json()

model_json_path = os.path.join(save_model_dir, store_model_json_name)

with open(model_json_path, "w") as json_file:

    json_file.write(model_json)

checkpoint_filepath = os.path.join(

    save_model_dir,

    "{0}_Ep{{epoch:02d}}_ValAcc{{val_acc:.3f}}_ValLoss{{val_loss:.2f}}.h5"

    .format(save_model_dir)

)

callbacks = [

    # tf.keras.callbacks.EarlyStopping(

    #     monitor='val_loss',

    #     patience=15,

    #     verbose=1),

    tf.keras.callbacks.ModelCheckpoint(

        filepath=checkpoint_filepath,

        monitor='val_acc',

        save_best_only=False,

        save_weights_only=True,

        verbose=1,

        save_freq="epoch"),

    tf.keras.callbacks.ReduceLROnPlateau(

        monitor='val_loss',

        factor=0.5, # usually 0.1

        patience=10,

        verbose=1),

    tf.keras.callbacks.CSVLogger(

        filename=os.path.join(

            save_model_dir,

            '{}.csv'.format(save_model_dir)

        ),

        append=False,

        separator=','),

]

train_steps = np.ceil(nb_train_samples / batch_size)

print("len train_generator: ", len(train_generator))

val_steps = np.ceil(nb_validation_samples / batch_size)

print("len validation_generator: ", len(validation_generator))

train_history = model.fit(

    train_generator,

    steps_per_epoch=train_steps,

    epochs=epochs,

    validation_data=validation_generator,

    validation_steps=val_steps,

    callbacks=callbacks)