"""
This module extracts out a bunch of the Auditory Scene Analysis (ASA)
logic, which has grown to be a little unwieldy in the AudioSegment class.
"""
import multiprocessing
import numpy as np
import scipy.signal as signal
import sys

def _plot(frequencies, spect, title, fn=None):
    import matplotlib.pyplot as plt
    for freq, (index, row) in zip(frequencies[::-1], enumerate(spect[::-1, :])):
        plt.subplot(spect.shape[0], 1, index + 1)
        if index == 0:
            plt.title(title)
        plt.ylabel("{0:.0f}".format(freq))
        plt.plot(row)
        if fn is not None:
            fn(plt, freq, index, row)

def visualize_time_domain(seg, title=""):
    import matplotlib.pyplot as plt
    plt.plot(seg)
    plt.title(title)
    plt.show()
    plt.clf()

def visualize(spect, frequencies, title=""):
    import matplotlib.pyplot as plt
    _plot(frequencies, spect, title)
    plt.show()
    plt.clf()

def visualize_peaks_and_valleys(peaks, valleys, spect, frequencies):
    import matplotlib.pyplot as plt
    # Reverse everything to make it have the low fs at the bottom of the figure
    peaks = peaks[::-1, :]
    valleys = valleys[::-1, :]
    def _plot_pandv(p, _freq, index, row):
        these_peaks = peaks[index]
        peak_values = these_peaks * row  # Mask off anything that isn't a peak
        these_valleys = valleys[index]
        valley_values = these_valleys * row  # Mask off anything that isn't a valley
        p.plot(peak_values, 'ro')
        p.plot(valley_values, 'bo')
    _plot(frequencies, spect, "Peaks (red) and Valleys (blue)", _plot_pandv)
    plt.show()
    plt.clf()

def visualize_fronts(onsets, offsets, spect, frequencies):
    import matplotlib.pyplot as plt
    # Reverse everything to make it have the low fs at the bottom of the figure
    onsets = onsets[::-1, :]
    offsets = offsets[::-1, :]
    def _plot_fronts(p, _freq, index, row):
        # Cycle through all the different onsets and offsets and plot them each
        colors = ('b', 'g', 'r', 'c', 'm', 'y', 'k')
        nonzero_indexes_onsets = np.reshape(np.where(onsets[index, :] != 0), (-1,))
        for x in nonzero_indexes_onsets:
            id = int(onsets[index][x])
            p.axvline(x=x, color=colors[id % len(colors)], linestyle='--')
        nonzero_indexes_offsets = np.reshape(np.where(offsets[index, :] != 0), (-1,))
        for x in nonzero_indexes_offsets:
            id = int(offsets[index][x])
            p.axvline(x=x, color=colors[id % len(colors)], linestyle='-')
    _plot(frequencies, spect, "Fronts", _plot_fronts)
    plt.show()
    plt.clf()

def visualize_segmentation_mask(segmentation, spect, frequencies, mode='new'):
    import matplotlib.pyplot as plt
    if mode.lower() == 'old':
        # Reverse segmentation mask's frequency dimension so that low fs is at the bottome
        segmentation = np.copy(segmentation)
        segmentation = segmentation[::-1, :]
        def _plot_seg(p, freq, index, row):
            colors = ('b', 'g', 'r', 'c', 'm', 'y', 'k')
            segment_ids_in_this_frequency = np.reshape(np.where(segmentation[index, :] != 0), (-1,))
            for x in segment_ids_in_this_frequency:
                id = int(segmentation[index][x])
                plot_a_line = False
                if x == 0:  # if this is the first item in the frequency, then it must be a start of the mask
                    plot_a_line = True
                elif int(segmentation[index][x - 1]) != id:  # this is the start of a mask in this frequency
                    plot_a_line = True
                elif x == len(row) - 1:  # this is the very last sample, it must be the end of a mask
                    plot_a_line = True
                elif int(segmentation[index][x + 1]) != id:  # this is the last sample of the mask in this frequency
                    plot_a_line = True

                if plot_a_line:
                    p.axvline(x=x, color=colors[id % len(colors)], linestyle="--")
        _plot(frequencies, spect, "Segmentation Mask", _plot_seg)
    elif mode.lower() == 'new':
        times = np.arange(spect.shape[1])
        spect = spect * 10000
        plt.pcolormesh(times, frequencies, segmentation)
        plt.show()
    else:
        raise ValueError("Mode must be one of ('new', 'old')")
    plt.show()
    plt.clf()

def _compute_peaks_or_valleys_of_first_derivative(s, do_peaks=True):
    """
    Takes a spectrogram and returns a 2D array of the form:

    0 0 0 1 0 0 1 0 0 0 1   <-- Frequency 0
    0 0 1 0 0 0 0 0 0 1 0   <-- Frequency 1
    0 0 0 0 0 0 1 0 1 0 0   <-- Frequency 2
    *** Time axis *******

    Where a 1 means that the value in that time bin in the spectrogram corresponds to
    a peak/valley in the first derivative.

    This function is used as part of the ASA algorithm and is not meant to be used publicly.
    """
    # Get the first derivative of each frequency in the time domain
    gradient = np.nan_to_num(np.apply_along_axis(np.gradient, 1, s), copy=False)

    # Calculate the value we will use for determinig whether something is an event or not
    threshold = np.squeeze(np.nanmean(gradient, axis=1) + np.nanstd(gradient, axis=1))

    # Look for relative extrema along the time dimension
    half_window = 4
    if do_peaks:
        indexes = [signal.argrelextrema(gradient[i, :], np.greater, order=half_window)[0] for i in range(gradient.shape[0])]
    else:
        indexes = [signal.argrelextrema(gradient[i, :], np.less, order=half_window)[0] for i in range(gradient.shape[0])]

    # indexes should now contain the indexes of possible extrema
    # But we need to filter out values that are not large enough, and we want the end result
    # to be a 1 or 0 mask corresponding to locations of extrema
    extrema = np.zeros(s.shape)
    for row_index, index_array in enumerate(indexes):
        # Each index_array is a list of indexes corresponding to all the extrema in a given row
        for col_index in index_array:
            if do_peaks and (gradient[row_index, col_index] > threshold[row_index]):
                extrema[row_index, col_index] = 1
            elif not do_peaks:
                # Note that we do not remove under-threshold values from the offsets - these will be taken care of later in the algo
                extrema[row_index, col_index] = 1
    return extrema, gradient

def _correlate_onsets_and_offsets(onsets, offsets, gradients):
    """
    Takes an array of onsets and an array of offsets, of the shape [nfrequencies, nsamples], where
    each item in these arrays is either a 0 (not an on/offset) or a 1 (a possible on/offset).

    This function returns a new offsets array, where there is a one-to-one correlation between
    onsets and offsets, such that each onset has exactly one offset that occurs after it in
    the time domain (the second dimension of the array).

    The gradients array is used to decide which offset to use in the case of multiple possibilities.
    """
    # For each freq channel:
    for freq_index, (ons, offs) in enumerate(zip(onsets[:, :], offsets[:, :])):
        # Scan along onsets[f, :] until we find the first 1
        indexes_of_all_ones = np.reshape(np.where(ons == 1), (-1,))

        # Zero out anything in the offsets up to (and including) this point
        # since we can't have an offset before the first onset
        last_idx = indexes_of_all_ones[0]
        offs[0:last_idx + 1] = 0

        if len(indexes_of_all_ones > 1):
            # Do the rest of this only if we have more than one onset in this frequency band
            for next_idx in indexes_of_all_ones[1:]:
                # Get all the indexes of possible offsets from onset index to next onset index
                offset_choices = offs[last_idx:next_idx]
                offset_choice_indexes = np.where(offset_choices == 1)

                # Assert that there is at least one offset choice
                if not np.any(offset_choices):
                    continue
                assert np.any(offset_choices), "Offsets from {} to {} only include zeros".format(last_idx, next_idx)

                # If we have more than one choice, the offset index is the one that corresponds to the most negative gradient value
                # Convert the offset_choice_indexes to indexes in the whole offset array, rather than just the offset_choices array
                offset_choice_indexes = np.reshape(last_idx + offset_choice_indexes, (-1,))
                assert np.all(offsets[freq_index, offset_choice_indexes])
                gradient_values = gradients[freq_index, offset_choice_indexes]
                index_of_largest_from_gradient_values = np.where(gradient_values == np.min(gradient_values))[0]
                index_of_largest_offset_choice = offset_choice_indexes[index_of_largest_from_gradient_values]
                assert offsets[freq_index, index_of_largest_offset_choice] == 1

                # Zero the others
                offsets[freq_index, offset_choice_indexes] = 0
                offsets[freq_index, index_of_largest_offset_choice] = 1
                last_idx = next_idx
        else:
            # We only have one onset in this frequency band, so the offset will be the very last sample
            offsets[freq_index, :] = 0
            offsets[freq_index, -1] = 1
    return offsets

def _form_onset_offset_fronts(ons_or_offs, sample_rate_hz, threshold_ms=20):
    """
    Takes an array of onsets or offsets (shape = [nfrequencies, nsamples], where a 1 corresponds to an on/offset,
    and samples are 0 otherwise), and returns a new array of the same shape, where each 1 has been replaced by
    either a 0, if the on/offset has been discarded, or a non-zero positive integer, such that
    each front within the array has a unique ID - for example, all 2s in the array will be the front for on/offset
    front 2, and all the 15s will be the front for on/offset front 15, etc.

    Due to implementation details, there will be no 1 IDs.
    """
    threshold_s = threshold_ms / 1000
    threshold_samples = sample_rate_hz * threshold_s

    ons_or_offs = np.copy(ons_or_offs)

    claimed = []
    this_id = 2
    # For each frequency,
    for frequency_index, row in enumerate(ons_or_offs[:, :]):
        ones = np.reshape(np.where(row == 1), (-1,))

        # for each 1 in that frequency,
        for top_level_frequency_one_index in ones:
            claimed.append((frequency_index, top_level_frequency_one_index))

            found_a_front = False
            # for each frequencies[i:],
            for other_frequency_index, other_row in enumerate(ons_or_offs[frequency_index + 1:, :], start=frequency_index + 1):

                # for each non-claimed 1 which is less than theshold_ms away in time,
                upper_limit_index = top_level_frequency_one_index + threshold_samples
                lower_limit_index = top_level_frequency_one_index - threshold_samples
                other_ones = np.reshape(np.where(other_row == 1), (-1,))  # Get the indexes of all the 1s in row
                tmp = np.reshape(np.where((other_ones >= lower_limit_index)  # Get the indexes in the other_ones array of all items in bounds
                                        & (other_ones <= upper_limit_index)), (-1,))
                other_ones = other_ones[tmp]  # Get the indexes of all the 1s in the row that are in bounds
                if len(other_ones) > 0:
                    unclaimed_idx = other_ones[0]  # Take the first one
                    claimed.append((other_frequency_index, unclaimed_idx))
                elif len(claimed) < 3:
                    # revert the top-most 1 to 0
                    ons_or_offs[frequency_index, top_level_frequency_one_index] = 0
                    claimed = []
                    break  # Break from the for-each-frequencies[i:] loop so we can move on to the next item in the top-most freq
                elif len(claimed) >= 3:
                    found_a_front = True
                    # this group of so-far-claimed forms a front
                    claimed_as_indexes = tuple(np.array(claimed).T)
                    ons_or_offs[claimed_as_indexes] = this_id
                    this_id += 1
                    claimed = []
                    break  # Move on to the next item in the top-most array
            # If we never found a frequency that did not have a matching offset, handle that case here
            if len(claimed) >= 3:
                claimed_as_indexes = tuple(np.array(claimed).T)
                ons_or_offs[claimed_as_indexes] = this_id
                this_id += 1
                claimed = []
            elif found_a_front:
                this_id += 1
            else:
                ons_or_offs[frequency_index, top_level_frequency_one_index] = 0
                claimed = []

    return ons_or_offs

def _lookup_offset_by_onset_idx(onset_idx, onsets, offsets):
    """
    Takes an onset index (freq, sample) and returns the offset index (freq, sample)
    such that frequency index is the same, and sample index is the minimum of all
    offsets ocurring after the given onset. If there are no offsets after the given
    onset in that frequency channel, the final sample in that channel is returned.
    """
    assert len(onset_idx) == 2, "Onset_idx must be a tuple of the form (freq_idx, sample_idx)"
    frequency_idx, sample_idx = onset_idx
    offset_sample_idxs = np.reshape(np.where(offsets[frequency_idx, :] == 1), (-1,))
    # get the offsets which occur after onset
    offset_sample_idxs = offset_sample_idxs[offset_sample_idxs > sample_idx]
    if len(offset_sample_idxs) == 0:
        # There is no offset in this frequency that occurs after the onset, just return the last sample
        chosen_offset_sample_idx = offsets.shape[1] - 1
        assert offsets[frequency_idx, chosen_offset_sample_idx] == 0
    else:
        # Return the closest offset to the onset
        chosen_offset_sample_idx = offset_sample_idxs[0]
        assert offsets[frequency_idx, chosen_offset_sample_idx] != 0
    return frequency_idx, chosen_offset_sample_idx

def _get_front_idxs_from_id(fronts, id):
    """
    Return a list of tuples of the form (frequency_idx, sample_idx),
    corresponding to all the indexes of the given front.
    """
    if id == -1:
        # This is the only special case.
        # -1 is the index of the catch-all final column offset front.
        freq_idxs = np.arange(fronts.shape[0], dtype=np.int64)
        sample_idxs = np.ones(len(freq_idxs), dtype=np.int64) * (fronts.shape[1] - 1)
    else:
        freq_idxs, sample_idxs = np.where(fronts == id)
    return [(f, i) for f, i in zip(freq_idxs, sample_idxs)]

def _choose_front_id_from_candidates(candidate_offset_front_ids, offset_fronts, offsets_corresponding_to_onsets):
    """
    Returns a front ID which is the id of the offset front that contains the most overlap
    with offsets that correspond to the given onset front ID.
    """
    noverlaps = []  # will contain tuples of the form (number_overlapping, offset_front_id)
    for offset_front_id in candidate_offset_front_ids:
        offset_front_f_idxs, offset_front_s_idxs = np.where(offset_fronts == offset_front_id)
        offset_front_idxs = [(f, i) for f, i in zip(offset_front_f_idxs, offset_front_s_idxs)]
        noverlap_this_id = len(set(offset_front_idxs).symmetric_difference(set(offsets_corresponding_to_onsets)))
        noverlaps.append((noverlap_this_id, offset_front_id))
    _overlapped, chosen_offset_front_id = max(noverlaps, key=lambda t: t[0])
    return int(chosen_offset_front_id)

def _get_offset_front_id_after_onset_sample_idx(onset_sample_idx, offset_fronts):
    """
    Returns the offset_front_id which corresponds to the offset front which occurs
    first entirely after the given onset sample_idx.
    """
    # get all the offset_front_ids
    offset_front_ids = [i for i in np.unique(offset_fronts) if i != 0]

    best_id_so_far = -1
    closest_offset_sample_idx = sys.maxsize
    for offset_front_id in offset_front_ids:
        # get all that offset front's indexes
        offset_front_idxs = _get_front_idxs_from_id(offset_fronts, offset_front_id)

        # get the sample indexes
        offset_front_sample_idxs = [s for _f, s in offset_front_idxs]

        # if each sample index is greater than onset_sample_idx, keep this offset front if it is the best one so far
        min_sample_idx = min(offset_front_sample_idxs)
        if min_sample_idx > onset_sample_idx and min_sample_idx < closest_offset_sample_idx:
            closest_offset_sample_idx = min_sample_idx
            best_id_so_far = offset_front_id

    assert best_id_so_far > 1 or best_id_so_far == -1
    return best_id_so_far

def _get_offset_front_id_after_onset_front(onset_front_id, onset_fronts, offset_fronts):
    """
    Get the ID corresponding to the offset which occurs first after the given onset_front_id.
    By `first` I mean the front which contains the offset which is closest to the latest point
    in the onset front. By `after`, I mean that the offset must contain only offsets which
    occur after the latest onset in the onset front.

    If there is no appropriate offset front, the id returned is -1.
    """
    # get the onset idxs for this front
    onset_idxs = _get_front_idxs_from_id(onset_fronts, onset_front_id)

    # get the sample idxs for this front
    onset_sample_idxs = [s for _f, s in onset_idxs]

    # get the latest onset in this onset front
    latest_onset_in_front = max(onset_sample_idxs)

    offset_front_id_after_this_onset_front = _get_offset_front_id_after_onset_sample_idx(latest_onset_in_front, offset_fronts)

    return int(offset_front_id_after_this_onset_front)

def _match_offset_front_id_to_onset_front_id(onset_front_id, onset_fronts, offset_fronts, onsets, offsets):
    """
    Find all offset fronts which are composed of at least one offset which corresponds to one of the onsets in the
    given onset front.
    The offset front which contains the most of such offsets is the match.
    If there are no such offset fronts, return -1.
    """
    # find all offset fronts which are composed of at least one offset which corresponds to one of the onsets in the onset front
    # the offset front which contains the most of such offsets is the match

    # get the onsets that make up front_id
    onset_idxs = _get_front_idxs_from_id(onset_fronts, onset_front_id)

    # get the offsets that match the onsets in front_id
    offset_idxs = [_lookup_offset_by_onset_idx(i, onsets, offsets) for i in onset_idxs]

    # get all offset_fronts which contain at least one of these offsets
    candidate_offset_front_ids = set([int(offset_fronts[f, i]) for f, i in offset_idxs])

    # It is possible that offset_idxs contains offset indexes that correspond to offsets that did not
    # get formed into a front - those will have a front ID of 0. Remove them.
    candidate_offset_front_ids = [id for id in candidate_offset_front_ids if id != 0]

    if candidate_offset_front_ids:
        chosen_offset_front_id = _choose_front_id_from_candidates(candidate_offset_front_ids, offset_fronts, offset_idxs)
    else:
        chosen_offset_front_id = _get_offset_front_id_after_onset_front(onset_front_id, onset_fronts, offset_fronts)

    return chosen_offset_front_id

def _get_consecutive_portions_of_front(front):
    """
    Yields lists of the form [(f, s), (f, s)], one at a time from the given front (which is a list of the same form),
    such that each list yielded is consecutive in frequency.
    """
    last_f = None
    ls = []
    for f, s in front:
        if last_f is not None and f != last_f + 1:
            yield ls
            ls = []
        ls.append((f, s))
        last_f = f
    yield ls

def _get_consecutive_and_overlapping_fronts(onset_fronts, offset_fronts, onset_front_id, offset_front_id):
    """
    Gets an onset_front and an offset_front such that they both occupy at least some of the same
    frequency channels, then returns the portion of each that overlaps with the other.
    """
    # Get the onset front of interest
    onset_front = _get_front_idxs_from_id(onset_fronts, onset_front_id)

    # Get the offset front of interest
    offset_front = _get_front_idxs_from_id(offset_fronts, offset_front_id)

    # Keep trying consecutive portions of this onset front until we find a consecutive portion
    # that overlaps with part of the offset front
    consecutive_portions_of_onset_front = [c for c in _get_consecutive_portions_of_front(onset_front)]
    for consecutive_portion_of_onset_front in consecutive_portions_of_onset_front:
        # Only get the segment of this front that overlaps in frequencies with the onset front of interest
        onset_front_frequency_indexes = [f for f, _ in consecutive_portion_of_onset_front]
        overlapping_offset_front = [(f, s) for f, s in offset_front if f in onset_front_frequency_indexes]

        # Only get as much of this overlapping portion as is actually consecutive
        for consecutive_portion_of_offset_front in _get_consecutive_portions_of_front(overlapping_offset_front):
            if consecutive_portion_of_offset_front:
                # Just return the first one we get - if we get any it means we found a portion of overlap
                return consecutive_portion_of_onset_front, consecutive_portion_of_offset_front
    return [], []  # These two fronts have no overlap

def _update_segmentation_mask(segmentation_mask, onset_fronts, offset_fronts, onset_front_id, offset_front_id_most_overlap):
    """
    Returns an updated segmentation mask such that the input `segmentation_mask` has been updated by segmenting between
    `onset_front_id` and `offset_front_id`, as found in `onset_fronts` and `offset_fronts`, respectively.

    This function also returns the onset_fronts and offset_fronts matrices, updated so that any fronts that are of
    less than 3 channels wide are removed.

    This function also returns a boolean value indicating whether the onset channel went to completion.

    Specifically, segments by doing the following:

    - Going across frequencies in the onset_front,
    - add the segment mask ID (the onset front ID) to all samples between the onset_front and the offset_front,
      if the offset_front is in that frequency.

    Possible scenarios:

    Fronts line up completely:

    ::

        |   |       S S S
        |   |  =>   S S S
        |   |       S S S
        |   |       S S S

    Onset front starts before offset front:

    ::

        |           |
        |   |       S S S
        |   |  =>   S S S
        |   |       S S S

    Onset front ends after offset front:

    ::

        |   |       S S S
        |   |  =>   S S S
        |   |       S S S
        |           |

    Onset front starts before and ends after offset front:

    ::

        |           |
        |   |  =>   S S S
        |   |       S S S
        |           |

    The above three options in reverse:

    ::

            |       |S S|           |
        |S S|       |S S|       |S S|
        |S S|       |S S|       |S S|
        |S S|           |           |

    There is one last scenario:

    ::

        |   |
        \   /
         \ /
         / \
        |   |

    Where the offset and onset fronts cross one another. If this happens, we simply
    reverse the indices and accept:

    ::

        |sss|
        \sss/
         \s/
         /s\
        |sss|

    The other option would be to destroy the offset front from the crossover point on, and
    then search for a new offset front for the rest of the onset front.
    """
    # Get the portions of the onset and offset fronts that overlap and are consecutive
    onset_front_overlap, offset_front_overlap = _get_consecutive_and_overlapping_fronts(onset_fronts, offset_fronts, onset_front_id, offset_front_id_most_overlap)
    onset_front = _get_front_idxs_from_id(onset_fronts, onset_front_id)
    offset_front = _get_front_idxs_from_id(offset_fronts, offset_front_id_most_overlap)
    msg = "Onset front {} and offset front {} result in consecutive overlapping portions of (on) {} and (off) {}, one of which is empty".format(
        onset_front, offset_front, onset_front_overlap, offset_front_overlap
    )
    assert onset_front_overlap, msg
    assert offset_front_overlap, msg
    onset_front = onset_front_overlap
    offset_front = offset_front_overlap

    # Figure out which frequencies will go in the segment
    flow_on, _slow_on = onset_front[0]
    fhigh_on, _shigh_on = onset_front[-1]
    flow_off, _slow_off = offset_front[0]
    fhigh_off, _shigh_off = offset_front[-1]
    flow = max(flow_on, flow_off)
    fhigh = min(fhigh_on, fhigh_off)

    # Update all the masks with the segment
    for fidx, _freqchan in enumerate(segmentation_mask[flow:fhigh + 1, :], start=flow):
        assert fidx >= flow, "Frequency index is {}, but we should have started at {}".format(fidx, flow)
        assert (fidx - flow) < len(onset_front), "Frequency index {} minus starting frequency {} is too large for nfrequencies {} in onset front {}".format(
            fidx, flow, len(onset_front), onset_front
        )
        assert (fidx - flow) < len(offset_front), "Frequency index {} minus starting frequency {} is too large for nfrequencies {} in offset front {}".format(
            fidx, flow, len(offset_front), offset_front
        )
        _, beg = onset_front[fidx - flow]
        _, end = offset_front[fidx - flow]
        if beg > end:
            end, beg = beg, end
        assert end >= beg
        segmentation_mask[fidx, beg:end + 1] = onset_front_id
        onset_fronts[fidx, (beg + 1):(end + 1)] = 0
        offset_fronts[fidx, (beg + 1):(end + 1)] = 0
    nfreqs_used_in_onset_front = (fidx - flow) + 1

    # Update the other masks to delete fronts that have been used
    indexes = np.arange(flow, fhigh + 1, 1, dtype=np.int64)
    onset_front_sample_idxs_across_freqs = np.array([s for _, s in onset_front])
    onset_front_sample_idxs_across_freqs_up_to_break = onset_front_sample_idxs_across_freqs[:nfreqs_used_in_onset_front]
    offset_front_sample_idxs_across_freqs = np.array([s for _, s in offset_front])
    offset_front_sample_idxs_across_freqs_up_to_break = offset_front_sample_idxs_across_freqs[:nfreqs_used_in_onset_front]

    ## Remove the offset front from where we started to where we ended
    offset_fronts[indexes[:nfreqs_used_in_onset_front], offset_front_sample_idxs_across_freqs_up_to_break] = 0

    ## Remove the onset front from where we started to where we ended
    onset_fronts[indexes[:nfreqs_used_in_onset_front], onset_front_sample_idxs_across_freqs_up_to_break] = 0

    # Determine if we matched the entire onset front by checking if there is any more of this onset front in onset_fronts
    whole_onset_front_matched = onset_front_id not in np.unique(onset_fronts)

    return whole_onset_front_matched

def _front_id_from_idx(front, index):
    """
    Returns the front ID found in `front` at the given `index`.

    :param front:               An onset or offset front array of shape [nfrequencies, nsamples]
    :index:                     A tuple of the form (frequency index, sample index)
    :returns:                   The ID of the front or -1 if not found in `front` and the item at `onsets_or_offsets[index]`
                                is not a 1.
    """
    fidx, sidx = index
    id = front[fidx, sidx]
    if id == 0:
        return -1
    else:
        return id

def _get_front_ids_one_at_a_time(onset_fronts):
    """
    Yields one onset front ID at a time until they are gone. All the onset fronts from a
    frequency channel are yielded, then all of the next channel's, etc., though one at a time.
    """
    yielded_so_far = set()
    for row in onset_fronts:
        for id in row:
            if id != 0 and id not in yielded_so_far:
                yield id
                yielded_so_far.add(id)

def _get_corresponding_offsets(onset_fronts, onset_front_id, onsets, offsets):
    """
    Gets the offsets that occur as close as possible to the onsets in the given onset-front.
    """
    corresponding_offsets = []
    for index in _get_front_idxs_from_id(onset_fronts, onset_front_id):
        offset_fidx, offset_sidx = _lookup_offset_by_onset_idx(index, onsets, offsets)
        corresponding_offsets.append((offset_fidx, offset_sidx))
    return corresponding_offsets

def _get_all_offset_fronts_from_offsets(offset_fronts, corresponding_offsets):
    """
    Returns all the offset fronts that are composed of at least one of the given offset indexes.
    Also returns a dict of the form {offset_front_id: ntimes saw}
    """
    all_offset_fronts_of_interest = []
    ids_ntimes_seen = {}
    for offset_index in corresponding_offsets:
        offset_id = _front_id_from_idx(offset_fronts, offset_index)
        if offset_id not in ids_ntimes_seen:
            offset_front_idxs = _get_front_idxs_from_id(offset_fronts, offset_id)
            all_offset_fronts_of_interest.append(offset_front_idxs)
            ids_ntimes_seen[offset_id] = 1
        else:
            ids_ntimes_seen[offset_id] += 1
    return all_offset_fronts_of_interest, ids_ntimes_seen

def _remove_overlaps(segmentation_mask, fronts):
    """
    Removes all points in the fronts that overlap with the segmentation mask.
    """
    fidxs, sidxs = np.where((segmentation_mask != fronts) & (segmentation_mask != 0) & (fronts != 0))
    fronts[fidxs, sidxs] = 0

def _match_fronts(onset_fronts, offset_fronts, onsets, offsets, debug=False):
    """
    Returns a segmentation mask, which looks like this:
    frequency 1: 0 0 4 4 4 4 4 0 0 5 5 5
    frequency 2: 0 4 4 4 4 4 0 0 0 0 5 5
    frequency 3: 0 4 4 4 4 4 4 4 5 5 5 5

    That is, each item in the array is either a 0 (not part of a segment) or a positive
    integer which indicates which segment the sample in that frequency band belongs to.
    """
    def printd(*args, **kwargs):
        if debug:
            print(*args, **kwargs)

    # Make copies of everything, so we can do whatever we want with them
    onset_fronts = np.copy(onset_fronts)
    offset_fronts = np.copy(offset_fronts)
    onsets = np.copy(onsets)
    offsets = np.copy(offsets)

    # This is what we will return
    segmentation_mask = np.zeros_like(onset_fronts)

    # - Take the first frequency in the onset_fronts matrix
    #     [ s s s s s s s s s] <-- This frequency
    #     [ s s s s s s s s s]
    #     [ s s s s s s s s s]
    #     [ s s s s s s s s s]
    #     [ s s s s s s s s s]

    # - Follow it along in time like this:

    #     first sample    last sample
    #       v      -->      v
    #     [ s s s s s s s s s]
    #     [ s s s s s s s s s]
    #     [ s s s s s s s s s]
    #     [ s s s s s s s s s]
    #     [ s s s s s s s s s]

    # until you get to the first onset front in that frequency

    #     Here it is!
    #         v
    #     [ . O . . . . . . .]
    #     [ . . O . . . . . .]
    #     [ . . O . . . . . .]
    #     [ . O . . . . . . .]
    #     [ O . . . . . . . .]

    resulting_onset_fronts = np.copy(onset_fronts)
    printd("    -> Dealing with onset fronts...")
    for onset_front_id in _get_front_ids_one_at_a_time(onset_fronts):
        printd("      -> Dealing with onset front", int(onset_front_id))
        front_is_complete = False
        while not front_is_complete:
            # - Now, starting at this onset front in each frequency, find that onset's corresponding offset

            #     [ . O . . . . F . .]
            #     [ . . O . . . F . .]
            #     [ . . O . F . . . .]
            #     [ . O . F . . . . .]
            #     [ O F . . . . . . .]

            corresponding_offsets = _get_corresponding_offsets(resulting_onset_fronts, onset_front_id, onsets, offsets)

            # It is possible that onset_front_id has been removed from resulting_onset_fronts,
            # if so, skip it and move on to the next onset front (we are iterating over the original
            # to keep the iterator valid)
            if not corresponding_offsets:
                break

            # - Get all the offset fronts that are composed of at least one of these offset times

            #     [ . O . . . . 1 . .]
            #     [ . . O 3 . . 1 . .]
            #     [ . . O 3 F . 1 . .]
            #     [ . O . 3 . . . 1 .]
            #     [ O F 3 . . . . . .]

            _all_offset_fronts_of_interest, ids_ntimes_seen = _get_all_offset_fronts_from_offsets(offset_fronts, corresponding_offsets)

            # - Check how many of these offset times each of the offset fronts are composed of:

            #     [ . O . . . . Y . .]
            #     [ . . O 3 . . Y . .]
            #     [ . . O 3 F . 1 . .]
            #     [ . O . X . . . 1 .]
            #     [ O F 3 . . . . . .]

            # In this example, offset front 1 is made up of 4 offset times, 2 of which (the Y's) are offset times
            # that correspond to onsets in the onset front we are currently dealing with. Meanwhile, offset
            # front 3 is made up of 4 offset times, only one of which (the X) is one of the offsets that corresponds
            # to the onset front.

            # - Choose the offset front which matches the most offset time candidates. In this example, offset front 1
            #   is chosen because it has 2 of these offset times.
            #   If there is a tie, we choose the ID with the lower number
            ntimes_seen_sorted = sorted([(k, v) for k, v in ids_ntimes_seen.items()], key=lambda tup: (-1 * tup[1], tup[0]))
            assert len(ntimes_seen_sorted) > 0, "We somehow got an empty dict of offset front IDs"

            # Only use the special final front (the -1, catch-all front composed of final samples in each frequency) if necessary
            offset_front_id, _ntimes_seen = ntimes_seen_sorted[0]
            if offset_front_id == -1 and len(ntimes_seen_sorted) > 1:
                offset_front_id, _ntimes_seen = ntimes_seen_sorted[1]
            offset_front_id_most_overlap = offset_front_id

            # - Finally, update the segmentation mask to follow the offset
            #   front from where it first overlaps in frequency with the onset front to where it ends or to where
            #   the onset front ends, whichever happens first.

            #     [ . S S S S S S . .]
            #     [ . . S S S S S . .]
            #     [ . . S S S S S . .]
            #     [ . S S S S S S S .]
            #     [ O F 3 . . . . . .]  <-- This frequency has not yet been matched with an offset front
            front_is_complete = _update_segmentation_mask(segmentation_mask,
                                                            resulting_onset_fronts,
                                                            offset_fronts,
                                                            onset_front_id,
                                                            offset_front_id_most_overlap)

            # Remove any onsets that are covered by the new segmentation mask
            _remove_overlaps(segmentation_mask, resulting_onset_fronts)

            # Remove any offsets that are covered by the new segmentaion mask
            _remove_overlaps(segmentation_mask, offset_fronts)

            # - Repeat this algorithm, restarting in the first frequency channel that did not match (the last frequency in
            #   the above example). Do this until you have finished with this onset front.

        # - Repeat for each onset front in the rest of this frequency
        # - Repeat for each frequency

    return segmentation_mask

def _remove_fronts_that_are_too_small(fronts, size):
    """
    Removes all fronts from `fronts` which are strictly smaller than
    `size` consecutive frequencies in length.
    """
    ids = np.unique(fronts)
    for id in ids:
        if id == 0 or id == -1:
            continue
        front = _get_front_idxs_from_id(fronts, id)
        if len(front) < size:
            indexes = ([f for f, _ in front], [s for _, s in front])
            fronts[indexes] = 0

def _break_poorly_matched_fronts(fronts, threshold=0.1, threshold_overlap_samples=3):
    """
    For each onset front, for each frequency in that front, break the onset front if the signals
    between this frequency's onset and the next frequency's onset are not similar enough.

    Specifically:
    If we have the following two frequency channels, and the two O's are part of the same onset front,

    ::

        [ . O . . . . . . . . . . ]
        [ . . . . O . . . . . . . ]

    We compare the signals x and y:

    ::

        [ . x x x x . . . . . . . ]
        [ . y y y y . . . . . . . ]

    And if they are not sufficiently similar (via a DSP correlation algorithm), we break the onset
    front between these two channels.

    Once this is done, remove any onset fronts that are less than 3 channels wide.
    """
    assert threshold_overlap_samples > 0, "Number of samples of overlap must be greater than zero"
    breaks_after = {}
    for front_id in _get_front_ids_one_at_a_time(fronts):
        front = _get_front_idxs_from_id(fronts, front_id)
        for i, (f, s) in enumerate(front):
            if i < len(front) - 1:
                # Get the signal from f, s to f, s+1 and the signal from f+1, s to f+1, s+1
                next_f, next_s = front[i + 1]
                low_s = min(s, next_s)
                high_s = max(s, next_s)
                sig_this_f = fronts[f, low_s:high_s]
                sig_next_f = fronts[next_f, low_s:high_s]
                assert len(sig_next_f) == len(sig_this_f)

                if len(sig_next_f) > threshold_overlap_samples:
                    # If these two signals are not sufficiently close in form, this front should be broken up
                    correlation = signal.correlate(sig_this_f, sig_next_f, mode='same')
                    assert len(correlation) > 0
                    correlation = correlation / max(correlation + 1E-9)
                    similarity = np.sum(correlation) / len(correlation)
                    # TODO: the above stuff probably needs to be figured out
                    if similarity < threshold:
                        if front_id in breaks_after:
                            breaks_after[front_id].append((f, s))
                        else:
                            breaks_after[front_id] = []

    # Now update the fronts matrix by breaking up any fronts at the points we just identified
    # and assign the newly created fronts new IDs
    taken_ids = sorted(np.unique(fronts))
    next_id = taken_ids[-1] + 1
    for id in breaks_after.keys():
        for f, s in breaks_after[id]:
            fidxs, sidxs = np.where(fronts == id)
            idxs_greater_than_f = [fidx for fidx in fidxs if fidx > f]
            start = len(sidxs) - len(idxs_greater_than_f)
            indexes = (idxs_greater_than_f, sidxs[start:])
            fronts[indexes] = next_id
            next_id += 1

    _remove_fronts_that_are_too_small(fronts, 3)

def _update_segmentation_mask_if_overlap(toupdate, other, id, otherid):
    """
    Merges the segments specified by `id` (found in `toupdate`) and `otherid`
    (found in `other`) if they overlap at all. Updates `toupdate` accordingly.
    """
    # If there is any overlap or touching, merge the two, otherwise just return
    yourmask = other == otherid
    mymask = toupdate == id
    overlap_exists = np.any(yourmask & mymask)
    if not overlap_exists:
        return

    yourfidxs, yoursidxs = np.where(other == otherid)
    toupdate[yourfidxs, yoursidxs] = id

def _segments_are_adjacent(seg1, seg2):
    """
    Checks if seg1 and seg2 are adjacent at any point. Each is a tuple of the form
    (fidxs, sidxs).
    """
    # TODO: This is unnacceptably slow
    lsf1, lss1 = seg1
    lsf2, lss2 = seg2
    for i, f1 in enumerate(lsf1):
        for j, f2 in enumerate(lsf2):
            if f1 <= f2 + 1 and f1 >= f2 - 1:
                # Frequencies are a match, are samples?
                if lss1[i] <= lss2[j] + 1 and lss1[i] >= lss2[j] - 1:
                    return True
    return False

def _merge_adjacent_segments(mask):
    """
    Merges all segments in `mask` which are touching.
    """
    mask_ids = [id for id in np.unique(mask) if id != 0]
    for id in mask_ids:
        myfidxs, mysidxs = np.where(mask == id)
        for other in mask_ids:  # Ugh, brute force O(N^2) algorithm.. gross..
            if id == other:
                continue
            else:
                other_fidxs, other_sidxs = np.where(mask == other)
                if _segments_are_adjacent((myfidxs, mysidxs), (other_fidxs, other_sidxs)):
                    mask[other_fidxs, other_sidxs] = id  # This may lead to additional adjacencies, but we only do this once - otherwise too much clustering

def _integrate_segmentation_masks(segmasks):
    """
    `segmasks` should be in sorted order of [coarsest, ..., finest].

    Integrates the given list of segmentation masks together to form one segmentation mask
    by having each segment subsume ones that exist in the finer masks.
    """
    if len(segmasks) == 1:
        return segmasks

    assert len(segmasks) > 0, "Passed in empty list of segmentation masks"
    coarse_mask = np.copy(segmasks[0])
    mask_ids = [id for id in np.unique(coarse_mask) if id != 0]
    for id in mask_ids:
        for mask in segmasks[1:]:
            finer_ids = [i for i in np.unique(mask) if i != 0]
            for finer_id in finer_ids:
                _update_segmentation_mask_if_overlap(coarse_mask, mask, id, finer_id)

    # Lastly, merge all adjacent blocks, but just kidding, since this algorithm is waaaay to slow
    #_merge_adjacent_segments(coarse_mask)
    return coarse_mask

def _separate_masks_task(id, threshold, mask):
    idxs = np.where(mask == id)
    if len(idxs[0]) > threshold:
        m = np.zeros_like(mask)
        m[idxs] = id
        return m
    else:
        return None

def _separate_masks(mask, threshold=0.025):
    """
    Returns a list of segmentation masks each of the same dimension as the input one,
    but where they each have exactly one segment in them and all other samples in them
    are zeroed.

    Only bothers to return segments that are larger in total area than `threshold * mask.size`.
    """
    try:
        ncpus = multiprocessing.cpu_count()
    except NotImplementedError:
        ncpus = 2

    with multiprocessing.Pool(processes=ncpus) as pool:
        mask_ids = [id for id in np.unique(mask) if id != 0]
        thresholds = [threshold * mask.size for _ in range(len(mask_ids))]
        masks = [mask for _ in range(len(mask_ids))]
        ms = pool.starmap(_separate_masks_task, zip(mask_ids, thresholds, masks))
    return [m for m in ms if m is not None]

def _get_downsampled_indexes(arr, factor):
    fractional_component = factor - int(factor)
    indexes_to_keep = []
    overflow_counter = 0.0
    for index in range(arr.shape[1]):
        # if we overflowed, skip this item
        if overflow_counter >= 1.0:
            overflow_counter = 0.0
            continue

        # if the integer component of the factor is satisfied, skip this item
        if int(factor) != 1 and index % int(factor) == 0:
            overflow_counter += fractional_component
            continue
        elif int(factor) == 1:
            overflow_counter += fractional_component

        # Otherwise, add this index
        indexes_to_keep.append(index)
    return indexes_to_keep

def _downsample_one_or_the_other(mask, mask_indexes, stft, stft_indexes):
    """
    Takes the given `mask` and `stft`, which must be matrices of shape `frequencies, times`
    and downsamples one of them into the other one's times, so that the time dimensions
    are equal. Leaves the frequency dimension untouched.
    """
    assert len(mask.shape) == 2, "Expected a two-dimensional `mask`, but got one of {} dimensions.".format(len(mask.shape))
    assert len(stft.shape) == 2, "Expected a two-dimensional `stft`, but got one of {} dimensions.".format(len(stft.shape))

    if mask.shape[1] > stft.shape[1]:
        downsample_factor = mask.shape[1] / stft.shape[1]
        indexes = _get_downsampled_indexes(mask, downsample_factor)
        mask = mask[:, indexes]
        mask_indexes = np.array(indexes)
    elif mask.shape[1] < stft.shape[1]:
        downsample_factor = stft.shape[1] / mask.shape[1]
        indexes = _get_downsampled_indexes(stft, downsample_factor)
        stft = stft[:, indexes]
        stft_indexes = np.array(indexes)

    return mask, mask_indexes, stft, stft_indexes

def _map_segmentation_mask_to_stft_domain(mask, times, frequencies, stft_times, stft_frequencies):
    """
    Maps the given `mask`, which is in domain (`frequencies`, `times`) to the new domain (`stft_frequencies`, `stft_times`)
    and returns the result.
    """
    assert mask.shape == (frequencies.shape[0], times.shape[0]), "Times is shape {} and frequencies is shape {}, but mask is shaped {}".format(
        times.shape, frequencies.shape, mask.shape
    )
    result = np.zeros((stft_frequencies.shape[0], stft_times.shape[0]))

    if len(stft_times) > len(times):
        all_j = [j for j in range(len(stft_times))]
        idxs  = [int(i) for i in np.linspace(0, len(times) - 1, num=len(stft_times))]
        all_i = [all_j[idx] for idx in idxs]
    else:
        all_i = [i for i in range(len(times))]
        idxs  = [int(i) for i in np.linspace(0, len(stft_times) - 1, num=len(times))]
        all_j = [all_i[idx] for idx in idxs]

    for i, j in zip(all_i, all_j):
        result[:, j] = np.interp(stft_frequencies, frequencies, mask[:, i])

    return result

def _asa_task(q, masks, stft, sample_width, frame_rate, nsamples_for_each_fft):
    """
    Worker for the ASA algorithm's multiprocessing step.
    """
    # Convert each mask to (1 or 0) rather than (ID or 0)
    for mask in masks:
        mask = np.where(mask > 0, 1, 0)

    # Multiply the masks against STFTs
    masks = [mask * stft for mask in masks]

    nparrs = []
    dtype_dict = {1: np.int8, 2: np.int16, 4: np.int32}
    dtype = dtype_dict[sample_width]
    for m in masks:
        _times, nparr = signal.istft(m, frame_rate, nperseg=nsamples_for_each_fft)
        nparr = nparr.astype(dtype)
        nparrs.append(nparr)

    for m in nparrs:
        q.put(m)
    q.put("DONE")


"""
This module contains a bunch of functions that are integral to the
auditory event detection algorithm used by AudioSegment. We refactored
them to here because they aren't really useful on their own, and they
take up brainspace by being in the AudioSegment class.
"""
import random

def _get_filter_indices(seg, start_as_yes, prob_raw_yes, ms_per_input, model, transition_matrix, model_stats):
    """
    Runs a Markov Decision Process over the given `seg` in chunks of `ms_per_input`, yielding `True` if
    this `ms_per_input` chunk has been classified as positive (1) and `False` if this chunk has been
    classified as negative (0).

    :param seg:                 The AudioSegment to apply this algorithm to.
    :param start_as_yes:        If True, the first `ms_per_input` chunk will be classified as positive.
    :param prob_raw_yes:        The raw probability of finding the event in any given independently sampled `ms_per_input`.
    :param ms_per_input:        The number of ms of AudioSegment to be fed into the model at a time.
    :param model:               The model, which must hava predict() function, which takes an AudioSegment of `ms_per_input`
                                number of ms and which outputs 1 if the audio event is detected in that input, 0 if not.
    :param transition_matrix:   An iterable of the form: [p(yes->no), p(no->yes)].
    :param model_stats:         An iterable of the form: [p(reality=1|output=1), p(reality=1|output=0)].
    :yields:                    `True` if the event has been classified in this chunk, `False` otherwise.
    """
    filter_triggered = 1 if start_as_yes else 0
    prob_raw_no = 1.0 - prob_raw_yes
    for segment, _timestamp in seg.generate_frames_as_segments(ms_per_input):
        yield filter_triggered
        observation = int(round(model.predict(segment)))
        assert observation == 1 or observation == 0, "The given model did not output a 1 or a 0, output: "\
                + str(observation)
        prob_hyp_yes_given_last_hyp = 1.0 - transition_matrix[0] if filter_triggered else transition_matrix[1]
        prob_hyp_no_given_last_hyp  = transition_matrix[0] if filter_triggered else 1.0 - transition_matrix[1]
        prob_hyp_yes_given_data = model_stats[0] if observation == 1 else model_stats[1]
        prob_hyp_no_given_data = 1.0 - model_stats[0] if observation == 1 else 1.0 - model_stats[1]
        hypothesis_yes = prob_raw_yes * prob_hyp_yes_given_last_hyp * prob_hyp_yes_given_data
        hypothesis_no  = prob_raw_no * prob_hyp_no_given_last_hyp  * prob_hyp_no_given_data
        # make a list of ints - each is 0 or 1. The number of 1s is hypotheis_yes * 100
        # the number of 0s is hypothesis_no * 100
        distribution = [1 for i in range(int(round(hypothesis_yes * 100)))]
        distribution.extend([0 for i in range(int(round(hypothesis_no * 100)))])
        # shuffle
        random.shuffle(distribution)
        filter_triggered = random.choice(distribution)

def _group_filter_values(seg, filter_indices, ms_per_input):
    """
    Takes a list of 1s and 0s and returns a list of tuples of the form:
    ['y/n', timestamp].
    """
    ret = []
    for filter_value, (_segment, timestamp) in zip(filter_indices, seg.generate_frames_as_segments(ms_per_input)):
        if filter_value == 1:
            if len(ret) > 0 and ret[-1][0] == 'n':
                ret.append(['y', timestamp])  # The last one was different, so we create a new one
            elif len(ret) > 0 and ret[-1][0] == 'y':
                ret[-1][1] = timestamp  # The last one was the same as this one, so just update the timestamp
            else:
                ret.append(['y', timestamp])  # This is the first one
        else:
            if len(ret) > 0 and ret[-1][0] == 'n':
                ret[-1][1] = timestamp
            elif len(ret) > 0 and ret[-1][0] == 'y':
                ret.append(['n', timestamp])
            else:
                ret.append(['n', timestamp])
    return ret

def _homogeneity_filter(ls, window_size):
    """
    Takes `ls` (a list of 1s and 0s) and smoothes it so that adjacent values are more likely
    to be the same.

    :param ls:          A list of 1s and 0s to smooth.
    :param window_size: How large the smoothing kernel is.
    :returns:           A list of 1s and 0s, but smoother.
    """
    # TODO: This is fine way to do this, but it seems like it might be faster and better to do a Gaussian convolution followed by rounding
    k = window_size
    i = k
    while i <= len(ls) - k:
        # Get a window of k items
        window = [ls[i + j] for j in range(k)]
        # Change the items in the window to be more like the mode of that window
        mode = 1 if sum(window) >= k / 2 else 0
        for j in range(k):
            ls[i+j] = mode
        i += k
    return ls


"""
Convenience functions for using Numpy/Scipy filters in the audio domain.
"""
import numpy as np
import scipy.signal as signal

def bandpass_filter(data, low, high, fs, order=5):
    """
    Does a bandpass filter over the given data.

    :param data: The data (numpy array) to be filtered.
    :param low: The low cutoff in Hz.
    :param high: The high cutoff in Hz.
    :param fs: The sample rate (in Hz) of the data.
    :param order: The order of the filter. The higher the order, the tighter the roll-off.
    :returns: Filtered data (numpy array).
    """
    nyq = 0.5 * fs
    low = low / nyq
    high = high / nyq
    b, a = signal.butter(order, [low, high], btype='band')
    y = signal.lfilter(b, a, data)
    return y

def lowpass_filter(data, cutoff, fs, order=5):
    """
    Does a lowpass filter over the given data.

    :param data: The data (numpy array) to be filtered.
    :param cutoff: The high cutoff in Hz.
    :param fs: The sample rate in Hz of the data.
    :param order: The order of the filter. The higher the order, the tighter the roll-off.
    :returns: Filtered data (numpy array).
    """
    nyq = 0.5 * fs
    normal_cutoff = cutoff / nyq
    b, a = signal.butter(order, normal_cutoff, btype='low', analog=False)
    y = signal.lfilter(b, a, data)
    return y


from BlackBoxAuditing.model_factories.AbstractModelFactory import AbstractModelFactory
from BlackBoxAuditing.model_factories.AbstractModelVisitor import AbstractModelVisitor

from sklearn import tree
import numpy as np
import copy

class ModelFactory(AbstractModelFactory):
    def __init__(self, *args, **kwargs):
        # Default DecisionTreeClassifier params
        self.max_depth = None
        self.min_samples_split = 2
        self.min_samples_leaf = 1
        self.max_features = None  # try "sqrt", or "log2"
        self.random_state = 123
        self.init_shuffle = True
        self.feats_to_ignore = []

        if "options" in kwargs:
            options = kwargs["options"]
            if "max_depth" in options:
                self.max_depth = options.pop("max_depth")
            if "min_samples_split" in options:
                self.min_samples_split = options.pop("min_samples_split")
            if "min_samples_leaf" in options:
                self.min_samples_leaf = options.pop("min_samples_leaf")
            if "max_features" in options:
                self.max_features = options.pop("max_features")
            if "random_state" in options:
                self.random_state = options.pop("random_state")

        super(ModelFactory, self).__init__(*args, **kwargs)
        self.verbose_factory_name = "DecisionTree"

        # Maps each header to the values for it's column.
        self.col_vals = {header:{row[i] for row in self.all_data} for i, header in enumerate(self.headers)}
        self.response_index = self.headers.index(self.response_header)
        self.num_outcomes = len(self.col_vals[self.response_header])

        # Mark any categorical features for column expansion.
        # Categorical features are transfered into -1, 1 binary features for each value.
        self.columns_to_expand = []
        for i, header in enumerate(self.headers):
            categorical = any(type(val)==str for val in self.col_vals[header])
            if i == self.response_index or categorical:
                self.columns_to_expand.append(header)

        # All non-categorical features will be standardized (mean=0, variance=1).
        self.standardizers = {}
        for i, header in enumerate(self.headers):
            if any (type(row[i])==str for row in self.all_data):
                mean = None
                std_dev = None
            else:
                mean = sum(row[i] for row in self.all_data)/float(len(self.all_data))
                std_dev = np.sqrt(sum(((row[i] - mean) ** 2) for row in self.all_data)/float(len(self.all_data)))
            self.standardizers[header] = {'mean':mean, 'std_dev':std_dev}

        # If the response column is shifted by expanding categorical features,
        # update the response index.
        response_col_shift = 0
        for header in self.headers[:self.response_index]:
            if header in self.columns_to_expand:
                response_col_shift += len(self.col_vals[header]) - 1
        self.adjusted_response_index = self.response_index + response_col_shift

        # Map each outcome to an index
        self.outcome_trans_dict = {val:i for i, val in enumerate(self.col_vals[self.response_header])}

    def build(self, train_set):
        # prepare train data for learning
        if self.init_shuffle == True:
            np.random.seed(123)
            np.random.shuffle(train_set)
        expanded_and_stdized_train_set, self.expanded_headers = expand_and_standardize_dataset(self.response_index, self.response_header, train_set, self.col_vals, self.headers, self.standardizers, self.feats_to_ignore, self.columns_to_expand, self.outcome_trans_dict)
        train_matrix, train_outcomes = list_to_tf_input(expanded_and_stdized_train_set, self.adjusted_response_index, self.num_outcomes)
        train_size, num_features = train_matrix.shape

        clf = tree.DecisionTreeClassifier(
                max_depth = self.max_depth,\
                min_samples_split = self.min_samples_split,\
                min_samples_leaf = self.min_samples_leaf,\
                max_features = self.max_features,\
                random_state = self.random_state)

        clf.fit(train_matrix, train_outcomes)
        model_name = "DecisionTreeModel"

        return(ModelVisitor(clf, model_name, self.response_header, self.response_index, self.adjusted_response_index, self.num_outcomes, self.outcome_trans_dict, self.headers, self.expanded_headers, self.standardizers, self.col_vals, self.feats_to_ignore, self.columns_to_expand))

class ModelVisitor(AbstractModelVisitor):

  def __init__(self, clf, model_name, response_header, response_index, adjusted_response_index, num_outcomes, outcome_trans_dict, headers, expanded_headers, standardizers, train_col_vals, feats_to_ignore, columns_to_expand):
    super(ModelVisitor,self).__init__(model_name)
    self.response_index = response_index
    self.response_header = response_header
    self.adjusted_response_index = adjusted_response_index
    self.num_outcomes = num_outcomes
    self.outcome_trans_dict = outcome_trans_dict
    self.headers = headers
    self.expanded_headers = expanded_headers
    self.train_expanded_headers = expanded_headers
    self.train_col_vals = train_col_vals
    self.standardizers = standardizers
    self.feats_to_ignore = feats_to_ignore
    self.columns_to_expand = columns_to_expand
    self.clf = clf

  def test(self, test_set, test_name=""):
    expanded_and_stdized_test_set, self.test_set_expanded_headers = expand_and_standardize_dataset(self.response_index, self.response_header, test_set, self.train_col_vals, self.headers, self.standardizers, self.feats_to_ignore, self.columns_to_expand, self.outcome_trans_dict)
    if self.test_set_expanded_headers != self.train_expanded_headers:
      raise ValueError('Feature dimensions do not align!')
    else:
      test_matrix, test_labels = list_to_tf_input(expanded_and_stdized_test_set, self.adjusted_response_index, self.num_outcomes)

    predictions = self.clf.predict(test_matrix)
    predictions, test_labels = predictions.tolist(), test_labels.tolist() # cast from numpy array to python list
    test_labels = [(corr[0], corr[1]) for corr in test_labels]
    predictions = [row.index(1) for row in predictions]
    predictions_dict = {i:key for key,i in list(self.outcome_trans_dict.items())}
    predictions = [predictions_dict[pred] for pred in predictions]
    return list(zip([row[self.response_index] for row in test_set], predictions))

def list_to_tf_input(data, response_index, num_outcomes):
  """
  Separates the outcome feature from the data and creates the onehot vector for each row.
  """
  matrix = np.matrix([row[:response_index] + row[response_index+1:] for row in data])
  outcomes = np.asarray([row[response_index] for row in data], dtype=np.uint8)
  outcomes_onehot = (np.arange(num_outcomes) == outcomes[:, None]).astype(np.float32)

  return matrix, outcomes_onehot

def expand_and_standardize_dataset(response_index, response_header, data_set, col_vals, headers, standardizers, feats_to_ignore, columns_to_expand, outcome_trans_dict):
  """
  Standardizes continuous features and expands categorical features.
  """
  # expand and standardize
  modified_set = []
  for row_index, row in enumerate(data_set):
    new_row = []
    for col_index, val in enumerate(row):
      header = headers[col_index]

      # Outcome feature -> index outcome
      if col_index == response_index:
        new_outcome = outcome_trans_dict[val]
        new_row.append(new_outcome)

      # Ignored feature -> pass
      elif header in feats_to_ignore:
        pass

      # Categorical feature -> create new binary column for each possible value of the column
      elif header in columns_to_expand:
        for poss_val in col_vals[header]:
          if val == poss_val:
            new_cat_val = 1.0
          else:
            new_cat_val = -1.0
          new_row.append(new_cat_val)

      # Continuous feature -> standardize value with respect to its column
      else:
        new_cont_val = float((val - standardizers[header]['mean']) / standardizers[header]['std_dev'])
        new_row.append(new_cont_val)

    modified_set.append(new_row)

  # update headers to reflect column expansion
  expanded_headers = []
  for header in headers:
    if header in feats_to_ignore:
      pass
    elif (header in columns_to_expand) and (header is not response_header):
      for poss_val in col_vals[header]:
        new_header = '{}_{}'.format(header,poss_val)
        expanded_headers.append(new_header)
    else:
      expanded_headers.append(header)

  return modified_set, expanded_headers

def test():
  test_expand_and_standardize_dataset()
  test_unseen_categorical_feature()
  test_categorical_model()
  test_categorical_response()
  test_list_to_tf_input()
  test_basic_model()
  test_simple_data()

def test_expand_and_standardize_dataset():
  """
  Tests if the function correctly converts categorical
  features into onehot features, correctly standardizes
  continuous features, and correctly converts the response
  feature to indexes.
  """
  headers = ['cont_feature', 'cat_feature', 'ignored', 'response']
  data = [[75., 'A', 'ingore1', 'X'], [25, 'B', 'ignore2', 'Y'], [75, 'C', 'ignore3', 'Z'], [25., 'A', 'ignore1', 'Y']]
  resp_index = 3
  resp_header = 'response'
  feats_to_ignore = 'ignored'
  col_vals = {'cont_feature':[75., 25.], 'cat_feature':['A','B','C'], 'ignored':['ignore1', 'ignore2', 'ignore3'], 'response':['X', 'Y', 'Z']}
  columns_to_expand = ['cat_feature', 'ignored', 'response']
  outcome_trans_dict = {'X':0, 'Y':1, 'Z':2}
  standardizers = {'cont_feature':{'mean':50, 'std_dev':25}}
  new_data, new_headers = expand_and_standardize_dataset(resp_index, resp_header, data, col_vals, headers, standardizers, feats_to_ignore, columns_to_expand, outcome_trans_dict)
  correct_data = [[1.0, 1.0, -1.0, -1.0, 0], [-1.0, -1.0, 1.0, -1.0, 1], [1.0, -1.0, -1.0, 1.0, 2], [-1.0, 1.0, -1.0, -1.0, 1]]
  correct_headers = ['cont_feature', 'cat_feature_A', 'cat_feature_B', 'cat_feature_C', 'response']
  print('expanding and standardizing data correctly? --', new_data==correct_data)
  print('expanding headers correctly? --', new_headers==correct_headers)

def test_unseen_categorical_feature():
  """
  Tests if the model can correctly handle unfamiliar values
  for categorical features. The model should not create new onehot
  features for values in the test set and should instead treat
  these values as 'not any known value for that column'.
  """
  headers = ["predictor 1", "predictor 2", "predictor3", "response"]
  response = "response"
  random_features = ['H', 5.5, 'I']
  train_set = [[random_features[i%3],'B',1,'C'] for i in range(1,50)] + [['A','A',-1,'D'] for i in range(1,50)]
  train_set_copy = copy.copy(train_set)
  new_features = ['J',9]
  test_set = [[new_features[i%2],'B',1,'C'] for i in range(1,49)] + [['A','A',-1,'D'] for i in range(1,50)]
  all_data = train_set + test_set

  factory = ModelFactory(all_data, headers, response, name_prefix="test")
  model = factory.build(train_set)
  print("factory builds ModelVisitor? -- ", isinstance(model, ModelVisitor))

  predictions = model.test(test_set)
  print("handling unseen categorical values correctly? -- ", all([pred[0] == pred[1] for pred in predictions]))

def test_list_to_tf_input():
  data = [[0,0],[0,1],[0,2]]
  tf_matrix, tf_onehot = list_to_tf_input(data, 1, 3)
  correct_matrix = [[0],[0],[0]]
  correct_onehot = [[1,0,0], [0,1,0], [0,0,1]]
  print("list_to_tf_input matrix correct? --",np.array_equal(tf_matrix, correct_matrix))
  print("list_to_tf_input onehot correct? --",np.array_equal(tf_onehot, correct_onehot))

def test_basic_model():
  headers = ["predictor 1", "predictor 2", "response"]
  response = "response"
  train_set = [[i,0,0] for i in range(1,50)] + [[0,i,1] for i in range(1,50)]
  test_set = [[i,0,0] for i in range(1,50)] + [[0,i,1] for i in range(1,50)]
  all_data = train_set + test_set

  factory = ModelFactory(all_data, headers, response, name_prefix="test")

  model = factory.build(train_set)

  predictions = model.test(test_set)
  intended_predictions = [pred[0] for pred in predictions]
  actual_predictions = [pred[1] for pred in predictions]
  print("predicting numeric categories correctly? -- ", actual_predictions == intended_predictions)

def test_categorical_response():
  headers = ["predictor 1", "predictor 2", "response"]
  response = "response"
  train_set = [[i,0,"A"] for i in range(1,50)] + [[0,i,"B"] for i in range(1,50)]
  test_set = [[i,0,"A"] for i in range(1,50)] + [[0,i,"C"] for i in range(1,50)]
  train_set_copy = copy.copy(train_set)
  test_set_copy = copy.copy(test_set)
  all_data = train_set + test_set

  factory = ModelFactory(all_data, headers, response, name_prefix="test")
  model = factory.build(train_set)

  predictions = model.test(test_set)
  resp_index = headers.index(response)
  intended_predictions = [(test_row[resp_index], train_row[resp_index]) for train_row, test_row in zip(train_set_copy,test_set_copy)]
  print("predicting string-categories correctly? -- ", predictions == intended_predictions)

def test_categorical_model():
  """
  Tests whether the model can correctly learn on the train set.
  The model should have poor accuracy on the test set
  because the outcomes do not match those of the train set.
  The model should incorrectly predict response 'B' for predictor 'B'
  when response 'C' is the correct response.
  """
  headers = ["predictor", "response"]
  response = "response"
  train_set = [["A","A"] for i in range(1,50)] + [["B","B"] for i in range(1,50)]
  train_set_copy = copy.copy(train_set)
  test_set = [["A","A"] for i in range(1,50)] + [["B","C"] for i in range(1,50)]
  all_data = train_set + test_set
  factory = ModelFactory(all_data, headers, response, name_prefix="test", options={})

  model = factory.build(train_set)

  predictions = model.test(test_set)
  resp_index = headers.index(response)
  intended_predictions = [(test_row[resp_index], train_row[resp_index]) for train_row, test_row in zip(train_set_copy,test_set)]
  print("predicting string-categories correctly? -- ", predictions == intended_predictions)

def test_simple_data():
    headers = ["predictor_1", "predictor_2", "response"]
    response = "response"

    train_set = [["A","B","C"] for i in range(1,50)] + [["B","A","D"] for i in range(1,50)] + [["B","B","E"] for i in range(1,50)] + [["A", "A", "E"] for i in range(1,50)] + [["C", "D", "F"] for i in range(1,10)]
    train_set_copy = copy.copy(train_set)
    test_set = [["A","B","C"] for i in range(1,50)] + [["B","A","D"] for i in range(1,50)] + [["B","B","E"] for i in range(1,50)] + [["A", "A", "E"] for i in range(1,50)] + [["A","X","Y"] for i in range(1,10)]
    all_data = train_set + test_set
    factory = ModelFactory(all_data, headers, response, name_prefix="test", options={})

    model = factory.build(train_set)
    predictions = model.test(test_set)
    resp_index = headers.index(response)

    intended_predictions = [(test_row[resp_index], train_row[resp_index]) for train_row, test_row in zip(train_set_copy,test_set)]
    last_few_preds = all(pred[0] == pred[1] for pred in predictions[-9:])
    all_but_last_few = all(pred[0] == pred[1] for pred in predictions[:-9])
    print("handling simple data correctly? -- ", (all_but_last_few and not last_few_preds) == True)

if __name__=='__main__':
  test()


import networkx as nx
from networkx.algorithms.flow import (preflow_push, edmonds_karp, shortest_augmenting_path)
from collections import defaultdict
import random

class CategoricalFeature:
  def __init__(self, data, name="no_name"):
    self.data = data
    self.name = name
    d1=defaultdict(int) #bin_data
    d2=defaultdict(int) #bin_index_dict
    d3=defaultdict(list) #bin_fulldata
    d4=defaultdict(int) #bin_index_dict_reverse
    d5=defaultdict(int) #category_count

    n = len(self.data)
    count = 0
    for i in range(0,n):
      obs = self.data[i]
      if obs in d2: pass  # if obs (i.e. category) is alreay a KEY in bin_index_data then don't do anything
      else:
        d2[obs] = count #bin_index_dict inits the KEY: category, with VALUE: count
        d4[count] = obs #bin_index_dict_reverse does the opposite
        count += 1
      bin_idx = d2[obs]
      d1[bin_idx] += 1 #add 1 to the obs category idex in bin_data
      d5[obs] += 1 #add 1 to the obs category NAME in category_count
      d3[bin_idx].append(i) #add obs to the list of obs with that category in bin_fulldata

    self.bin_data = d1
    self.category_count = d5
    self.num_bins = len(list(d1.items()))
    self.bin_fulldata = d3
    self.bin_index_dict = d2
    self.bin_index_dict_reverse = d4

  def create_graph(self, count_generator): #creates graph given a CategoricalFeature object
    DG=nx.DiGraph() #using networkx package
    bin_list = list(self.bin_data.items())
    bin_index_dict_reverse = self.bin_index_dict_reverse
    k = self.num_bins
    total_items = 0
    for i, amt in bin_list:
        total_items += amt
    DG.add_node('s')
    DG.add_node('t')
    for i in range(0, k): #lefthand side nodes have capacity = number of observations in category i
      DG.add_node(i)
      DG.add_edge('s', i, capacity=bin_list[i][1], weight=0)
    for i in range(k, 2*k): #righthand side nodes have capacity = DESIRED number of observations in category i
      DG.add_node(i)
      cat = bin_index_dict_reverse[i-k]
      desired_count = count_generator(cat)
      DG.add_edge(i, 't', capacity=desired_count, weight=0)
    #Add special node to hold overflow
    DG.add_node(2*k)
    DG.add_edge(2*k, 't', capacity=total_items, weight=0)
    for i in range(0, k):
      for j in range(k,2*k): #for each edge from a lefthand side node to a righhand side node:
        if (i+k)==j:  #IF they represent the same category, the edge weight is 0
            DG.add_edge(i, j, weight=0)
        else: #IF they represent different categories, the edge weight is 1
          DG.add_edge(i, j, weight=1)
      #THIS IS THE OVERFLOW NODE!!
      DG.add_edge(i, 2*k, weight=2)
    return DG

  def repair(self, DG): #new_feature = repair_feature(feature, create_graph(feature))
    mincostFlow = nx.max_flow_min_cost(DG, 's', 't') #max_flow_min_cost returns Dictionary of dictionaries. Keyed by nodes such that mincostFlow[u][v] is the flow edge (u,v)
    bin_dict = self.bin_fulldata
    index_dict = self.bin_index_dict_reverse
    size_data = len(self.data)
    repair_bin_dict = {}
    repair_data = [0]*size_data #initialize repaired data to be 0. If there are zero's after we fill it in the those observations belong in the overflow, "no category"
    k = self.num_bins
    overflow = 0
    for i in range(0,k): #for each lefthand side node i
      overflow += mincostFlow[i][2*k]
      for j in range(k, 2*k): #for each righthand side node j
        edgeflow = mincostFlow[i][j] #get the int (edgeflow) representing the amount of observations going from node i to j
        group = random.sample(bin_dict[i], int(edgeflow)) #randomly sample x (edgeflow) unique elements from the list of observations in that category.
        q=j-k #q is the category index for a righhand side node
        for elem in group: #for each element in the randomly selected group list
          bin_dict[i].remove(elem) #remove the element from the list of observation in that category
          repair_data[elem] = index_dict[q] #Mutate repair data at the index of the observation (elem) with its new category (it was 0) which is the category index for the righthand side node it flows to
        if q in repair_bin_dict: #if the category index is already keyed
          repair_bin_dict[q].extend(group) #extend the list of observations with a new list of observations in that category
        else:
          repair_bin_dict[q] = group #otherwise key that category index and set it's value as the group list in that category
    new_feature = CategoricalFeature(repair_data) #initialize our new_feature (repaired feature)
    new_feature.bin_fulldata = repair_bin_dict
    return new_feature, overflow

def test():
  random.seed(10)
  test_feature = CategoricalFeature(["A","B","C","D","D","D","C","B","A","C","B","A"])
  desired_count_dict = {"A": 1, "B": 2, "C": 2, "D": 3}
  desired_category_count = lambda category : desired_count_dict[category]
  DG = test_feature.create_graph(desired_category_count)
  new_feature, overflow = test_feature.repair(DG)
  edges = [(0, 8, {'weight': 2}), (0, 4, {'weight': 0}), (0, 5, {'weight': 1}), (0, 6, {'weight': 1}), (0, 7, {'weight': 1}), (1, 8, {'weight': 2}), (1, 4, {'weight': 1}), (1, 5, {'weight': 0}), (1, 6, {'weight': 1}), (1, 7, {'weight': 1}), (2, 8, {'weight': 2}), (2, 4, {'weight': 1}), (2, 5, {'weight': 1}), (2, 6, {'weight': 0}), (2, 7, {'weight': 1}), (3, 8, {'weight': 2}), (3, 4, {'weight': 1}), (3, 5, {'weight': 1}), (3, 6, {'weight': 1}), (3, 7, {'weight': 0}), (4, 't', {'capacity': 1, 'weight': 0}), (5, 't', {'capacity': 2, 'weight': 0}), (6, 't', {'capacity': 2, 'weight': 0}), (7, 't', {'capacity': 3, 'weight': 0}), (8, 't', {'weight': 0}), ('s', 0, {'capacity': 3, 'weight': 0}), ('s', 1, {'capacity': 3, 'weight': 0}), ('s', 2, {'capacity': 3, 'weight': 0}), ('s', 3, {'capacity': 3, 'weight': 0})]

  new_data = [0, 'B', 'C', 'D', 'D', 'D', 'C', 'B', 0, 0, 0, 'A']
  print("CategoricalFeature has correct number of categories?", 4 == test_feature.num_bins)
  print("Directed Graph has correct edges and edge weights?", equal_ignore_order(DG.edges(data=True), edges))
  print("mincostFlow has correct overflow?", overflow == 4)
  print("mincostFlow has correct output data?", new_feature.data == new_data)

def equal_ignore_order(a, b):
  """
  Used to check whether the two edge lists have the same edges
  when elements are neither hashable nor sortable.
  """
  unmatched = list(b)
  for element in a:
    try:
      unmatched.remove(element)
    except ValueError:
      return False
  return not unmatched

if __name__=="__main__": test()