import librosa
import librosa.display
from scipy.spatial import distance
import numpy
import matplotlib.pyplot as plt


# calculates the averages of a librosa stft output
def calc_stft_averages(stft):
    return [numpy.average(t) for t in stft]


# calculates the euclidean distance of two full signal stft ouputs
def calc_euclidean_distance(stft1, stft2):
    stft1_averages = calc_stft_averages(stft1)
    stft2_averages = calc_stft_averages(stft2)
    euclidean_distances = []
    for index in range(0, (len(stft1_averages) - 1)):
        euclidean_distances.append(distance.euclidean(stft1_averages[index], stft2_averages[index]))
    return numpy.average(euclidean_distances)


# calculates segment scores for a list of beat clips
def score_segments(clips, segment_length, midpoint, density, bias_mode=False):
    scores = []
    # only compute scores for first or last third of song depending on bias mode
    if bias_mode:
        #precalculate stfts of needed clips
        clip_stfts = [numpy.abs(librosa.stft(clips[i])) for i in range(0, int(len(clips) * 0.33))]
        for si in range(0, (int(len(clip_stfts) * 0.33)-int(segment_length * density))):
            if midpoint > 2:
                score_first = numpy.average(
                    [calc_euclidean_distance(clip_stfts[si], clip_stfts[si+i]) for i in range(1, (midpoint*density - 1), density)])
                score_second = numpy.average([calc_euclidean_distance(clip_stfts[si], clip_stfts[si+i]) for i in
                                          range((midpoint*density), (segment_length*density - 1), density)])
            else:
                score_first = calc_euclidean_distance(clip_stfts[si], clip_stfts[si+1*density])
                score_second = numpy.average([calc_euclidean_distance(clip_stfts[si], clip_stfts[si+i]) for i in
                                          range((midpoint*density) , (segment_length*density - 1), density)])

            score = score_first + score_second

            scores.append(score)

        #create dummy values
        for si in range(0, int(len(clips) * 0.66)):
            scores.append(1000)
    else:
        #precalculate stfts of needed clips
        clip_stfts = [numpy.abs(librosa.stft(clips[i])) for i in range(int(len(clips) * 0.66), (len(clips) - 1))]
        #create dummy values to
        for si in range(0, (int(len(clips) * 0.66))):
            scores.append(1000)
        for si in range(0, (len(clip_stfts) - int(segment_length * density))):
            if midpoint > 2:
                score_first = numpy.average(
                    [calc_euclidean_distance(clip_stfts[si], clip_stfts[si+i]) for i in range(1, (midpoint*density - 1), density)])
                score_second = numpy.average([calc_euclidean_distance(clip_stfts[si], clip_stfts[si+i]) for i in
                                          range((midpoint*density), (int(segment_length*density) - 1), density)])
            else:
                score_first = calc_euclidean_distance(clip_stfts[si], clip_stfts[si+1*density])
                score_second = numpy.average([calc_euclidean_distance(clip_stfts[si], clip_stfts[si+i]) for i in
                                          range((midpoint*density), (segment_length*density - 1), density)])

            score = score_first + score_second

            scores.append(score)
    return scores