from __future__ import print_function
import os
from keras.utils import plot_model
from keras import backend as K
from keras import callbacks
import numpy as np
# import pdb

from .hed import hed
from .hed_data_parser import DataParser


def generate_minibatches(dataParser, train=True):
    # pdb.set_trace()
    while True:
        if train:
            batch_ids = np.random.choice(dataParser.training_ids, dataParser.batch_size_train)
        else:
            batch_ids = np.random.choice(dataParser.validation_ids, dataParser.batch_size_train * 2)
        ims, ems, _ = dataParser.get_batch(batch_ids)
        yield (ims, [ems, ems, ems, ems, ems, ems])


######
if __name__ == "__main__":
    # params
    model_name = 'HEDSeg'
    model_dir = os.path.join('_model', 'checkpoints', model_name)
    csv_fn = os.path.join(model_dir, 'train_log.csv')
    checkpoint_fn = os.path.join(model_dir, 'checkpoint.{epoch:02d}-{val_loss:.2f}.hdf5')

    batch_size_train = 10

    # environment
    K.set_image_data_format('channels_last')
    K.image_data_format()
    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
    if not os.path.isdir(model_dir):
        os.makedirs(model_dir)

    # prepare data
    dataParser = DataParser(batch_size_train)

    # model
    model = hed()
    # plot_model(model, to_file=os.path.join(model_dir, 'model.pdf'), show_shapes=True)

    # training
    # call backs
    checkpointer = callbacks.ModelCheckpoint(filepath=checkpoint_fn, verbose=1, save_best_only=True)
    csv_logger = callbacks.CSVLogger(csv_fn, append=True, separator=';')
    tensorboard = callbacks.TensorBoard(log_dir=model_dir, histogram_freq=0, batch_size=10,
                                        write_graph=False, write_grads=True, write_images=False)

    train_history = model.fit_generator(
        generate_minibatches(dataParser, ),
        # max_q_size=40, workers=1,
        steps_per_epoch=dataParser.steps_per_epoch,  # batch size
        epochs=2048 * 2,
        validation_data=generate_minibatches(dataParser, train=False),
        validation_steps=dataParser.validation_steps,
        callbacks=[checkpointer, csv_logger, tensorboard])

    # pdb.set_trace()

    print(train_history)
