from datetime import datetime
from model import MyDeepModel, create_submission
from data_loader import read_testset, read_trainset, DataGenerator

import keras as K

from sklearn.model_selection import ShuffleSplit


# from K_applications.resnet import ResNet50
from keras.applications.inception_v3 import InceptionV3
from keras.applications.inception_resnet_v2 import InceptionResNetV2, preprocess_input
from keras.applications.densenet import DenseNet121
from keras.applications.mobilenet_v2 import MobileNetV2

test_images_dir = '../../data/stage_1_test_images/'
train_images_dir = '../../data/stage_1_train_images/'
trainset_filename = "../../data/stage_1_train.csv"
testset_filename = "../../stage_1_sample_submission.csv"
num_epochs = 10
img_shape = (256,256,3)
batch_size=32
TRAINING =True # If False, then just load model and predict

engine=InceptionV3
model_filename="InceptionV3_{}.hdf5".format(datetime.now().strftime('%Y_%m_%d_%H_%M_%S'))
#model_filename="wrapper_2019_11_02_22_06_45.hdf5"

# obtain model
model = MyDeepModel(engine=engine, input_dims=img_shape, batch_size=batch_size,
                    learning_rate=5e-4,
                    num_epochs=num_epochs, decay_rate=0.8,
                    decay_steps=1,
                    weights="imagenet", verbose=1,
                    train_image_dir=train_images_dir,
                    model_filename=model_filename)

model.load("epoch2.hdf5")

#model.load(model_filename)  # Use previous checkpoint


if (TRAINING == True):

    df = read_trainset(trainset_filename)
    ss = ShuffleSplit(n_splits=10, test_size=0.1, random_state=816).split(df.index)
    # lets go for the first fold only
    train_idx, valid_idx = next(ss)

    # Train the model
    model.fit_model(df.iloc[train_idx], df.iloc[valid_idx])


test_df = read_testset(testset_filename)
test_generator = DataGenerator(test_df.index, None, 1, img_shape, test_images_dir)
best_model = K.models.load_model(model.model_filename, compile=False)

prediction_df = create_submission(best_model, test_generator, test_df)