from math import *
import numpy as np
from scipy.io import wavfile
from scipy import signal
import matplotlib.pyplot as plt
import subprocess
import wave as wv
import struct
import librosa
import heapq
import scipy
import os
import random
from pathlib import Path
from time import sleep
from datetime import timedelta

def adjust_timings(raw_data, snapped_data, indexes, thr=100):
    """
    adjusts weirdly snapped notes
    """
    current = 0

    while(current < len(indexes)):
    
        if(current < len(indexes) - 3 and current % 2 == 1): # on a 1/4 beat
            if(snapped_data[indexes[current]] > thr and snapped_data[indexes[current+1]] > thr and snapped_data[indexes[current+2]] > thr and snapped_data[indexes[current+3]] <= thr):
                # -XXX_
                snapped_data[indexes[current+3]] = snapped_data[indexes[current+2]]
                snapped_data[indexes[current+2]] = 0

        if(current > 0 and current < len(indexes) - 1 and current % 2 == 1):
            if(snapped_data[indexes[current]] > thr and (snapped_data[indexes[current+1]] < thr or snapped_data[indexes[current-1]] < thr)):
                #_X_
                '''if(snapped_data[indexes[current-1]] < thr and raw_data[indexes[current-1]] > raw_data[indexes[current+1]]):
                    snapped_data[indexes[current-1]] = snapped_data[indexes[current]]
                else:
                    snapped_data[indexes[current+1]] = snapped_data[indexes[current]]'''

                snapped_data[indexes[current]] = 0

        current += 1

    print("Resnap done")
    return snapped_data



def snap(data, sample_rate, bpm, divisor, show=False):
    # adjust time amplitudes to match the given BPM 
    new = [0 for x in range(int(1000*len(data)/sample_rate))] # 1pt per millisecond
    print("old =", len(data))
    print("len =", 1000*len(data)/sample_rate)
    k = 0
    t = 0
    percent = 0
    for i in range(len(data)):

        while(t < i/sample_rate):
            t = k/(bpm*divisor)
            k += 60
        '''
        if(np.abs(i/sample_rate - k/(bpm*divisor)) > np.abs(i/sample_rate - (k-60)/(bpm*divisor))):
            k -= 60
            t = k/(bpm*divisor)'''

        if(i%(len(data)//100) == 0):
            print(percent, "%")
            percent += 1

        if(int(t*1000) < len(new)):
            new[int(t*1000)] = max(data[i], new[int(t*1000)])
        else:
            new[len(new)-1] = max(data[i], new[len(new)-1])

    if(show):
        t = [j/1000 for j in range(len(new))]
        plt.plot(t, new)
        plt.xlabel("Time (e)")
        plt.ylabel("Amplitude")
        plt.grid()
        plt.show()

    return new

def snap2(data, sample_rate, bpm, first_offset=0, div=4, show=False, adjust=False):
    """
    data : list(int)
    sample_rate : int
    bpm = float
    """

    song_len = int(len(data)/sample_rate)

    indexes = []
    app = True

    reduced = [0 for i in range(song_len*1000)]
    new = [0 for i in range(song_len*1000)]

    # build the reduced version
    for i in range(len(data)):
        x = int(i*1000/sample_rate)
        if(x < len(reduced)):
            reduced[x] = max(reduced[x], data[i])


    print("Build done")
    # snap
    k = 0
    current_t = first_offset

    while(current_t < 0):
        k += 1
        current_t = first_offset + k*60/(bpm*div)

    for j in range(len(new)):
        if(j/1000 > current_t):
            k += 1
            current_t = first_offset + k*60/(bpm*div)
            app = True

        y = int(current_t*1000)
        if(y < len(new)):
            new[y] = max(new[y], reduced[j])

            if(app):
                indexes.append(y)
                app = False 

    print("Snap done")
    
    if(adjust):
        print("Len :", len(indexes))

        new = adjust_timings(reduced, new, indexes)

    
    if(show):
        new2 = [0.9 if new[i] != 0 else 0 for i in range(len(new))]

        t = [j/1000+first_offset for j in range(len(new))]
        beats_1 = [0 for j in range(len(new))]
        beats_2 = [0 for k in range(len(new))]
        beats_4 = [0 for l in range(len(new))]

        k = 0
        current_t = first_offset

        while(current_t < 0):
            k += 1
            current_t = first_offset + k*60/(bpm*div)

        while(1000*current_t < len(new)):
            beats_4[int(1000*current_t)] = 0.9

            if(k % 2 == 0):
                beats_2[int(1000*current_t)] = 0.92

            if(k % 4 == 0):
                beats_1[int(1000*current_t)] = 0.94

            k += 1
            current_t = first_offset + k*60/(bpm*div)

        plt.plot(t, new2, "bo")
        plt.plot(t, beats_4, "r-")
        plt.plot(t, beats_2, "y-")
        plt.plot(t, beats_1, "g-")
        plt.xlabel("Time (s)")
        plt.ylabel("Amplitude")
        plt.grid()
        plt.show()

    return new

def filter_peaks(data, sample_rate=44100, thr=1000):
    tdata = []
    times = []

    for i in range(len(data)):
        if data[i] > thr:
            tdata.append(data[i])
            times.append(i/sample_rate)

    return (tdata, times)
    '''
    times is in seconds
    '''

def get_spacing(data, sample_rate=44100, show=False, retrieve=False):
    tdata, times = filter_peaks(data, sample_rate=sample_rate)

    absc = [i for i in range(len(times))]
    dt = [0]

    for i in range(1, len(times)):
        dt.append(1000*(times[i]-times[i-1]))

    if(show):
        plt.plot(absc, dt)
        plt.xlabel("x")
        plt.ylabel("T(peak x) - T(peak x-1) (ms)")
        plt.grid()
        plt.show()

    if(retrieve):
        return dt

    '''
    post-condition : 
    - dt[i] = time(peak number i) - time(peak number i-1)
    - dt is in ms
    '''

def avg(data, i, j):
    res = 0
    for e in range(i, min(len(data), j)):
        res += data[e]
    return (res/(min(len(data), j) - i))

def snap3(data, sample_rate=44100, mintime=10, initial_plot=False, after_plot=False):
    '''
    explaination : 
    1) get the time differences (cf get_spacing)
    2) for eack peak : 2 cases
        - if it's farther than mintime (in ms) :
            > calculate the weighted mean if all elements in temp_list
            > place a note at that mean
            > empty temp_list
            > push the current peak to temp_list
        - else :
            > push the current peak to temp_list
    '''
    data_peaks, peak_times = filter_peaks(data, sample_rate=sample_rate)
    time_diff = get_spacing(data, show=initial_plot, retrieve=True)

    res_peaks = []
    res_times = []

    segments = []
    seglen = []

    current_left = 0

    for i in range(len(peak_times)):
        if(time_diff[i] > mintime):
            segments.append([current_left, i])
            seglen.append(peak_times[i]-peak_times[current_left])
            res_peaks.append(avg(data_peaks, current_left, i))
            res_times.append(peak_times[i])
            current_left = i

    for i in range(len(segments)):
        print(segments[i], ":", seglen[i])

    if(after_plot):
        peakplot = []
        diffplot = []
        for x in range(len(peak_times)):
            peakplot.append(peak_times[x]-peak_times[x]/1000)
            peakplot.append(peak_times[x])
            peakplot.append(peak_times[x]+peak_times[x]/1000)
            diffplot.append(0)
            diffplot.append(time_diff[x])
            diffplot.append(0)

        plt.plot(res_times, res_peaks, "ro", label="placed beats")
        plt.plot(peakplot, diffplot, label="derivatine of time")
        plt.xlabel("t (s)")
        plt.ylabel(".")
        plt.legend(loc="upper left")
        plt.grid()
        plt.show()

    return (res_peaks, res_times)
    """ res_times is in seconds """

    