function get_maxima_minima(amplitudes)
    # Find all local maxima
    maxima = zeros(size(amplitudes))
    begin
        maxima_idx = findall(diff(sign.(diff(amplitudes))) .== -2) .+ 1
        maxima[maxima_idx] .= 1.0
        if amplitudes[1] > amplitudes[2]
            maxima[1] = 1.0
        end
        if amplitudes[end] > amplitudes[end-1]
            maxima[end] = 1.0
        end
    end
    maxima .*= amplitudes

    minima = ones(size(amplitudes)) .* Inf
    begin
        minima_idx = findall(diff(sign.(diff(amplitudes))) .== 2) .+ 1
        if amplitudes[1] < amplitudes[2]
            minima[1] = amplitudes[1]
        end
        if amplitudes[end] < amplitudes[end-1]
            maxima[end] = amplitudes[end]
        end
        minima[minima_idx] .= amplitudes[minima_idx]
    end

    return maxima, minima
end

function get_regions(amplitudes, minima, pitch, bin_size)

    region_size = 0.5 * pitch / bin_size + 1

    regions = []
    for k in 1:200
        idx = pitch * k / bin_size + 1
        idx_int = Int(round(idx))
        if idx_int >= length(amplitudes)
            break
        end

        region_start = max(1, Int(round(idx - region_size)))
        region_min = findmin(minima[region_start:idx_int-1])
        if region_min[1] == Inf
            region_min = findmin(amplitudes[region_start:idx_int-1])
        end
        region_start = region_min[2] + region_start - 1

        region_end = min(Int(round(idx + region_size)), length(amplitudes))
        region_min = findmin(minima[idx_int+1:region_end])
        if region_min[1] == Inf
            region_min = findmin(amplitudes[idx_int+1:region_end])
        end
        region_end = region_min[2] + idx_int - 1

        push!(regions, (idx, region_start:region_end))
    end

    idx = 1
    noise = []
    for i in 1:length(regions)
        push!(noise, idx:(regions[i][2][1]-1))
        idx = regions[i][2][end] + 1
    end
    push!(noise, (regions[end][2][end]+1):length(amplitudes))

    return regions, noise
end

function interpolate_FFT(k, s)
    N = length(s)
    sum([s[n] * exp(-im * 2 * π * k * n / N) for n in 1:N])
end

function linear_interpolation(s, k)
    x_1 = Int.(ceil.(k))
    x_0 = Int.(floor.(k))
    return (s[x_0] .* (x_1 .- k) .+ s[x_1] .* (k .- x_0)) / (x_1[1] .- x_0[1])
end

##################################################################################
using GLMakie, Makie.Colors
set_theme!(theme_black())
using WAV
using DSP
include("Utils.jl")

function get_song(file)
    y, Fs = wavread(string("Inputs/", file, ".wav"))
    low_note = Utils.getNote(1) #hz

    # Signal pre-processing
    y = [(y[i, 1] + y[i, 2]) / 2 for i = 1:length(y[:, 1])]
    function highpassfilter(signals, Fs, cutoff, order = 4)
        wdo = 2.0 * cutoff / Fs
        filth = digitalfilter(Highpass(wdo), Butterworth(order))
        filtfilt(filth, signals)
    end
    y = highpassfilter(y, Fs, low_note)
    return y, Fs
end

file = "Arpeggio"
y, Fs = get_song(file)

##################################################################################
include("Utils.jl")
using .Utils
include("SPA.jl")
using FFTW
using WAV

Wl = 2048
bin_size = Fs / Wl
hop_size = Int(floor(0.25 * Wl))
windows = Utils.hannWindow.(Utils.getWindows(y, Wl, hop_size = hop_size))


nyquist = Int(Wl / 2) + 1
fft_in = [fft(window)[1:nyquist] for window in windows]
fft_out = [zeros(ComplexF64, Wl) for n in 1:length(fft_in)]

######################################################################
# Do stuff with fft_in
for n in 1:length(fft_in)
    # Convert to polar
    r = abs.(fft_in[n])

    # Get frequencies in hz
    ω_0 = SPA.detect_pitch(fft_in[n], Fs)
    ω_end = Utils.getNote(Int(round(Utils.getNoteFromFreq(ω_0))))

    # Get ratio of difference
    Δω = (ω_end - ω_0) / ω_0

    # Decompose into regions
    maxima, minima = get_maxima_minima(r)
    regions, noise = get_regions(r, minima, ω_0, bin_size)

    # Shift frequencies
    for region in noise
        fft_out[n][region] .+= fft_in[n][region]
    end

    for region in regions
        Δbins = region[1] * Δω
        Δbins_int = Int(floor(Δbins))
        Δbins_frac = Δbins - Δbins_int

        region_intshift = region[2] .+ Δbins_int

        # Remove invalid indexes
        valid_idx = findall(x -> x >= 1 && x <= length(fft_in[n]), region_intshift)
        if length(valid_idx) > 0
            region_intshift = region_intshift[valid_idx]

            fft_region = linear_interpolation(r, region[2][valid_idx] .- Δbins_frac)

            # Phase correction
            Δϕ = ℯ^(im * Δω * hop_size)
            ϕ = ℯ .^ (im * angle.(fft_in[n][region[2][valid_idx]])) .* Δϕ

            fft_out[n][region_intshift] += fft_region .* ϕ
        end
    end
end
#####################################################################################

# Convert back to signal
for n in 1:length(fft_out)
    copyto!(fft_out[n], nyquist + 1, conj.(fft_out[n][nyquist-1:-1:2]))
end
signal = zeros(length(y) + (Wl - length(y) % Wl))
for n in 1:length(fft_out)
    range_start = (n - 1) * hop_step + 1
    range = range_start:range_start+Wl-1
    signal[range] .+= real.(ifft(fft_out[n]))
end
signal = signal[1:length(y)]
wavwrite(signal, string("Outputs/", file, "_modified.wav"), Fs = Fs)
lines(signal)