"""Using annotated onsets, segment source audio files into individual
audio files, with each containing a single 'note'.

Usage:
 segment_audio.py [options] <segment_index> <note_index> <note_audio_dir>

Example:

Arguments:
 segment_index   Input index file, as generated by `collect_data.py`.
 note_index      csv file which will contain the record of all notes
                 created.
 note_audio_dir  Directory where note audio will be stored.

Options:
 -h, --help      Print help.
 --dry-run       Don't actually run sox; just print the function call.
 -v, --verbose   Increase verbosity level.
 --limit=N_FILES  Limit to procesing N_FILES for testing.
 --duration=DUR  Desired (fixed) output duration, in DUR sec, of notes.
"""
from __future__ import print_function

import claudio
from docopt import docopt
import logging
import logging.config
import os
import pandas as pd
import sys
import time

import minst.logger
import minst.model as model
import minst.signal as signal
import minst.utils as utils

logger = logging.getLogger("split_audio_to_clips")

PRINT_PROGRESS = True


def audio_to_observations(index, audio_file, onsets_file, note_audio_dir,
                          file_ext='flac', note_duration=None, **meta):
    """Segment an audio file given an onset file, writing outputs to disk.

    Paramaters
    ----------
    audio_file : str
        Source audio file.

    onsets_file : str
        Path to a CSV file of cut points.

    note_audio_dir : str
        Path at which to write outputs.

    file_ext : str
        Desired output audio format for note files.

    note_duration : float, or default=None
        Desired duration of the output note files; if None, no fixed duration
        is applied.

    **meta : keyword args
        Additional record data to pass on to each observation; see
        model.Observation for more detail.

    Returns
    -------
    note_files : list of str
        Collection of paths on disk of generated outputs. These will take the
        following format: {note_audio_dir}/{index}.{file_ext}
    """

    # Get the soxi information on this file to get the Duration
    max_length = float(claudio.sox.soxi(audio_file, 'D'))

    # load the onset file.
    onsets = pd.read_csv(onsets_file, index_col=0)
    if onsets.empty:
        logger.warning(
            "Onset File is empty! We can't extract notes without "
            "onsets, so skipping: {}".format(os.path.basename(onsets_file)))
        return []

    # Append the duration to the end of the offsets so we can
    # do this by pairs.
    onsets.loc[onsets.size] = max_length

    # Make sure it's sorted by time now.
    # TODO: Do we really want to drop the index here?
    onsets = onsets.sort_values('time').reset_index(drop=True)
    logger.debug("Attempting to generate {} observations".format(len(onsets)))
    observations = []
    # for each pair of onsets
    for i in range(len(onsets) - 1):
        start_time = onsets.iloc[i]['time']
        if start_time < 0.0:
            start_time = 0.0
        end_time = onsets.iloc[i + 1]['time']
        if end_time > max_length:
            end_time = max_length

        clip_index = utils.generate_id(
            index, "{}".format(start_time), hash_len=6)

        rel_output_file = "{}.{}".format(clip_index, file_ext.strip('.'))
        output_file = os.path.join(note_audio_dir, rel_output_file)

        if signal.extract_clip(audio_file, output_file, start_time,
                               end_time, note_duration):
            obs = model.Observation(
                index=clip_index, audio_file=rel_output_file,
                source_index=index, start_time=start_time,
                duration=end_time - start_time, **meta)
            observations.append(obs)
            logger.debug("New Observation: {}".format(obs.to_builtin()))

    return observations


def audio_collection_to_observations(segment_index_file, note_index_file,
                                     note_audio_dir, limit_n_files=None,
                                     note_duration=None):
    """
    Parameters
    ----------
    segment_index_file : str
        Input file containing all pointers to audio files and
        onsets files.

    note_index_file: str
        Path to the output index file which will contain pointers
        to the output note audio, and the metadata relating to it.

    note_audio_dir : str
        Path to store the resulting audio file.

    Returns
    -------
    success : bool
        True if the method completed as expected.
    """
    logger.info("Begin audio collection segmentation")
    logger.debug("Loading segment index")
    segment_df = pd.read_csv(segment_index_file, index_col=0)
    logger.debug("loaded {} records.".format(len(segment_df)))

    if segment_df.empty:
        logger.warning(utils.colorize(
            "No data available in {}; exiting.".format(segment_index_file),
            color='red'))
        # Here, we sys.exit 0 so the makefile will continue to build
        # other datasets, even if this one
        return True

    # Drop rows that do not have onsets_files.
    segment_df = segment_df.loc[segment_df.onsets_file.dropna().index]
    utils.create_directory(note_audio_dir)
    count = 0
    observations = []
    for idx, row in segment_df.iterrows():
        if pd.isnull(row.onsets_file):
            logger.warning("No onset file for {} [{}]; moving on.".format(
                row.audio_file, row.dataset))
            continue
        observations += audio_to_observations(
            idx, row.audio_file, row.onsets_file, note_audio_dir,
            file_ext='flac', dataset=row.dataset, instrument=row.instrument,
            dynamic=row.dynamic, note_duration=note_duration)
        logger.debug("Generated {} observations ({} of {}).".format(
            len(observations), (count + 1), len(segment_df)))

        if PRINT_PROGRESS:
            print("Progress: {:0.1f}% ({} of {})\r".format(
                (((count + 1) / float(len(segment_df))) * 100.),
                (count + 1), len(segment_df)), end='')
            sys.stdout.flush()
        count += 1

        if limit_n_files and count >= limit_n_files:
            break

    if PRINT_PROGRESS:
        print()

    collection = model.Collection(observations)
    collection.to_dataframe().to_csv(note_index_file)
    logger.debug("Wrote note index to {} with {} records".format(
        note_index_file, len(collection)))
    logger.info("Completed audio collection segmentation")
    return os.path.exists(note_index_file)


if __name__ == "__main__":
    arguments = docopt(__doc__)

    level = 'INFO' if not arguments.get('--verbose') else 'DEBUG'
    logging.config.dictConfig(minst.logger.get_config(level))
    PRINT_PROGRESS = not arguments['--verbose']

    t0 = time.time()
    audio_collection_to_observations(
        arguments['<segment_index>'],
        arguments['<note_index>'],
        arguments['<note_audio_dir>'],
        int(arguments['--limit']) if arguments['--limit'] else None,
        float(arguments['--duration']) if arguments['--duration'] else None)
    t_end = time.time()
    print("segmented audio collection to observations completed in: {}s"
          "".format(t_end - t0))
