#!/usr/bin/env python3
import os
import numpy as np
import sys
import logging
import string
import random

PhysNetdir = os.environ['PhysNet']
sys.path.append(PhysNetdir)

import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
from shutil import copyfile
from datetime import datetime
from neural_network.NeuralNetwork import *
from neural_network.activation_fn import *
from training.Trainer        import *
from training.DataContainer import *
from training.DataProvider  import *
from training.DataQueue     import *

def train(args):
    tf.compat.v1.reset_default_graph()
    #used for creating a "unique" id for a run (almost impossible to generate the same twice)
    def id_generator(size=8, chars=string.ascii_uppercase + string.ascii_lowercase + string.digits):
        return ''.join(random.SystemRandom().choice(chars) for _ in range(size))

    logging.basicConfig(filename='train.log',level=logging.DEBUG)

    #create directories
    #a unique directory name is created for this run based on the input
    if args.physnet.restart is None:
        os.system('rm -rf '+args.mlmodelout+' > /dev/null 2>&1')
        directory = args.mlmodelout
    else:
        directory=args.mlmodelin
    logging.info("creating directories...")
    if not os.path.exists(directory):
        os.makedirs(directory)
    best_dir = os.path.join(directory, 'best')
    if not os.path.exists(best_dir):
        os.makedirs(best_dir)
    log_dir = os.path.join(directory, 'logs')
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    best_loss_file  = os.path.join(best_dir, 'best_loss.npz')
    best_checkpoint = os.path.join(best_dir, 'best_model.ckpt')
    step_checkpoint = os.path.join(log_dir,  'model.ckpt')

    #load dataset
    # logging.info("loading dataset...")
    # data = DataContainer(args.physnet.dataset)
    subtrain = DataContainer('PhysNet_subtrain.npz')
    validate = DataContainer('PhysNet_validate.npz')

    
    args.physnet.num_train = len(subtrain)
    args.physnet.num_valid = len(validate)
    #write config file (to restore command line arguments)
    logging.info("writing args to file...")
    with open(os.path.join(directory, 'config.txt'), 'w') as f:
        for k, v in args.data.physnet.data.items():
            if k[0] != '_' and k != 'batch_sizes':
                f.write('--'+ k + '='+ str(v) + "\n")

    #generate DataProvider (splits dataset into training, validation and test set based on seed)
    # data_provider = DataProvider(data, args.physnet.num_train, args.physnet.num_valid, args.physnet.batch_size, args.physnet.valid_batch_size, seed=args.physnet.seed)
    subtrain_provider = DataProvider(subtrain, args.physnet.num_train, 0, args.physnet.batch_size, 0, seed=args.physnet.seed)
    validate_provider = DataProvider(validate, 0, args.physnet.num_valid, 0, args.physnet.valid_batch_size, seed=args.physnet.seed)

    #create neural network
    logging.info("creating neural network...")
    nn = NeuralNetwork(F=args.physnet.num_features,           
                    K=args.physnet.num_basis,                
                    sr_cut=args.physnet.cutoff,              
                    num_blocks=args.physnet.num_blocks, 
                    num_residual_atomic=args.physnet.num_residual_atomic,
                    num_residual_interaction=args.physnet.num_residual_interaction,
                    num_residual_output=args.physnet.num_residual_output,
                    use_electrostatic=(args.physnet.use_electrostatic==1),
                    use_dispersion=(args.physnet.use_dispersion==1),
                    s6=args.physnet.grimme_s6,
                    s8=args.physnet.grimme_s8,
                    a1=args.physnet.grimme_a1,
                    a2=args.physnet.grimme_a2,
                    Eshift=subtrain_provider.EperA_mean,  
                    Escale=subtrain_provider.EperA_stdev,   
                    activation_fn=shifted_softplus, 
                    seed=None,
                    scope="neural_network")
    NASdict = {'Eshift': subtrain_provider.EperA_mean,'Escale': subtrain_provider.EperA_stdev}
    np.savez(os.path.join(directory, 'NAS.npz'), **NASdict)

    logging.info("prepare training...")
    #generate data queues for efficient training
    train_queue = DataQueue(subtrain_provider.next_batch, capacity=1000, scope="train_data_queue")
    valid_queue = DataQueue(validate_provider.next_valid_batch, capacity=args.physnet.num_valid//args.physnet.valid_batch_size, scope="valid_data_queue")

    #get values for training and validation set
    Eref_t, Earef_t, Fref_t, Z_t, Dref_t, Qref_t, Qaref_t, R_t, idx_i_t, idx_j_t, batch_seg_t = train_queue.dequeue_op
    Eref_v, Earef_v, Fref_v, Z_v, Dref_v, Qref_v, Qaref_v, R_v, idx_i_v, idx_j_v, batch_seg_v = valid_queue.dequeue_op

    #calculate all necessary quantities (unscaled partial charges, energies, forces)
    Ea_t, Qa_t, Dij_t, nhloss_t = nn.atomic_properties(Z_t, R_t, idx_i_t, idx_j_t)
    Ea_v, Qa_v, Dij_v, nhloss_v = nn.atomic_properties(Z_v, R_v, idx_i_v, idx_j_v)
    energy_t, forces_t = nn.energy_and_forces_from_atomic_properties(Ea_t, Qa_t, Dij_t, Z_t, R_t, idx_i_t, idx_j_t, Qref_t, batch_seg_t)
    energy_v, forces_v = nn.energy_and_forces_from_atomic_properties(Ea_v, Qa_v, Dij_v, Z_v, R_v, idx_i_v, idx_j_v, Qref_v, batch_seg_v)
    #total charge
    Qtot_t = tf.segment_sum(Qa_t, batch_seg_t)
    Qtot_v = tf.segment_sum(Qa_v, batch_seg_v)
    #dipole moment vector
    QR_t = tf.stack([Qa_t*R_t[:,0], Qa_t*R_t[:,1], Qa_t*R_t[:,2]],1)
    QR_v = tf.stack([Qa_v*R_v[:,0], Qa_v*R_v[:,1], Qa_v*R_v[:,2]],1)
    D_t = tf.segment_sum(QR_t, batch_seg_t)
    D_v = tf.segment_sum(QR_v, batch_seg_v)

    #function to calculate loss, mean squared error, mean absolute error between two values
    def calculate_errors(val1, val2, weights=1):
        with tf.name_scope("calculate_errors"):
            delta  = tf.abs(val1-val2)
            delta2 = delta**2
            mse    = tf.reduce_mean(delta2)
            mae    = tf.reduce_mean(delta)
            loss   = mae #mean absolute error loss
        return loss, mse, mae

    with tf.name_scope("loss"):
        #calculate energy, force, charge and dipole errors/loss
        #energy
        if subtrain.E is not None:
            eloss_t, emse_t, emae_t = calculate_errors(Eref_t, energy_t)
            eloss_v, emse_v, emae_v = calculate_errors(Eref_v, energy_v)
        else:
            eloss_t, emse_t, emae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)
            eloss_v, emse_v, emae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)
        #atomic energies
        if subtrain.Ea is not None:
            ealoss_t, eamse_t, eamae_t = calculate_errors(Earef_t, Ea_t)
            ealoss_v, eamse_v, eamae_v = calculate_errors(Earef_v, Ea_v)
        else:
            ealoss_t, eamse_t, eamae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)
            ealoss_v, eamse_v, eamae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)
        #forces
        if subtrain.F is not None:
            floss_t, fmse_t, fmae_t = calculate_errors(Fref_t, forces_t)
            floss_v, fmse_v, fmae_v = calculate_errors(Fref_v, forces_v)
        else:
            floss_t, fmse_t, fmae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)
            floss_v, fmse_v, fmae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)
        #charge
        if subtrain.Q is not None:     
            qloss_t, qmse_t, qmae_t = calculate_errors(Qref_t, Qtot_t)
            qloss_v, qmse_v, qmae_v = calculate_errors(Qref_v, Qtot_v)
        else:
            qloss_t, qmse_t, qmae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)
            qloss_v, qmse_v, qmae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)
        #atomic charges
        if subtrain.Qa is not None:
            qaloss_t, qamse_t, qamae_t = calculate_errors(Qaref_t, Qa_t)
            qaloss_v, qamse_v, qamae_v = calculate_errors(Qaref_v, Qa_v)
        else:
            qaloss_t, qamse_t, qamae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)
            qaloss_v, qamse_v, qamae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)
        #dipole
        if subtrain.D is not None:
            dloss_t, dmse_t, dmae_t = calculate_errors(Dref_t, D_t)
            dloss_v, dmse_v, dmae_v = calculate_errors(Dref_v, D_v)
        else:
            dloss_t, dmse_t, dmae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)
            dloss_v, dmse_v, dmae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

        #define additional variables (such that certain losses can be overwritten)
        eloss_train = eloss_t
        floss_train = floss_t
        qloss_train = qloss_t
        dloss_train = dloss_t
        eloss_valid = eloss_v
        floss_valid = floss_v
        qloss_valid = qloss_v
        dloss_valid = dloss_v

        #atomic energies are present, so they replace the normal energy loss
        if subtrain.Ea is not None:
            eloss_train = ealoss_t
            eloss_valid = ealoss_v

        #atomic charges are present, so they replace the normal charge loss and nullify dipole loss
        if subtrain.Qa is not None:
            qloss_train = qaloss_t
            qloss_valid = qaloss_v
            dloss_train = tf.constant(0.0)
            dloss_valid = tf.constant(0.0)

        #define loss function (used to train the model)
        l2loss = tf.reduce_mean(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) 
        loss_t = eloss_train + args.physnet.force_weight*floss_train + args.physnet.charge_weight*qloss_train + args.physnet.dipole_weight*dloss_train + args.physnet.nhlambda*nhloss_t + args.physnet.l2lambda*l2loss
        loss_v = eloss_valid + args.physnet.force_weight*floss_valid + args.physnet.charge_weight*qloss_valid + args.physnet.dipole_weight*dloss_valid + args.physnet.nhlambda*nhloss_v + args.physnet.l2lambda*l2loss

    #create trainer
    trainer  = Trainer(args.physnet.learning_rate, args.physnet.decay_steps, args.physnet.decay_rate, scope="trainer")
    with tf.name_scope("trainer_ops"):
        train_op = trainer.build_train_op(loss_t, args.physnet.ema_decay, args.physnet.max_norm)
        save_variable_backups_op = trainer.save_variable_backups()
        load_averaged_variables_op = trainer.load_averaged_variables()
        restore_variable_backups_op = trainer.restore_variable_backups()

    #creates a summary from key-value pairs given a dictionary
    def create_summary(dictionary):
        summary = tf.Summary()
        for key, value in dictionary.items():
            summary.value.add(tag=key, simple_value=value)
        return summary

    #create summary writer
    nn_summary_op = tf.summary.merge_all()
    summary_writer = tf.summary.FileWriter(logdir=log_dir, graph=tf.get_default_graph())

    #create saver
    with tf.name_scope("saver"):
        saver = tf.train.Saver(max_to_keep=50)

    #save/load best recorded loss (only the best model is saved)
    if os.path.isfile(best_loss_file):
        loss_file   = np.load(best_loss_file)
        best_loss   = loss_file["loss"].item()
        best_emae   = loss_file["emae"].item()
        best_ermse  = loss_file["ermse"].item()
        best_fmae   = loss_file["fmae"].item()
        best_frmse  = loss_file["frmse"].item()
        best_qmae   = loss_file["qmae"].item()
        best_qrmse  = loss_file["qrmse"].item()
        best_dmae   = loss_file["dmae"].item()
        best_drmse  = loss_file["drmse"].item()
        best_step   = loss_file["step"].item()
    else:
        best_loss  = np.Inf #initialize best loss to infinity
        best_emae  = np.Inf
        best_ermse = np.Inf
        best_fmae  = np.Inf
        best_frmse = np.Inf
        best_qmae  = np.Inf
        best_qrmse = np.Inf
        best_dmae  = np.Inf
        best_drmse = np.Inf
        best_step  = 0.
        np.savez(best_loss_file, loss=best_loss, emae=best_emae,   ermse=best_ermse, 
                                                fmae=best_fmae,   frmse=best_frmse, 
                                                qmae=best_qmae,   qrmse=best_qrmse, 
                                                dmae=best_dmae,   drmse=best_drmse, 
                                                step=best_step)

    #for calculating average performance on the training set
    def reset_averages():
        return 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0

    def update_averages(num, tmploss_avg, tmploss, emse_avg, emse, emae_avg, emae, fmse_avg, fmse, fmae_avg, fmae, 
                        qmse_avg, qmse, qmae_avg, qmae, dmse_avg, dmse, dmae_avg, dmae):
        num += 1
        tmploss_avg += (tmploss-tmploss_avg)/num
        emse_avg += (emse-emse_avg)/num
        emae_avg += (emae-emae_avg)/num
        fmse_avg += (fmse-fmse_avg)/num
        fmae_avg += (fmae-fmae_avg)/num
        qmse_avg += (qmse-qmse_avg)/num
        qmae_avg += (qmae-qmae_avg)/num
        dmse_avg += (dmse-dmse_avg)/num
        dmae_avg += (dmae-dmae_avg)/num
        return num, tmploss_avg, emse_avg, emae_avg, fmse_avg, fmae_avg, qmse_avg, qmae_avg, dmse_avg, dmae_avg

    #initialize training set error averages
    num_t, tmploss_avg_t, emse_avg_t, emae_avg_t, fmse_avg_t, fmae_avg_t, qmse_avg_t, qmae_avg_t, dmse_avg_t, dmae_avg_t = reset_averages()

    #create tensorflow session
    with tf.Session() as sess:
        if (args.physnet.record_run_metadata > 0):
            run_options  = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
            run_metadata = tf.RunMetadata()
        else:
            run_options  = None
            run_metadata = None

        #start data queues
        coord = tf.train.Coordinator()
        train_queue.create_thread(sess, coord)
        valid_queue.create_thread(sess, coord)

        #initialize variables
        tf.global_variables_initializer().run()

        #restore latest checkpoint
        checkpoint = tf.train.latest_checkpoint(log_dir)
        if checkpoint is not None:
            step = int(checkpoint.split('-')[-1]) #reads step from checkpoint filename
            saver.restore(sess, checkpoint)
            sess.run(trainer.global_step.assign(step))
        else:
            step = 0

        #training loop
        logging.info("starting training...")
        while not coord.should_stop():
            #finish training when maximum number of iterations is reached
            if step > args.physnet.max_steps:
                coord.request_stop()
                break

            #perform training step 
            step += 1
            _, tmploss, emse, emae, fmse, fmae, qmse, qmae, dmse, dmae = sess.run([train_op, loss_t, emse_t, emae_t, fmse_t, fmae_t, qmse_t, qmae_t, dmse_t, dmae_t], options=run_options, feed_dict={nn.keep_prob: args.physnet.keep_prob}, run_metadata=run_metadata)
            
            #update averages
            num_t, tmploss_avg_t, emse_avg_t, emae_avg_t, fmse_avg_t, fmae_avg_t, qmse_avg_t, qmae_avg_t, dmse_avg_t, dmae_avg_t = update_averages(num_t, tmploss_avg_t, tmploss, emse_avg_t, emse, emae_avg_t, emae, fmse_avg_t, fmse, fmae_avg_t, fmae, qmse_avg_t, qmse, qmae_avg_t, qmae, dmse_avg_t, dmse, dmae_avg_t, dmae)

            #save progress
            if (step % args.physnet.save_interval == 0):
                saver.save(sess, step_checkpoint, global_step=step)   

            #check performance on the validation set
            if (step % args.physnet.validation_interval == 0):
                #save backup variables and load averaged variables
                sess.run(save_variable_backups_op)
                sess.run(load_averaged_variables_op)

                #initialize averages to 0
                num_v, tmploss_avg_v, emse_avg_v, emae_avg_v, fmse_avg_v, fmae_avg_v, qmse_avg_v, qmae_avg_v, dmse_avg_v, dmae_avg_v = reset_averages()
                #compute averages
                for i in range(args.physnet.num_valid//args.physnet.valid_batch_size):
                    tmploss, emse, emae, fmse, fmae, qmse, qmae, dmse, dmae = sess.run([loss_v, emse_v, emae_v, fmse_v, fmae_v, qmse_v, qmae_v, dmse_v, dmae_v])
                    num_v, tmploss_avg_v, emse_avg_v, emae_avg_v, fmse_avg_v, fmae_avg_v, qmse_avg_v, qmae_avg_v, dmse_avg_v, dmae_avg_v = update_averages(num_v, tmploss_avg_v, tmploss, emse_avg_v, emse, emae_avg_v, emae, fmse_avg_v, fmse, fmae_avg_v, fmae, qmse_avg_v, qmse, qmae_avg_v, qmae, dmse_avg_v, dmse, dmae_avg_v, dmae)

                #store results in dictionary
                results = {}
                results["loss_valid"] = tmploss_avg_v
                if subtrain.E is not None:
                    results["energy_mae_valid"]  = emae_avg_v
                    results["energy_rmse_valid"] = np.sqrt(emse_avg_v)
                if subtrain.F is not None:
                    results["forces_mae_valid"]  = fmae_avg_v
                    results["forces_rmse_valid"] = np.sqrt(fmse_avg_v)
                if subtrain.Q is not None:
                    results["charge_mae_valid"]  = qmae_avg_v
                    results["charge_rmse_valid"] = np.sqrt(qmse_avg_v)
                if subtrain.D is not None:
                    results["dipole_mae_valid"]  = dmae_avg_v
                    results["dipole_rmse_valid"] = np.sqrt(dmse_avg_v)

                if results["loss_valid"] < best_loss:
                    best_loss   = results["loss_valid"]
                    if subtrain.E is not None:
                        best_emae   = results["energy_mae_valid"]
                        best_ermse  = results["energy_rmse_valid"]
                    else:
                        best_emae  = np.Inf
                        best_ermse = np.Inf
                    if subtrain.F is not None:
                        best_fmae   = results["forces_mae_valid"]
                        best_frmse  = results["forces_rmse_valid"]
                    else:
                        best_fmae  = np.Inf
                        best_frmse = np.Inf
                    if subtrain.Q is not None:
                        best_qmae   = results["charge_mae_valid"]
                        best_qrmse  = results["charge_rmse_valid"]
                    else:
                        best_qmae  = np.Inf
                        best_qrmse = np.Inf
                    if subtrain.D is not None:
                        best_dmae   = results["dipole_mae_valid"]
                        best_drmse  = results["dipole_rmse_valid"]
                    else:
                        best_dmae  = np.Inf
                        best_drmse = np.Inf
                    best_step = step
                    np.savez(best_loss_file, loss=best_loss, emae=best_emae,   ermse=best_ermse, 
                                            fmae=best_fmae,   frmse=best_frmse, 
                                            qmae=best_qmae,   qrmse=best_qrmse, 
                                            dmae=best_dmae,   drmse=best_drmse, 
                                            step=best_step)
                    nn.save(sess, best_checkpoint, global_step=step)
                results["loss_best"] = best_loss
                if subtrain.E is not None:
                    results["energy_mae_best"]  = best_emae
                    results["energy_rmse_best"] = best_ermse
                if subtrain.F is not None:
                    results["forces_mae_best"]  = best_fmae
                    results["forces_rmse_best"] = best_frmse
                if subtrain.Q is not None:
                    results["charge_mae_best"]  = best_qmae
                    results["charge_rmse_best"] = best_qrmse
                if subtrain.D is not None:
                    results["dipole_mae_best"]  = best_dmae
                    results["dipole_rmse_best"] = best_drmse
                summary = create_summary(results)
                summary_writer.add_summary(summary, global_step=step)

                #restore backup variables
                sess.run(restore_variable_backups_op)
    
            #generate summaries
            if (step % args.physnet.summary_interval == 0) and (step > 0): 
                results = {}            
                results["loss_train"] = tmploss_avg_t
                if subtrain.E is not None:
                    results["energy_mae_train"]  = emae_avg_t
                    results["energy_rmse_train"] = np.sqrt(emse_avg_t)
                if subtrain.F is not None:
                    results["forces_mae_train"]  = fmae_avg_t
                    results["forces_rmse_train"] = np.sqrt(fmse_avg_t)
                if subtrain.Q is not None:
                    results["charge_mae_train"]  = qmae_avg_t
                    results["charge_rmse_train"] = np.sqrt(qmse_avg_t)
                if subtrain.D is not None:
                    results["dipole_mae_train"]  = dmae_avg_t
                    results["dipole_rmse_train"] = np.sqrt(dmse_avg_t)
                num_t, tmploss_avg_t, emse_avg_t, emae_avg_t, fmse_avg_t, fmae_avg_t, qmse_avg_t, qmae_avg_t, dmse_avg_t, dmae_avg_t = reset_averages()

                summary = create_summary(results)
                summary_writer.add_summary(summary, global_step=step)
                nn_summary = nn_summary_op.eval()
                summary_writer.add_summary(nn_summary, global_step=step)
                if (args.physnet.record_run_metadata > 0):
                    summary_writer.add_run_metadata(run_metadata, 'step %d' % step, global_step=step)
                if subtrain.E is not None:
                    print(str(step)+'/'+str(args.physnet.max_steps), "loss:", results["loss_train"], "best:", best_loss, "emae:", results["energy_mae_train"], "best:", best_emae)
