include("Utils.jl")
include("Testing.jl")
using Flux
using FFTW
using LinearAlgebra
using BSON

####################
# TRAINING DATA
####################
function hannWindow(window)
    N = length(window)
    return [window[n] * 0.5 * (1 - cos(2π * n / N)) for n = 1:N]
end

function spectrum(window)
    S = fft(window)
    S = S[1:Int(length(S) / 2)]
    scale = 1 / length(window)
    global S = [sqrt(real(n)^2 + imag(n)^2) * scale for n in S]
end

function test_case(f, Fs, W)
    K = Int(floor(Fs / 2 / f)) - 1
    frequencies = [f * i for i = 1:K]
    test = Testing.generate_test(f, W, Fs)

    S = spectrum(hannWindow(test.signal))
    bin_size = Fs / 2 / length(S)

    correct_peaks = frequencies ./ bin_size .+ 1
    assumed_peaks = Int.(round.(correct_peaks))

    inputs = []
    answers = []
    for i = 1:length(assumed_peaks)
        idx = assumed_peaks[i]
        if S[idx-1] < S[idx] && S[idx] > S[idx+1]
            push!(inputs, S[[(idx - 1), (idx + 1)]] ./ S[idx])
            push!(answers, correct_peaks[i] - assumed_peaks[i])
        end
    end
    return inputs, answers
end

function get_training_data(N)
    W_list = [1024, 2048, 4096]
    Fs = 48000
    inputs = [[[] for n = 1:N] for W in W_list]
    answers = [[[] for n = 1:N] for W in W_list]
    frequencies = [[[] for n = 1:N] for W in W_list]

    for w = 1:length(W_list)
        W = W_list[w]
        bin_size = Fs / W
        low_note = bin_size * 4
        high_note = 4800

        note_range = Utils.getRange(low_note, high_note)
        Threads.@threads for n = 1:N
            for f in [note_range[(n-1)%length(note_range)+1], rand(low_note:0.01:high_note)]
                x, y = test_case(f, Fs, W)
            
                push!(inputs[w][n], x)
                push!(answers[w][n], y)
                push!(frequencies[w][n], f)
            end
        end
    end
    X = reduce(hcat, reduce(vcat, reduce(vcat, reduce(vcat, inputs))))
    Y = float.(reduce(vcat, reduce(vcat, reduce(vcat, answers))))
    frequencies = float.(reduce(vcat, reduce(vcat, reduce(vcat, answers))))
    return X, Y, frequencies
end

###############################
# MODEL
###############################
function get_params(param_sizes)
    W = Matrix{Float64}[]
    b = Vector{Float64}[]
    for layer in param_sizes
        fan_in = layer[1]
        fan_out = layer[2]
        limit = sqrt(6 / (fan_in + fan_out))
        push!(W, rand(-limit:0.001:limit, fan_in, fan_out))
        push!(b, zeros(layer[2]))
    end
    return W, b
end

function params_to_vector(W, b)
    reduce(vcat, [[W[i]..., b[i]...] for i = 1:length(W)])
end

function vector_to_params(β, param_sizes)
    W = Matrix{Float64}[]
    b = Vector{Float64}[]
    idx = 1
    for layer in param_sizes
        len = layer[1] * layer[2]
        push!(W, reshape(β[idx:idx+len-1], layer))
        idx += len
        push!(b, β[idx:idx+layer[2]-1])
        idx += layer[2]
    end
    return W, b
end

function training_model(x, β)
    W, b = vector_to_params(β, param_sizes)
    vec(model(x, W, b))
end

function model(x, W, b)
    W[3]' * tanh.(W[2]' * tanh.(W[1]' * x .+ b[1]) .+ b[2]) .+ b[3]
end

###############################
# Training
###############################
param_sizes = [(2, 3), (3, 4), (4, 1)]
X, Y, frequencies = get_training_data(2500)
W, b = get_params(param_sizes)
β = params_to_vector(W, b)

using LsqFit
fit = curve_fit(training_model, X, Y, β)
β = coef(fit)

W, b = vector_to_params(β, param_sizes)
BSON.@save "Models/Params.bson" W b

###############################
# EVALUATION
###############################
X, Y, frequencies = get_training_data(2500)
sort_order = sortperm(frequencies)
f_sort = frequencies[sort_order]
X_sort = X[:, sort_order]
Y_sort = Y[sort_order]
guesses = vec(model(X_sort, W, b))

err = Y_sort .- guesses
avg_err = sum(err) / length(err)
maximum(err)

using GLMakie
lines(err)















