"""
These set of functions help the algorithms of MSAF to read and write files
of the Segmentation Dataset.
"""
from collections import Counter
import datetime
import glob
import json
import logging
import numpy as np
import os
from threading import Thread

# Local stuff
import msaf
from msaf import jams2
from msaf import utils


class FileStruct:
    def __init__(self, audio_file):
        """Creates the entire file structure given the audio file."""
        self.ds_path = os.path.dirname(os.path.dirname(audio_file))
        self.audio_file = audio_file
        self.est_file = self._get_dataset_file(msaf.Dataset.estimations_dir,
                                               msaf.Dataset.estimations_ext)
        self.features_file = self._get_dataset_file(msaf.Dataset.features_dir,
                                                    msaf.Dataset.features_ext)
        self.ref_file = self._get_dataset_file(msaf.Dataset.references_dir,
                                               msaf.Dataset.references_ext)

    def _get_dataset_file(self, dir, ext):
        """Gets the desired dataset file."""
        audio_file_ext = "." + self.audio_file.split(".")[-1]
        base_file = os.path.basename(self.audio_file).replace(
            audio_file_ext, ext)
        return os.path.join(self.ds_path, dir, base_file)

    def __repr__(self):
        """Prints the file structure."""
        return "FileStruct(\n\tds_path=%s,\n\taudio_file=%s,\n\test_file=%s," \
            "\n\tfeatures_file=%s,\n\tref_file=%s\n)" % (
                self.ds_path, self.audio_file, self.est_file,
                self.features_file, self.ref_file)


def has_same_parameters(est_params, boundaries_id, labels_id, params):
    """Checks whether the parameters in params are the same as the estimated
    parameters in est_params."""
    K = 0
    for param_key in params.keys():
        if param_key in est_params.keys() and \
                est_params[param_key] == params[param_key] and \
                est_params["boundaries_id"] == boundaries_id and \
                ((labels_id is None and est_params["labels_id"] is None)
                or (est_params["labels_id"] == labels_id)):
            K += 1
    return K == len(params.keys())


def find_estimation(all_estimations, boundaries_id, labels_id, params,
                    est_file):
    """Finds the correct estimation from all the estimations contained in a
    JAMS file given the specified arguments.

    Parameters
    ----------
    all_estimations : list
        List of section Range Annotations from a JAMS file.
    boundaries_id : str
        Identifier of the algorithm used to compute the boundaries.
    labels_id : str
        Identifier of the algorithm used to compute the labels.
    params : dict
        Additional search parameters. E.g. {"feature" : "hpcp"}.
    est_file : str
        Path to the estimated file (JAMS file).

    Returns
    -------
    correct_est : RangeAnnotation
        Correct estimation found in all the estimations.
        None if it couldn't be found.
    corect_i : int
        Index of the estimation in the all_estimations list.
    """
    correct_est = None
    correct_i = -1
    found = False
    for i, estimation in enumerate(all_estimations):
        est_params = estimation.sandbox
        if has_same_parameters(est_params, boundaries_id, labels_id,
                               params) and not found:
            correct_est = estimation
            correct_i = i
            found = True
        elif has_same_parameters(est_params, boundaries_id, labels_id,
                                 params) and found:
            logging.warning("Multiple estimations match your parameters in "
                            "file %s" % est_file)
            correct_est = estimation
            correct_i = i
    return correct_est, correct_i


def read_estimations(est_file, boundaries_id, labels_id=None, **params):
    """Reads the estimations (boundaries and/or labels) from a jams file
    containing the estimations of an algorithm.

    Parameters
    ----------
    est_file : str
        Path to the estimated file (JAMS file).
    boundaries_id : str
        Identifier of the algorithm used to compute the boundaries.
    labels_id : str
        Identifier of the algorithm used to compute the labels.
    params : dict
        Additional search parameters. E.g. {"feature" : "hpcp"}.

    Returns
    -------
    boundaries : np.array((N,2))
        Array containing the estimated boundaries in intervals.
    labels : np.array(N)
        Array containing the estimated labels.
        Empty array if labels_id is None.
    """

    # Open file and read jams
    try:
        jam = jams2.load(est_file)
    except:
        logging.error("Could not open JAMS file %s" % est_file)
        return np.array([]), np.array([])

    # Get all the estimations for the sections
    all_estimations = jam.sections

    # Find correct estimation
    correct_est, i = find_estimation(all_estimations, boundaries_id, labels_id,
                                     params, est_file)
    if correct_est is None:
        logging.error("Could not find estimation in %s" % est_file)
        return np.array([]), np.array([])

    # Retrieve unique levels of segmentation
    levels = []
    for range in correct_est.data:
        levels.append(range.label.context)
    levels = list(set(levels))

    # Retrieve data
    all_boundaries = []
    all_labels = []
    for level in levels:
        boundaries = []
        labels = []
        for range in correct_est.data:
            if level == range.label.context:
                boundaries.append([range.start.value, range.end.value])
                if labels_id is not None:
                    labels.append(range.label.value)
        all_boundaries.append(np.asarray(boundaries))
        all_labels.append(np.asarray(labels, dtype=int))

    # If there is only one level, return np.arrays instead of lists
    if len(levels) == 1:
        all_boundaries = all_boundaries[0]
        all_labels = all_labels[0]

    return all_boundaries, all_labels


def get_algo_ids(est_file):
    """Gets the algorithm ids that are contained in the est_file."""
    with open(est_file, "r") as f:
        est_data = json.load(f)
        algo_ids = est_data["boundaries"].keys()
    return algo_ids


def read_references(audio_path, annotator_id=0):
    """Reads the boundary times and the labels.

    Parameters
    ----------
    audio_path : str
        Path to the audio file

    Returns
    -------
    ref_times : list
        List of boundary times
    ref_labels : list
        List of labels
    """
    # Dataset path
    ds_path = os.path.dirname(os.path.dirname(audio_path))

    # Read references
    jam_path = os.path.join(ds_path, msaf.Dataset.references_dir,
                            os.path.basename(audio_path)[:-4] +
                            msaf.Dataset.references_ext)
    ds_prefix = os.path.basename(audio_path).split("_")[0]

    # Get context
    if ds_prefix in msaf.prefix_dict.keys():
        context = msaf.prefix_dict[ds_prefix]
    else:
        context = "function"

    try:
        ref_inters, ref_labels = jams2.converters.load_jams_range(
            jam_path, "sections", context=context, annotator=annotator_id)
    except:
        logging.warning("Reference not found in %s" % jam_path)
        return []

    # Intervals to times
    ref_times = utils.intervals_to_times(ref_inters)

    return ref_times, ref_labels


def read_ref_labels(audio_path):
    """Reads the annotated labels from the given audio path."""
    ref_times, ref_labels = read_references(audio_path)
    return ref_labels


def read_ref_int_labels(audio_path):
    """Reads the annotated labels using unique integers as identifiers
    instead of strings."""
    ref_labels = read_ref_labels(audio_path)
    labels = []
    label_dict = {}
    k = 1
    for ref_label in ref_labels:
        if ref_label in label_dict.keys():
            labels.append(label_dict[ref_label])
        else:
            label_dict[ref_label] = k
            labels.append(k)
            k += 1
    return labels


def align_times(times, frames):
    """Aligns the times to the closes frame times (e.g. beats)."""
    dist = np.minimum.outer(times, frames)
    bound_frames = np.argmax(np.maximum(0, dist), axis=1)
    return np.unique(bound_frames)


def read_ref_bound_frames(audio_path, beats):
    """Reads the corresponding references file to retrieve the boundaries
        in frames."""
    ref_times, ref_labels = read_references(audio_path)
    # align with beats
    bound_frames = align_times(ref_times, beats)
    return bound_frames


def get_features(audio_path, annot_beats=False, framesync=False,
                 pre_features=None):
    """
    Gets the features of an audio file given the audio_path.

    Parameters
    ----------
    audio_path: str
        Path to the audio file.
    annot_beats: bool
        Whether to use annotated beats or not.
    framesync: bool
        Whether to use framesync features or not.
    pre_features: dict
        Pre computed features as a dictionary.
        `None` for reading them form the json file.

    Return
    ------
    C: np.array((N, 12))
        (Beat-sync) Chromagram
    M: np.array((N, 13))
        (Beat-sync) MFCC
    T: np.array((N, 6))
        (Beat-sync) Tonnetz
    cqt: np.array((N, msaf.Anal.cqt_bins))
        (Beat-sync) Constant-Q transform
    beats: np.array(T)
        Beats in seconds
    dur: float
        Song duration
    analysis : dict
        Parameters of analysis of track (e.g. sampling rate)
    """
    if pre_features is None:
        # Dataset path
        ds_path = os.path.dirname(os.path.dirname(audio_path))

        # Read Estimations
        features_path = os.path.join(ds_path, msaf.Dataset.features_dir,
            os.path.basename(audio_path)[:-4] + msaf.Dataset.features_ext)
        with open(features_path, "r") as f:
            feats = json.load(f)

        # Beat Synchronous Feats
        if framesync:
            feat_str = "framesync"
            beats = None
        else:
            if annot_beats:
                # Read references
                try:
                    annotation_path = os.path.join(
                        ds_path, msaf.Dataset.references_dir,
                        os.path.basename(audio_path)[:-4] +
                        msaf.Dataset.references_ext)
                    jam = jams2.load(annotation_path)
                except:
                    raise RuntimeError("No references found in file %s" %
                                    annotation_path)

                feat_str = "ann_beatsync"
                beats = []
                beat_data = jam.beats[0].data
                if beat_data == []:
                    raise ValueError
                for data in beat_data:
                    beats.append(data.time.value)
                beats = np.unique(beats)
            else:
                feat_str = "est_beatsync"
                beats = np.asarray(feats["beats"]["times"])
        C = np.asarray(feats[feat_str]["hpcp"])
        M = np.asarray(feats[feat_str]["mfcc"])
        T = np.asarray(feats[feat_str]["tonnetz"])
        cqt = np.asarray(feats[feat_str]["cqt"])
        analysis = feats["analysis"]
        dur = analysis["dur"]

        # Frame times might be shorter than the actual number of features.
        if framesync:
            frame_times = utils.get_time_frames(dur, analysis)
            C = C[:len(frame_times)]
            M = M[:len(frame_times)]
            T = T[:len(frame_times)]