#!/usr/bin/env python
# -*- coding: utf-8 -*-
import argparse
import glob
import json
import logging
import os
import sys
import time

SOURCE_DIR = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
sys.path.append(SOURCE_DIR)

import matplotlib
matplotlib.use('agg')
import numpy as np

import matplotlib.pyplot as plt

from alignment.features import (extract_features,
                                DEFAULT_FEATURE_KWARGS,
                                SAMPLE_RATE,
                                FRAME_SIZE,
                                HOP_SIZE)
from alignment.dtw import dtw_alignment, DEFAULT_DTW_KWARGS

from analysis.utils import score_beat_map, get_beat_times

logging.basicConfig(level=logging.INFO)
LOGGER = logging.getLogger(__name__)

def _perf_name(fn):
    """Gets name of the piece, replaces spaces with '_' and removes extension
    """
    return os.path.basename(fn).replace(' ', '_').replace('.flac', '')

def main(performance_fn, reference_fn, out_dir,
         beat_fn=None,
         sample_rate=SAMPLE_RATE,
         frame_size=FRAME_SIZE,
         hop_size=HOP_SIZE,
         make_plots=True,
         feature_kwargs=DEFAULT_FEATURE_KWARGS,
         dtw_kwargs=DEFAULT_DTW_KWARGS):
    """
    Perform an alignment

    Parameters
    ----------
    performance_fn : str
        Audio file with the performance
    reference_fn : str
        Audio file with the reference performance.
    out_dir : str
        Directory to store alignment results
    sample_rate : int
    frame_size : int
    hop_size : int
    make_plots : bool
    feature_kwargs : dict
    dtw_kwargs : dict
    """
    ###### Feature extraction ######
    features = feature_kwargs.pop('features')
    LOGGER.info('Exctracting {0} features from performance'.format(features))
    perf_features, perf_times = extract_features(performance_fn,
                                                 features=features,
                                                 feature_kwargs=feature_kwargs,
                                                 sample_rate=sample_rate,
                                                 frame_size=frame_size,
                                                 hop_size=hop_size)
    LOGGER.info('Extracting {0} features from reference'.format(features))
    ref_features, ref_times = extract_features(reference_fn,
                                               features=features,
                                               feature_kwargs=feature_kwargs,
                                               sample_rate=sample_rate,
                                               frame_size=frame_size,
                                               hop_size=hop_size)

    if beat_fn is not None:
        beats = np.loadtxt(beat_fn)
        score_beats = beats[:, 1]
        beats = beats[:, 0]

        beat_map, beat_index_map = score_beat_map(ref_times,
                                                  beats=beats,
                                                  score_beats=score_beats)
    else:
        beats, beat_map, beat_index_map, score_beats = None, None, None, None
    

    if make_plots:
        LOGGER.info('Plotting features')
        fig, axes = plt.subplots(2)

        pxticks = np.arange(0, len(perf_features), len(perf_features) // 10)
        pxtick_labels = ['{0:.0f}'.format(np.round(pt / 60)) for pt in perf_times[pxticks]]
        rxticks = np.arange(0, len(ref_features), len(ref_features) // 10)
        rxtick_labels = ['{0:.0f}'.format(np.round(pt / 60)) for pt in ref_times[rxticks]]

        axes[0].imshow(perf_features.T, aspect='auto',
                       interpolation='nearest',
                       origin='lowest',
                       cmap='BuPu')
        axes[0].set_title(_perf_name(performance_fn))
        axes[0].set_ylabel('Features')
        axes[0].set_xticks(pxticks)
        axes[0].set_xticklabels(pxtick_labels)
        axes[1].imshow(ref_features.T, aspect='auto',
                       interpolation='nearest',
                       origin='lowest',
                       cmap='BuPu')
        axes[1].set_xlabel('Time (minutes)')
        axes[1].set_ylabel('Features')
        axes[1].set_xticks(rxticks)
        axes[1].set_xticklabels(rxtick_labels)
        axes[1].set_title(_perf_name(reference_fn))
        plt.tight_layout()

        fig_fn = os.path.join(out_dir,
                              'features_p_{0}_r_{1}.pdf'.format(_perf_name(performance_fn),
                                                                _perf_name(reference_fn)))
        plt.savefig(fig_fn)
        plt.clf()
        plt.close()

    ###### Compute alignment ######
    LOGGER.info('Computing Alignment')
    # Compute pairwise cost
    wp, D = dtw_alignment(perf_features=perf_features,
                          ref_features=ref_features,
                          dtw_kwargs=dtw_kwargs)
    out_fn = os.path.join(out_dir,
                          'alignment_p_{0}_r_{1}.txt'.format(_perf_name(performance_fn),
                                                             _perf_name(reference_fn)))
    if beat_map is None:
        alignment = np.column_stack((perf_times[wp[:, 0]], ref_times[wp[:, 1]]))
    else:
        alignment = np.column_stack((perf_times[wp[:, 0]],
                                     ref_times[wp[:, 1]],
                                     beat_map(ref_times[wp[:, 1]])))
        est_beat_times = get_beat_times(alignment=alignment,
                                        s_beats=score_beats)
        np.savetxt(out_fn.replace('.txt', '_beat_times.txt'), est_beat_times)
    np.savetxt(out_fn, alignment, delimiter='\t')

    if make_plots:
        LOGGER.info('Plotting alignment')
        
        try:
            # if Distance is not matrix
            plt.matshow(D, aspect='auto',
                        interpolation='nearest',
                       cmap='gray',
                       origin='lowest')
        except:
            pass

        plt.plot(wp[:, 1],  wp[:, 0], label='Optimal path', color='y')

        plt.xlabel(_perf_name(reference_fn))
        plt.ylabel(_perf_name(performance_fn))
        plot_fn = os.path.join(out_dir,
                              'alignment_p_{0}_r_{1}.pdf'.format(_perf_name(performance_fn),
                                                                _perf_name(reference_fn)))
        plt.savefig(plot_fn)
        plt.clf()
        plt.close()
    

if __name__ == '__main__':

    parser = argparse.ArgumentParser('Align a piece')

    parser.add_argument('performance_fn', help='Audio file with the piece to be aligned')
    parser.add_argument('reference_fn', help='Audio file with the reference performance')
    parser.add_argument('out_dir', help='Directory for storing the outputs',
                        default='.')
    parser.add_argument('--sr', help='Sample rate (default {0}Hz)'.format(SAMPLE_RATE),
                        type=int, default=SAMPLE_RATE)
    parser.add_argument('--fs', help='Frame size (Default {0})'.format(FRAME_SIZE),
                        type=int, default=FRAME_SIZE)
    parser.add_argument('--hs', help='Hop size (Default {0})'.format(HOP_SIZE),
                        type=int, default=HOP_SIZE)
    parser.add_argument('--features-config',
                        help='Config file for computing audio features (JSON format)',
                        default=None)
    parser.add_argument('--no-plots', help='Do not generate plots',
                        action='store_true', default=False)
    parser.add_argument('--dtw-config', help='Config file for DTW (JSON format)',
                        default=None)
    parser.add_argument('--beat-fn', help='Beat file',
                        default=None)

    args = parser.parse_args()

    if not os.path.exists(args.out_dir):
        os.mkdir(args.out_dir)

    if args.features_config is not None:
        feature_kwargs = json.load(open(args.features_config))
    else:
        feature_kwargs = DEFAULT_FEATURE_KWARGS

    if args.dtw_config is not None:
        dtw_kwargs = json.load(open(args.dtw_config))
    else:
        dtw_kwargs = DEFAULT_DTW_KWARGS

    main(performance_fn=args.performance_fn,
         reference_fn=args.reference_fn,
         beat_fn=args.beat_fn,
         out_dir=args.out_dir,
         make_plots=not args.no_plots,
         sample_rate=args.sr,
         frame_size=args.fs,
         hop_size=args.hs,
         feature_kwargs=feature_kwargs,
         dtw_kwargs=dtw_kwargs)
    

    

