from __future__ import division, print_function

import matplotlib.pyplot as plt

import numpy as np

from tqdm import tqdm

from keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard, ReduceLROnPlateau, LearningRateScheduler

from keras import models

from graph import ECG_model

from config import get_config

from utils import *

def train(config, X, y, Xval=None, yval=None):

    classes = ['N','V','/','A','F','~']#,'L','R',f','j','E','a']#,'J','Q','e','S']

    print("Initial shapes - X:", X.shape, "y:", y.shape)

    print("Initial validation shapes - Xval:", Xval.shape if Xval is not None else None, "yval:", yval.shape if yval is not None else None)

    print("Any NaN in initial X:", np.any(np.isnan(X)), "y:", np.any(np.isnan(y)))

    Xe = np.expand_dims(X, axis=2)

    if not config.split:

        from sklearn.model_selection import train_test_split

        Xe, Xvale, y, yval = train_test_split(Xe, y, test_size=0.2, random_state=1)

    else:

        Xvale = np.expand_dims(Xval, axis=2)

        print("Data shapes before training - Xe:", Xe.shape, "y:", y.shape)

        print("Val shapes before training - Xvale:", Xvale.shape, "yval:", yval.shape)

        print("Final shapes - Xe:", Xe.shape, "y:", y.shape)

        print("Final val shapes - Xvale:", Xvale.shape, "yval:", yval.shape)

    if config.checkpoint_path is not None:

        model = models.load_model(config.checkpoint_path)

        initial_epoch = config.resume_epoch # put the resuming epoch

    else:

        model = ECG_model(config)

        initial_epoch = 0

    mkdir_recursive('models')

    #lr_decay_callback = LearningRateSchedulerPerBatch(lambda epoch: 0.1)

    # Validate input data

    if np.any(np.isnan(Xe)) or np.any(np.isnan(y)):

        raise ValueError("Input data contains None/NaN values")

    if np.any(np.isnan(Xvale)) or np.any(np.isnan(yval)):

        raise ValueError("Validation data contains None/NaN values")

    callbacks = [

            EarlyStopping(patience = config.patience, verbose=1),

            ReduceLROnPlateau(factor = 0.5, patience = 3, min_lr = 0.01, verbose=1),

            TensorBoard(log_dir='./logs', histogram_freq=0, write_graph=True, write_images=True),

            ModelCheckpoint('models/{}-latest.keras'.format(config.feature), monitor='val_loss', save_best_only=False, verbose=1, save_freq=10)

            # , lr_decay_callback

    ]

    model.fit(Xe, y,

            validation_data=(Xvale, yval),

            epochs=config.epochs,

            batch_size=config.batch,

            callbacks=callbacks,

            initial_epoch=initial_epoch)

    print_results(config, model, Xvale, yval, classes, )

    #return model

def main(config):

    print('feature:', config.feature)

    #np.random.seed(0)

    (X,y, Xval, yval) = loaddata(config.input_size, config.feature)

    print(X, y)

    train(config, X, y, Xval, yval)

if __name__=="__main__":

    config = get_config()

    main(config)