"""
Training script.

Created on 2019-01-28-12-21
Author: Stephan Rasp, raspstephan@gmail.com
"""

from cbrain.imports import *
from cbrain.utils import *
from cbrain.losses import *
from cbrain.data_generator_vae import DataGenerator
#from cbrain.models_new import *
from cbrain.learning_rate_schedule import LRUpdate
from cbrain.save_weights import save2txt, save_norm
from tensorflow.keras.callbacks import LearningRateScheduler
from tensorflow.keras.losses import mse
import json
from datetime import datetime,date
from cbrain.models_vae import VariationalAutoEncoder

def main(args):
    """Main training script."""

    if args.gpu is not None:
        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    limit_mem()

    logging.basicConfig(
    filename='saved_models\logs\model_{}_{}.txt'.format(args.exp_name,date.today()),
    filemode ='a',
    format='%(asctime)s %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p',
    level=logging.DEBUG
    )
            


    # Load output scaling dictionary
    out_scale_dict = load_pickle(args.output_dict)

    logging.info('Create training and validation data generators')
    train_gen = DataGenerator(
        data_fn=args.data_dir + args.train_fn,
        input_vars=args.inputs,
        output_vars=args.outputs,
        norm_fn=args.data_dir + args.norm_fn,
        input_transform=(args.input_sub, args.input_div),
        output_transform=out_scale_dict,
        batch_size=args.batch_size,
        shuffle=True,
        var_cut_off=args.var_cut_off
    )

    if args.valid_fn is not None:
        valid_gen = DataGenerator(
            data_fn=args.data_dir + args.valid_fn,
            input_vars=args.inputs,
            output_vars=args.outputs,
            norm_fn=args.data_dir + args.norm_fn,
            input_transform=(args.input_sub, args.input_div),
            output_transform=out_scale_dict,
            batch_size=args.batch_size * 10,
            shuffle=False,
            var_cut_off=args.var_cut_off
        )
    else:
        valid_gen = None

    logging.info('Build model')


## change from here
    '''
    vae = vae_model(
        input_shape= train_gen.n_inputs,
        latent_dim = args.latent_dim,
        encoder_layers = args.encoder_layers,
        decoder_layers = args.decoder_layers,
        activation = args.activation)
    '''
    lrs = LearningRateScheduler(LRUpdate(args.lr, args.lr_step, args.lr_divide))
    
    vae = VariationalAutoEncoder(original_dim = train_gen.n_inputs, intermediate_dim = args.intermediate_dim ,latent_dim = args.latent_dim, activation=args.activation)
    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
    mse_loss_fn = tf.keras.losses.MeanSquaredError()

    loss_metric = tf.keras.metrics.Mean()
    #vae.compile(optimizer, loss=mse_loss_fn, metr)

    
    #optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

    #epochs = 3

    #vae = VariationalAutoEncoder(784, 64, 32)

    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

    vae.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())
    #vae.fit(x_train, x_train, epochs=3, batch_size=64) 

    vae.fit_generator(
        train_gen, epochs=args.epochs, validation_data=valid_gen, callbacks=[lrs]
        )

    '''
    def mse_loss(y_true, y_pred):
        reconstruction_loss = mse(y_true, y_pred)
        reconstruction_loss *= 64 # vae.get_layer('encoder_input').input_shape
        z_mean = vae.get_layer('encoder').get_layer('z_mean').output
        z_log_var = vae.get_layer('encoder').get_layer('z_log_var').output
        kl_loss = 1 + z_log_var - K.square(z_mean) - K.exp(z_log_var)
        kl_loss = K.sum(kl_loss, axis=-1)
        kl_loss *= -0.5
        return K.mean(reconstruction_loss + kl_loss)


    print(vae.summary())

    logging.info('Compile model')
    #loss = args.loss
    metrics = [mse]

    vae.compile(args.optimizer, loss=mse_loss, metrics=metrics)

    lrs = LearningRateScheduler(LRUpdate(args.lr, args.lr_step, args.lr_divide))

    logging.info('Train model')
    vae.fit_generator(
        train_gen, epochs=args.epochs, validation_data=valid_gen, callbacks=[lrs]
        )
    '''


    #print(np.array([x[1] for x in valid_gen]).flatten())
    #print(model.predict_generator(valid_gen).flatten())
    # save valid(dev)/ prediction    
    #print(len([x[0][0] for x in valid_gen][0]))
    #pd.DataFrame({'Actual':np.array([x[1] for x in valid_gen]).flatten()/out_scale_dict['PRECT'],
    #            'Predicted':model.predict_generator(valid_gen).flatten()/out_scale_dict['PRECT']}).to_csv('saved_models\logs\prediction.csv')

    if args.exp_name is not None:
        exp_dir = args.model_dir + args.exp_name + '/'
        os.makedirs(exp_dir, exist_ok=True)
        model_fn = exp_dir + 'model_vae.h5'
        logging.info(f'Saving model as {model_fn}')
        #vae.save(model_fn)
        #tf.saved_model.save(vae, exp_dir + "/" + "test/")
        vae.save_weights(exp_dir + "/" + "test/", save_format='tf')
            
        '''
        if args.save_txt:
            weights_fn = exp_dir + 'vae_weights.h5'
            logging.info(f'Saving weights as {weights_fn}')
            vae.save_weights(weights_fn)
            save2txt(weights_fn, exp_dir)
        
        #save_norm(train_gen.input_transform, train_gen.output_transform, exp_dir)
        '''
    logging.info('Done!')


# Create command line interface
if __name__ == '__main__':
    p = ArgParser()
    p.add('-c', '--config_file', default='config.yml', is_config_file=True, help='Path to config file.')

    # Data arguments
    p.add('--data_dir', type=str, help='Path to preprocessed data files.')
    p.add('--inputs', type=str, nargs='+', help='List of input variables.')
    p.add('--outputs', type=str, nargs='+', help='List of output variables.')
    p.add('--train_fn', type=str, help='File name of training file.')
    p.add('--norm_fn', type=str, help='File name of normalization file.')
    p.add('--input_sub', type=str, help='What to subtract from input array. E.g. "mean"')
    p.add('--input_div', type=str, help='What to divide input array by. E.g. "maxrs"')
    p.add('--output_dict', type=str, help='Output scaling dictionary.')
    p.add('--var_cut_off', type=json.loads, help='Input variable cut off for upper levels.')
    p.add('--latent_dim', type = int, default = 32, help='Size of the latent vector')
    p.add('--intermediate_dim', type = int, nargs='+', help='Intermediate layer sizes.')
    #p.add('--decoder_layers', type = int, nargs='+', help='Decoder layer sizes.')

    p.add('--valid_fn', type=str, default=None, help='File name of training file.')

    # Neural network hyperparameteris
    p.add('--batch_size', type=int, default=1024, help='Batch size of training generator.')
    #p.add('--hidden_layers', type=int, nargs='+', help='Hidden layer sizes.')
    p.add('--activation', type=str, default='LeakyReLU', help='Activation function.')
    p.add('--optimizer', type=str, default='adam', help='Optimizer.')
    #p.add('--conservation_layer', dest='conservation_layer', action='store_true', help='Add conservation layer.')
    #p.set_defaults(conservation_layer=False)

    # Loss parameters
    p.add('--loss', type=str, default='mse', help='Loss function.')
    #p.add('--conservation_metrics', dest='conservation_metrics', action='store_true', help='Add conservation metrics.')
    #p.set_defaults(conservation_metrics=False)
    #p.add('--alpha_mass', type=float, default=0.25, help='If weak_loss, weight of mass loss.')
    #p.add('--alpha_ent', type=float, default=0.25, help='If weak_loss, weight of ent loss.')
    #p.add('--noadiab', dest='noadiab', action='store_true',help='noadiab')
    #p.set_defaults(noadiab=False)

    # Learning rate schedule
    p.add('--lr', type=float, default=0.001, help='Initial learning rate.')
    p.add('--lr_step', type=int, default=2, help='Divide every step epochs.')
    p.add('--lr_divide', type=float, default=5, help='Divide by this number.')
    p.add('--epochs', type=int, default=10, help='Number of epochs.')

    # Save parameters
    p.add('--exp_name', type=str, default="test_dnn_vae", help='Experiment identifier.')
    p.add('--model_dir', type=str, default='./saved_models/', help='Model save path.')
    p.add('--save_txt', dest='save_txt', action='store_true', help='Save F90 txt files.')
    p.set_defaults(save_txt=True)

    p.add('--gpu', type=str, default=None, help='Which GPU to use.')

    args = p.parse_args()
    main(args)



