import numpy as np
import os
import sys
import glob
import uproot as ur
import matplotlib.pyplot as plt
import time
import seaborn as sns
import tensorflow as tf
from graph_nets import utils_np
from graph_nets import utils_tf
from graph_nets.graphs import GraphsTuple
import sonnet as snt
import argparse
import yaml
import logging
import tensorflow as tf

from modules.mpdatagen import MPGraphDataGenerator
import modules.multiOutBlock_wWeightedRegress as models
sns.set_context('poster')

if __name__=="__main__":
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', default='configs/default.yaml')
    args = parser.parse_args()

    config = yaml.load(open(args.config))

    data_config = config['data']
    model_config = config['model']
    train_config = config['training']

    data_dir = data_config['data_dir']
    num_train_files = data_config['num_train_files']
    num_val_files = data_config['num_val_files']
    batch_size = data_config['batch_size']
    shuffle = data_config['shuffle']
    num_procs = data_config['num_procs']
    preprocess = data_config['preprocess']
    output_dir = data_config['output_dir']
    already_preprocessed = data_config['already_preprocessed']

    concat_input = model_config['concat_input']

    epochs = train_config['epochs']
    learning_rate = train_config['learning_rate']
    alpha = train_config['alpha']
    os.environ['CUDA_VISIBLE_DEVICES'] = str(train_config['gpu'])
    log_freq = train_config['log_freq']
    save_dir = train_config['save_dir'] + '/Block_'+time.strftime("%Y%m%d_%H%M")+'_'+args.config.replace('.yaml','').split('/')[-1]
    os.makedirs(save_dir, exist_ok=True)
    yaml.dump(config, open(save_dir + '/config.yaml', 'w'))

    logging.basicConfig(level=logging.INFO, 
                        format='%(message)s', 
                        filename=save_dir + '/output.log')
    logging.info('Using config file {}'.format(args.config))
    # logging.info('Running training for {} with concant_input: {}\n'.format(particle_type, concat_input))

    pi0_files = np.sort(glob.glob(data_dir+'*graphs.v01*/*pi0*/*root'))
    pion_files = np.sort(glob.glob(data_dir+'*graphs.v01*/*pion*/*root'))
    train_start = 10
    train_end = train_start + num_train_files
    val_end = train_end + num_val_files
    pi0_train_files = pi0_files[train_start:train_end]
    pi0_val_files = pi0_files[train_end:val_end]
    pion_train_files = pion_files[train_start:train_end]
    pion_val_files = pion_files[train_end:val_end]

    train_output_dir = None
    val_output_dir = None
            
    # Get Data
    if preprocess:
        train_output_dir = output_dir + '/train/'
        val_output_dir = output_dir + '/val/'

        if already_preprocessed:
            train_files = np.sort(glob.glob(train_output_dir+'*.p'))[:num_train_files]
            val_files = np.sort(glob.glob(val_output_dir+'*.p'))[:num_val_files]

            pi0_train_files = train_files
            pi0_val_files = val_files
            pion_train_files = None
            pion_val_files = None


            train_output_dir = None
            val_output_dir = None

    data_gen_train = MPGraphDataGenerator(pi0_file_list=pi0_train_files,
                                        pion_file_list=pion_train_files,
                                        cellGeo_file=data_dir+'graph_examples/cell_geo.root',
                                        batch_size=batch_size,
                                        shuffle=shuffle,
                                        num_procs=num_procs,
                                        preprocess=preprocess,
                                        output_dir=train_output_dir)

    data_gen_val = MPGraphDataGenerator(pi0_file_list=pi0_val_files,
                                        pion_file_list=pion_val_files,
                                        cellGeo_file=data_dir+'graph_examples/cell_geo.root',
                                        batch_size=batch_size,
                                        shuffle=shuffle,
                                        num_procs=num_procs,
                                        preprocess=preprocess,
                                        output_dir=val_output_dir)

    if preprocess and not already_preprocessed:
        exit()

    # Optimizer.
    #optimizer = snt.optimizers.Adam(learning_rate)
    optimizer = tf.keras.optimizers.Adam(learning_rate)

    model = models.MultiOutBlockWeightedRegressModel(global_output_size=1, num_outputs=2, model_config=model_config)

    training_loss_epoch = []
    training_loss_regress_epoch = []
    training_loss_class_epoch = []
    val_loss_epoch = []
    val_loss_regress_epoch = []
    val_loss_class_epoch = []
    
    checkpoint = tf.train.Checkpoint(module=model)
    checkpoint_prefix = os.path.join(save_dir, 'latest_model')
    latest = tf.train.latest_checkpoint(save_dir)
    if latest is not None:
        checkpoint.restore(latest)
    else:
        checkpoint.save(checkpoint_prefix)

    def convert_to_tuple(graphs):
        nodes = []
        edges = []
        globals = []
        senders = []
        receivers = []
        n_node = []
        n_edge = []
        offset = 0

        for graph in graphs:
            nodes.append(graph['nodes'])
            edges.append(graph['edges'])
            globals.append([graph['globals']])
            senders.append(graph['senders'] + offset)
            receivers.append(graph['receivers'] + offset)
            n_node.append(graph['nodes'].shape[:1])
            n_edge.append(graph['edges'].shape[:1])

            offset += len(graph['nodes'])

        nodes = tf.convert_to_tensor(np.concatenate(nodes))
        edges = tf.convert_to_tensor(np.concatenate(edges))
        globals = tf.convert_to_tensor(np.concatenate(globals))
        senders = tf.convert_to_tensor(np.concatenate(senders))
        receivers = tf.convert_to_tensor(np.concatenate(receivers))
        n_node = tf.convert_to_tensor(np.concatenate(n_node))
        n_edge = tf.convert_to_tensor(np.concatenate(n_edge))

        graph = GraphsTuple(
                nodes=nodes,
                edges=edges,
                globals=globals,
                senders=senders,
                receivers=receivers,
                n_node=n_node,
                n_edge=n_edge
            )

        return graph
       
    def get_batch(data_iter):
        for graphs, targets in data_iter:
            graphs = convert_to_tuple(graphs)
            targets = tf.convert_to_tensor(targets)
            
            yield graphs, targets

    samp_graph, samp_target = next(get_batch(data_gen_train.generator()))
    data_gen_train.kill_procs()
    graph_spec = utils_tf.specs_from_graphs_tuple(samp_graph, True, True, True)
    
    mae_loss = tf.keras.losses.MeanAbsoluteError()
    bce_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)

    def loss_fn(targets, regress_preds, class_preds):
        regress_loss = mae_loss(targets[:,:1], regress_preds)
        class_loss = bce_loss(targets[:,1:], class_preds)
        combined_loss = alpha*regress_loss + (1 - alpha)*class_loss 
        return regress_loss, class_loss, combined_loss

    @tf.function(input_signature=[graph_spec, tf.TensorSpec(shape=[None,2], dtype=tf.float32)])
    def train_step(graphs, targets):
        with tf.GradientTape() as tape:
            regress_output, class_output = model(graphs)
            regress_preds = regress_output.globals
            class_preds = class_output.globals
            regress_loss, class_loss, loss = loss_fn(targets, regress_preds, class_preds)
        
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        return regress_loss, class_loss, loss

    @tf.function(input_signature=[graph_spec, tf.TensorSpec(shape=[None,2], dtype=tf.float32)])
    def val_step(graphs, targets):
        regress_output, class_output = model(graphs)
        regress_preds = regress_output.globals
        class_preds = class_output.globals
        regress_loss, class_loss, loss = loss_fn(targets, regress_preds, class_preds)

        return regress_loss, class_loss, loss, regress_preds, class_preds

    curr_loss = 1e5
    for e in range(epochs):

        logging.info('\nStarting epoch: {}'.format(e))
        print('\nStarting epoch: {}'.format(e))
        epoch_start = time.time()

        training_loss = []
        training_loss_regress = []
        training_loss_class = []
        val_loss = []
        val_loss_regress = []
        val_loss_class = []

        # Train
        logging.info('Training...')
        i = 1
        for graph_data_tr, targets_tr in get_batch(data_gen_train.generator()):#train_iter):
            start = time.time()
            #if i==1:
            losses_tr_rg, losses_tr_cl, losses_tr = train_step(graph_data_tr, targets_tr)
            end = time.time()

            training_loss.append(losses_tr.numpy())
            training_loss_regress.append(losses_tr_rg.numpy())
            training_loss_class.append(losses_tr_cl.numpy())

            if not (i-1)%log_freq:
                logging.info('Iter: {:04d}, Tr_loss_mean: {:.4f}, Tr_loss_rg_mean: {:.4f}, Tr_loss_cl_mean: {:.4f}, Took {:.3f}secs'. \
                      format(i, 
                             np.mean(training_loss), 
                             np.mean(training_loss_regress), 
                             np.mean(training_loss_class), 
                             end-start))
                # logging.info('Took {:.3f} secs'.format(end-start))
            
            i += 1 

        training_loss_epoch.append(training_loss)
        training_loss_regress_epoch.append(training_loss_regress)
        training_loss_class_epoch.append(training_loss_class)
        training_end = time.time()

        # validate
        logging.info('\nValidation...')
        i = 1
        all_targets = []
        all_outputs = []
        for graph_data_val, targets_val in get_batch(data_gen_val.generator()):#val_iter):
            start = time.time()
            losses_val_rg, losses_val_cl, losses_val, regress_vals, class_vals = val_step(graph_data_val, targets_val)
            end = time.time()

            targets_val = targets_val.numpy()
            regress_vals = regress_vals.numpy()
            class_vals = class_vals.numpy()

            targets_val[:,0] = 10**targets_val[:,0]
            regress_vals = 10**regress_vals
            # targets_val[:,1] = 1 / (1 + np.exp(targets_val[:,1]))
            class_vals =  tf.math.sigmoid(class_vals)   # 1 / (1 + np.exp(class_vals))

            output_vals = np.hstack([regress_vals, class_vals])

            val_loss.append(losses_val.numpy())
            val_loss_regress.append(losses_val_rg.numpy())
            val_loss_class.append(losses_val_cl.numpy())

            all_targets.append(targets_val)
            all_outputs.append(output_vals)

            if not (i-1)%log_freq:
                logging.info('Iter: {:04d}, Val_loss_mean: {:.4f}, Val_loss_rg_mean: {:.4f}, Val_loss_cl_mean: {:.4f}, Took {:.3f}secs'. \
                      format(i, 
                             np.mean(val_loss), 
                             np.mean(val_loss_regress), 
                             np.mean(val_loss_class), 
                             end-start))
                # logging.info('Took {:.3f} secs'.format(end-start))
            
            i += 1

        epoch_end = time.time()


        all_targets = np.concatenate(all_targets)
        all_outputs = np.concatenate(all_outputs)
        
        val_loss_epoch.append(val_loss)
        val_loss_regress_epoch.append(val_loss_regress)
        val_loss_class_epoch.append(val_loss_class)
    
        np.savez(save_dir+'/losses', 
                training=training_loss_epoch, validation=val_loss_epoch,
                training_regress=training_loss_regress_epoch, validation_regress=val_loss_regress_epoch,
                training_class=training_loss_class_epoch, validation_class=val_loss_class_epoch,
                )
        # checkpoint.save(checkpoint_prefix)
        
        val_mins = int((epoch_end - training_end)/60)
        val_secs = int((epoch_end - training_end)%60)
        training_mins = int((training_end - epoch_start)/60)
        training_secs = int((training_end - epoch_start)%60)

        logging.info('\nEpoch {} ended\nTraining: {:2d}:{:02d}\nValidation: {:2d}:{:02d}'. \
                     format(e, training_mins, training_secs, val_mins, val_secs))
        print('\nEpoch {} ended\nTraining: {:2d}:{:02d}\nValidation: {:2d}:{:02d}'. \
              format(e, training_mins, training_secs, val_mins, val_secs))

        if np.mean(val_loss)<curr_loss:
            logging.info('\nLoss decreased from {:.4f} to {:.4f}'.format(curr_loss, np.mean(val_loss)))
            logging.info('Checkpointing and saving predictions to:\n{}'.format(save_dir))
            print('\nLoss decreased from {:.4f} to {:.4f}'.format(curr_loss, np.mean(val_loss)))
            print('Checkpointing and saving predictions to:\n{}'.format(save_dir))
            curr_loss = np.mean(val_loss)
            np.savez(save_dir+'/predictions', 
                    targets=all_targets, 
                    outputs=all_outputs)
            checkpoint.save(checkpoint_prefix)
        else: 
            logging.info('\nLoss didnt decrease from {:.4f}'.format(curr_loss))
            print('\nLoss didnt decrease from {:.4f}'.format(curr_loss))

        if not (e+1)%20:
            optimizer.learning_rate = optimizer.learning_rate/10
            logging.info('\nLearning rate decreased to: {:.3e}'.format(optimizer.learning_rate.value()))
            print('\nLearning rate decreased to: {:.3e}'.format(optimizer.learning_rate.value()))


tensorflow.convert_to_tensor numpy.hstack numpy.savez tensorflow.train.latest_checkpoint tensorflow.train.Checkpoint tensorflow.keras.losses.BinaryCrossentropy numpy.concatenate tensorflow.math.sigmoid tensorflow.GradientTape tensorflow.keras.optimizers.Adam numpy.mean tensorflow.TensorSpec tensorflow.keras.losses.MeanAbsoluteError
[(21, 'seaborn.set_context', 'sns.set_context', (['"""poster"""'], {}), True, 'import seaborn as sns\n'), (25, 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), False, 'import argparse\n'), (53, 'os.makedirs', 'os.makedirs', (['save_dir'], {'exist_ok': '(True)'}), False, 'import os\n'), (56, 'logging.basicConfig', 'logging.basicConfig', ([], {'level': 'logging.INFO', 'format': '"""%(message)s"""', 'filename': "(save_dir + '/output.log')"}), False, 'import logging\n'), (93, 'modules.mpdatagen.MPGraphDataGenerator', 'MPGraphDataGenerator', ([], {'pi0_file_list': 'pi0_train_files', 'pion_file_list': 'pion_train_files', 'cellGeo_file': "(data_dir + 'graph_examples/cell_geo.root')", 'batch_size': 'batch_size', 'shuffle': 'shuffle', 'num_procs': 'num_procs', 'preprocess': 'preprocess', 'output_dir': 'train_output_dir'}), False, 'from modules.mpdatagen import MPGraphDataGenerator\n'), (102, 'modules.mpdatagen.MPGraphDataGenerator', 'MPGraphDataGenerator', ([], {'pi0_file_list': 'pi0_val_files', 'pion_file_list': 'pion_val_files', 'cellGeo_file': "(data_dir + 'graph_examples/cell_geo.root')", 'batch_size': 'batch_size', 'shuffle': 'shuffle', 'num_procs': 'num_procs', 'preprocess': 'preprocess', 'output_dir': 'val_output_dir'}), False, 'from modules.mpdatagen import MPGraphDataGenerator\n'), (116, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', (['learning_rate'], {}), True, 'import tensorflow as tf\n'), (118, 'modules.multiOutBlock_wWeightedRegress.MultiOutBlockWeightedRegressModel', 'models.MultiOutBlockWeightedRegressModel', ([], {'global_output_size': '(1)', 'num_outputs': '(2)', 'model_config': 'model_config'}), True, 'import modules.multiOutBlock_wWeightedRegress as models\n'), (127, 'tensorflow.train.Checkpoint', 'tf.train.Checkpoint', ([], {'module': 'model'}), True, 'import tensorflow as tf\n'), (128, 'os.path.join', 'os.path.join', (['save_dir', '"""latest_model"""'], {}), False, 'import os\n'), (129, 'tensorflow.train.latest_checkpoint', 'tf.train.latest_checkpoint', (['save_dir'], {}), True, 'import tensorflow as tf\n'), (185, 'graph_nets.utils_tf.specs_from_graphs_tuple', 'utils_tf.specs_from_graphs_tuple', (['samp_graph', '(True)', '(True)', '(True)'], {}), False, 'from graph_nets import utils_tf\n'), (187, 'tensorflow.keras.losses.MeanAbsoluteError', 'tf.keras.losses.MeanAbsoluteError', ([], {}), True, 'import tensorflow as tf\n'), (188, 'tensorflow.keras.losses.BinaryCrossentropy', 'tf.keras.losses.BinaryCrossentropy', ([], {'from_logits': '(True)'}), True, 'import tensorflow as tf\n'), (62, 'glob.glob', 'glob.glob', (["(data_dir + '*graphs.v01*/*pi0*/*root')"], {}), False, 'import glob\n'), (63, 'glob.glob', 'glob.glob', (["(data_dir + '*graphs.v01*/*pion*/*root')"], {}), False, 'import glob\n'), (164, 'graph_nets.graphs.GraphsTuple', 'GraphsTuple', ([], {'nodes': 'nodes', 'edges': 'edges', 'globals': 'globals', 'senders': 'senders', 'receivers': 'receivers', 'n_node': 'n_node', 'n_edge': 'n_edge'}), False, 'from graph_nets.graphs import GraphsTuple\n'), (223, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (233, 'logging.info', 'logging.info', (['"""Training..."""'], {}), False, 'import logging\n'), (259, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (262, 'logging.info', 'logging.info', (['"""\nValidation..."""'], {}), False, 'import logging\n'), (300, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (303, 'numpy.concatenate', 'np.concatenate', (['all_targets'], {}), True, 'import numpy as np\n'), (304, 'numpy.concatenate', 'np.concatenate', (['all_outputs'], {}), True, 'import numpy as np\n'), (310, 'numpy.savez', 'np.savez', (["(save_dir + '/losses')"], {'training': 'training_loss_epoch', 'validation': 'val_loss_epoch', 'training_regress': 'training_loss_regress_epoch', 'validation_regress': 'val_loss_regress_epoch', 'training_class': 'training_loss_class_epoch', 'validation_class': 'val_loss_class_epoch'}), True, 'import numpy as np\n'), (156, 'numpy.concatenate', 'np.concatenate', (['nodes'], {}), True, 'import numpy as np\n'), (157, 'numpy.concatenate', 'np.concatenate', (['edges'], {}), True, 'import numpy as np\n'), (158, 'numpy.concatenate', 'np.concatenate', (['globals'], {}), True, 'import numpy as np\n'), (159, 'numpy.concatenate', 'np.concatenate', (['senders'], {}), True, 'import numpy as np\n'), (160, 'numpy.concatenate', 'np.concatenate', (['receivers'], {}), True, 'import numpy as np\n'), (161, 'numpy.concatenate', 'np.concatenate', (['n_node'], {}), True, 'import numpy as np\n'), (162, 'numpy.concatenate', 'np.concatenate', (['n_edge'], {}), True, 'import numpy as np\n'), (179, 'tensorflow.convert_to_tensor', 'tf.convert_to_tensor', (['targets'], {}), True, 'import tensorflow as tf\n'), (198, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (236, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (239, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (267, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (269, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (278, 'tensorflow.math.sigmoid', 'tf.math.sigmoid', (['class_vals'], {}), True, 'import tensorflow as tf\n'), (280, 'numpy.hstack', 'np.hstack', (['[regress_vals, class_vals]'], {}), True, 'import numpy as np\n'), (327, 'numpy.mean', 'np.mean', (['val_loss'], {}), True, 'import numpy as np\n'), (332, 'numpy.mean', 'np.mean', (['val_loss'], {}), True, 'import numpy as np\n'), (333, 'numpy.savez', 'np.savez', (["(save_dir + '/predictions')"], {'targets': 'all_targets', 'outputs': 'all_outputs'}), True, 'import numpy as np\n'), (52, 'time.strftime', 'time.strftime', (['"""%Y%m%d_%H%M"""'], {}), False, 'import time\n'), (196, 'tensorflow.TensorSpec', 'tf.TensorSpec', ([], {'shape': '[None, 2]', 'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (209, 'tensorflow.TensorSpec', 'tf.TensorSpec', ([], {'shape': '[None, 2]', 'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (81, 'glob.glob', 'glob.glob', (["(train_output_dir + '*.p')"], {}), False, 'import glob\n'), (82, 'glob.glob', 'glob.glob', (["(val_output_dir + '*.p')"], {}), False, 'import glob\n'), (328, 'numpy.mean', 'np.mean', (['val_loss'], {}), True, 'import numpy as np\n'), (330, 'numpy.mean', 'np.mean', (['val_loss'], {}), True, 'import numpy as np\n'), (248, 'numpy.mean', 'np.mean', (['training_loss'], {}), True, 'import numpy as np\n'), (249, 'numpy.mean', 'np.mean', (['training_loss_regress'], {}), True, 'import numpy as np\n'), (250, 'numpy.mean', 'np.mean', (['training_loss_class'], {}), True, 'import numpy as np\n'), (292, 'numpy.mean', 'np.mean', (['val_loss'], {}), True, 'import numpy as np\n'), (293, 'numpy.mean', 'np.mean', (['val_loss_regress'], {}), True, 'import numpy as np\n'), (294, 'numpy.mean', 'np.mean', (['val_loss_class'], {}), True, 'import numpy as np\n')]