"""This script trains a model."""
import os
import logging
import argparse
from pprint import pformat
import numpy as np
import scipy.stats
import tensorflow as tf
from musegan.config import LOGLEVEL, LOG_FORMAT
from musegan.data import load_data, get_dataset, get_samples
from musegan.metrics import get_save_metric_ops
from musegan.model import Model
from musegan.utils import make_sure_path_exists, load_yaml
from musegan.utils import backup_src, update_not_none, setup_loggers
LOGGER = logging.getLogger("musegan.train")

def parse_arguments():
    """Parse and return the command line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument('--exp_dir', help="Directory to save all the results.")
    parser.add_argument('--params', help="Path to the model parameter file.")
    parser.add_argument('--config', help="Path to the configuration file.")
    parser.add_argument('--gpu', '--gpu_device_num', type=str, default="0",
                        help="The GPU device number to use.")
    parser.add_argument('--n_jobs', type=int,
                        help="Number of parallel calls to use for input "
                             "pipeline. Set to 1 to disable multiprocessing.")
    args = parser.parse_args()
    return args

def setup_dirs(config):
    """Setup an experiment directory structure and update the `params`
    dictionary with the directory paths."""
    # Get experiment directory structure
    config['exp_dir'] = os.path.realpath(config['exp_dir'])
    config['src_dir'] = os.path.join(config['exp_dir'], 'src')
    config['eval_dir'] = os.path.join(config['exp_dir'], 'eval')
    config['model_dir'] = os.path.join(config['exp_dir'], 'model')
    config['sample_dir'] = os.path.join(config['exp_dir'], 'samples')
    config['log_dir'] = os.path.join(config['exp_dir'], 'logs', 'train')

    # Make sure directories exist
    for key in ('log_dir', 'model_dir', 'sample_dir', 'src_dir'):
        make_sure_path_exists(config[key])

def setup():
    """Parse command line arguments, load model parameters, load configurations,
    setup environment and setup loggers."""
    # Parse the command line arguments
    args = parse_arguments()

    # Load parameters
    params = load_yaml(args.params)
    if params.get('is_accompaniment') and params.get('condition_track_idx') is None:
        raise TypeError("`condition_track_idx` cannot be None type in "
                        "accompaniment mode.")

    # Load configurations
    config = load_yaml(args.config)
    update_not_none(config, vars(args))

    # Set unspecified schedule steps to default values
    for target in (config['learning_rate_schedule'], config['slope_schedule']):
        if target['start'] is None:
            target['start'] = 0
        if target['end'] is None:
            target['end'] = config['steps']

    # Setup experiment directories and update them to configuration dictionary
    setup_dirs(config)

    # Setup loggers
    del logging.getLogger('tensorflow').handlers[0]
    setup_loggers(config['log_dir'])

    # Setup GPUs
    os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']

    # Backup source code
    backup_src(config['src_dir'])

    return params, config

def load_training_data(params, config):
    """Load and return the training data."""
    # Load data
    if params['is_conditional']:
        raise ValueError("Not supported yet.")
    else:
        labels = None
    LOGGER.info("Loading training data.")
    data = load_data(config['data_source'], config['data_filename'])
    LOGGER.info("Training data size: %d", len(data))

    # Build dataset
    LOGGER.info("Building dataset.")
    dataset = get_dataset(
        data, labels, config['batch_size'], params['data_shape'],
        config['use_random_transpose'], config['n_jobs'])

    # Create iterator
    if params['is_conditional']:
        train_x, train_y = dataset.make_one_shot_iterator().get_next()
    else:
        train_x, train_y = dataset.make_one_shot_iterator().get_next(), None

    return train_x, train_y

def load_or_create_samples(params, config):
    """Load or create the samples used as the sampler inputs."""
    # Load sample_z
    LOGGER.info("Loading sample_z.")
    sample_z_path = os.path.join(config['model_dir'], 'sample_z.npy')
    if os.path.exists(sample_z_path):
        sample_z = np.load(sample_z_path)
        if sample_z.shape[1] != params['latent_dim']:
            LOGGER.info("Loaded sample_z has wrong shape")
            resample = True
        else:
            resample = False
    else:
        LOGGER.info("File for sample_z not found")
        resample = True

    # Draw new sample_z
    if resample:
        LOGGER.info("Drawing new sample_z.")
        sample_z = scipy.stats.truncnorm.rvs(
            -2, 2, size=(np.prod(config['sample_grid']), params['latent_dim']))
        make_sure_path_exists(config['model_dir'])
        np.save(sample_z_path, sample_z)

    if params.get('is_accompaniment'):
        # Load sample_x
        LOGGER.info("Loading sample_x.")
        sample_x_path = os.path.join(config['model_dir'], 'sample_x.npy')
        if os.path.exists(sample_x_path):
            sample_x = np.load(sample_x_path)
            if sample_x.shape[1:] != params['data_shape']:
                LOGGER.info("Loaded sample_x has wrong shape")
                resample = True
            else:
                resample = False
        else:
            LOGGER.info("File for sample_x not found")
            resample = True

        # Draw new sample_x
        if resample:
            LOGGER.info("Drawing new sample_x.")
            data = load_data(config['data_source'], config['data_filename'])
            sample_x = get_samples(
                np.prod(config['sample_grid']), data,
                use_random_transpose = config['use_random_transpose'])
            make_sure_path_exists(config['model_dir'])
            np.save(sample_x_path, sample_x)
    else:
        sample_x = None

    return sample_x, None, sample_z

def main():
    """Main function."""
    # Setup
    logging.basicConfig(level=LOGLEVEL, format=LOG_FORMAT)
    params, config = setup()
    LOGGER.info("Using parameters:\n%s", pformat(params))
    LOGGER.info("Using configurations:\n%s", pformat(config))

    # ================================== Data ==================================
    # Load training data
    train_x, _ = load_training_data(params, config)

    # ================================= Model ==================================
    # Build model
    model = Model(params)
    if params.get('is_accompaniment'):
        train_c = tf.expand_dims(
            train_x[..., params['condition_track_idx']], -1)
        train_nodes = model(
            x=train_x, c=train_c, mode='train', params=params, config=config)
    else:
        train_nodes = model(
            x=train_x, mode='train', params=params, config=config)

    # Log number of parameters in the model
    def get_n_params(var_list):
        """Return the number of variables in a variable list."""
        return int(np.sum([np.product(
            [x.value for x in var.get_shape()]) for var in var_list]))

    LOGGER.info("Number of trainable parameters in {}: {:,}".format(
        model.name, get_n_params(tf.trainable_variables(model.name))))
    for component in model.components:
        LOGGER.info("Number of trainable parameters in {}: {:,}".format(
            component.name, get_n_params(tf.trainable_variables(
                model.name + '/' + component.name))))

    # ================================ Sampler =================================
    if config['save_samples_steps'] > 0:
        # Get sampler inputs
        sample_x, sample_y, sample_z = load_or_create_samples(params, config)

        # Create sampler configurations
        sampler_config = {
            'result_dir': config['sample_dir'],
            'suffix': tf.as_string(train_nodes['gen_step']),
            'image_grid': config['sample_grid'],
            'colormap': np.array(config['colormap']).T,
            'midi': config['midi'],
            'collect_save_arrays_op': config['save_array_samples'],
            'collect_save_images_op': config['save_image_samples'],
            'collect_save_pianorolls_op': config['save_pianoroll_samples']}

        # Get prediction nodes
        placeholder_z = tf.placeholder(tf.float32, shape=sample_z.shape)
        placeholder_y = None
        if params.get('is_accompaniment'):
            c_shape = np.append(sample_x.shape[:-1], 1)
            placeholder_c = tf.placeholder(tf.float32, shape=c_shape)
            predict_nodes = model(
                z=placeholder_z, y=placeholder_y, c=placeholder_c,
                mode='predict', params=params, config=sampler_config)
        else:
            predict_nodes = model(
                z=placeholder_z, y=placeholder_y, mode='predict', params=params,
                config=sampler_config)

        # Get sampler op
        sampler_op = tf.group([
            predict_nodes[key] for key in (
                'save_arrays_op', 'save_images_op', 'save_pianorolls_op')
            if key in predict_nodes])
        sampler_op_no_pianoroll = tf.group([
            predict_nodes[key] for key in ('save_arrays_op', 'save_images_op')
            if key in predict_nodes])

    # ================================ Metrics =================================
    if config['evaluate_steps'] > 0:
        binarized = tf.round(.5 * (predict_nodes['fake_x'] + 1.)) ## 从-1 到 1 ？ 数值取整
        save_metric_ops = get_save_metric_ops(
            binarized, params['beat_resolution'], train_nodes['gen_step'],
            config['eval_dir'])
        save_metrics_op = tf.group(save_metric_ops)

    # ========================== Training Preparation ==========================
    # Get tensorflow session config
    tf_config = tf.ConfigProto()
    tf_config.gpu_options.allow_growth = True

    # Training hooks
    global_step = tf.train.get_global_step()
    steps_per_iter = config['n_dis_updates_per_gen_update'] + 1
    hooks = [tf.train.NanTensorHook(train_nodes['loss'])]

    # Tensor logger
    tensor_logger = {
        'step': train_nodes['gen_step'],
        'gen_loss': train_nodes['gen_loss'],
        'dis_loss': train_nodes['dis_loss']}
    step_logger = open(os.path.join(config['log_dir'], 'step.log'), 'w')

    # ======================= Monitored Training Session =======================
    LOGGER.info("Training start.")
    with tf.train.MonitoredTrainingSession(
        save_checkpoint_steps=config['save_checkpoint_steps'] * steps_per_iter,
        save_summaries_steps=config['save_summaries_steps'] * steps_per_iter,
        checkpoint_dir=config['model_dir'], log_step_count_steps=0,
        hooks=hooks, config=tf_config) as sess:

        # Get global step value
        step = tf.train.global_step(sess, global_step)
        if step == 0:
            step_logger.write('# step, gen_loss, dis_loss\n')

        # ============================== Training ==============================
        if step >= config['steps']:
            LOGGER.info("Global step has already exceeded total steps.")
            step_logger.close()
            return

        # Training iteration
        while step < config['steps']:

            # Train the discriminator
            if step < 10:
                n_dis_updates = 10 * config['n_dis_updates_per_gen_update']
            else:
                n_dis_updates = config['n_dis_updates_per_gen_update']
            for _ in range(n_dis_updates):
                sess.run(train_nodes['train_ops']['dis'])

            # Train the generator
            log_loss_steps = config['log_loss_steps'] or 100
            if (step + 1) % log_loss_steps == 0:
                step, _, tensor_logger_values = sess.run([
                    train_nodes['gen_step'], train_nodes['train_ops']['gen'],
                    tensor_logger])
                # Logger
                if config['log_loss_steps'] > 0:
                    LOGGER.info("step={}, {}".format(
                        tensor_logger_values['step'], ', '.join([
                            '{}={: 8.4E}'.format(key, value)
                            for key, value in tensor_logger_values.items()
                            if key != 'step'])))
                step_logger.write("{}, {: 10.6E}, {: 10.6E}\n".format(
                    tensor_logger_values['step'],
                    tensor_logger_values['gen_loss'],
                    tensor_logger_values['dis_loss']))
            else:
                step, _ = sess.run([
                    train_nodes['gen_step'], train_nodes['train_ops']['gen']])

            # Run sampler
            if ((config['save_samples_steps'] > 0)
                    and (step % config['save_samples_steps'] == 0)):
                LOGGER.info("Running sampler")
                feed_dict_sampler = {placeholder_z: sample_z}
                if params.get('is_accompaniment'):
                    feed_dict_sampler[placeholder_c] = np.expand_dims(
                        sample_x[..., params['condition_track_idx']], -1)
                if step < 3000:
                    sess.run(
                        sampler_op_no_pianoroll, feed_dict=feed_dict_sampler)
                else:
                    sess.run(sampler_op, feed_dict=feed_dict_sampler)

            # Run evaluation
            if ((config['evaluate_steps'] > 0)
                    and (step % config['evaluate_steps'] == 0)):
                LOGGER.info("Running evaluation")
                feed_dict_evaluation = {
                    placeholder_z: scipy.stats.truncnorm.rvs(-2, 2, size=(
                        np.prod(config['sample_grid']), params['latent_dim']))}
                if params.get('is_accompaniment'):
                    feed_dict_evaluation[placeholder_c] = np.expand_dims(
                        sample_x[..., params['condition_track_idx']], -1)
                sess.run(save_metrics_op, feed_dict=feed_dict_evaluation)

            # Stop training if stopping criterion suggests
            if sess.should_stop():
                break

    LOGGER.info("Training end")
    step_logger.close()

if __name__ == "__main__":
    main()