#=
    Implementation of Spectral Peak Analysis from
    "High Accuracy and Octave Error Immune Pitch Detection Algorithms"
    by M. Dziubinski and B. Kostek
=#
module SPA
    export detect_pitch

    using Flux: hardtanh
    using BSON
    include("Utils.jl")
    using .Utils

    global W = BSON.load("Models/Params.bson", @__MODULE__)[:W]
    global b = BSON.load("Models/Params.bson", @__MODULE__)[:b]
    function model(x)
        W[3]' * tanh.(W[2]' * tanh.(W[1]' * x .+ b[1]) .+ b[2]) .+ b[3]
    end
    global const d = 4

    function upscale_peak_idx(S, idx, maxima)
        if idx > 1 && idx < length(maxima) && maxima[idx] > 0
            scale = 1 / S[idx]
            x = S[[idx-1,idx+1]] * scale
            return idx + model(x)[1]
        else
            return idx
        end
    end

    function select_F_fund(S, idx_F_M, bin_size, M, maxima)
        # Select F_M
        M = min(M, Int(floor(idx_F_M / d)))
        F_M = Utils.idxToBin(upscale_peak_idx(S, idx_F_M, maxima), bin_size)
        
        # Equation (1)
        F_fund = [F_M / m for m = 1:M]

        return F_fund
    end

    function spa(S, F_fund, Fs, M, maxima)
        bin_size = Fs / 2 / length(S)
        max_freq = Fs / 2 - bin_size

        # Equation (2)
        K = Int(floor(max_freq / F_fund[1]))

        # Calculate scores
        scores = zeros(M)
        for m = 1:M
            # Equation (4) and (5)
            L = Int(floor(Fs / 2 / F_fund[m])) - K
            H_maxset = [sum([S[Utils.binIdx((i + j) * F_fund[m], bin_size)] for i = 1:K]) for j = 0:L-1]
            K_0 = L > 0 ? findmax(H_maxset)[2] : 1
            set = [F_fund[m] * k for k = K_0:K_0+K-1]
        
            # Equation (6) except it's an integral instead of a sum
            set_idx = Utils.binIdx.(set, bin_size)
            scores[m] = ∫(set_idx, maxima[set_idx])
        
            # Equation (7a) and (7b) except it's an integral instead of a sum
            half_f = F_fund[m] / 2
            F_L_ranges = [max(1, Utils.binIdx(set[k] - half_f, bin_size)):(set_idx[k]-2) for k = 1:K]
            F_R_ranges = [(set_idx[k]+2):min(length(maxima), Utils.binIdx(set[k] + half_f, bin_size)) for k = 1:K]
            F_L_ranges = F_L_ranges[findall(x -> length(x) > 0, F_L_ranges)]
            F_R_ranges = F_R_ranges[findall(x -> length(x) > 0, F_R_ranges)]
            F_L_max = [findmax(maxima[range])[2] + range[1] - 1 for range in F_L_ranges]
            F_R_max = [findmax(maxima[range])[2] + range[1] - 1 for range in F_R_ranges]
            scores[m] -= ∫(F_L_max, maxima[F_L_max])
            scores[m] -= ∫(F_R_max, maxima[F_R_max])
        end
    
        return scores
    end

    function detect_pitch(S, Fs; num_candidates = 5, M = 20)
        # Only care about amplitudes
        S = abs.(S)
        
        # Find all local maxima
        maxima = zeros(size(S))
        begin
            maxima_idx = findall(diff(sign.(diff(S))) .== -2) .+ 1
            maxima[maxima_idx] .= 1.0
            if S[1] > S[2]
                maxima[1] = 1.0
            end
            if S[end] > S[end-1]
                maxima[end] = 1.0
            end
        end
        maxima .*= S

        # Select F_fund
        bin_size = Fs / 2 / length(S)

        # HACK: Testing multiple F_Ms is just to get around transient detection, and will be removed in the next phase
        idx_F_M = findmax(S[d:end], 1:num_candidates)[2] .+ d .- 1
        F_fund = [select_F_fund(S, idx, bin_size, M, maxima) for idx in idx_F_M]
        M = [length(F_fund[n]) for n in 1:length(F_fund)]

        # Gather evidence
        scores = [spa(S, F_fund[n], Fs, M[n], maxima) for n in 1:length(F_fund)]
        best_idx = [findmax(scores[n])[2] for n in 1:length(scores)]
        best_set = findmax([scores[n][best_idx[n]] for n in 1:length(scores)])[2]

        return F_fund[best_set][best_idx[best_set]]
    end
end