using DSP, FFTW


"""
    function especplot(t, x, Mx, R, Nx, Ny; col="")
Desenha o gráfico de um conjunto de pontos `(t, x)` sobre um espectograma com `Nx` pontos na horizontal e `Ny` pontos na vertical, e frequência de amostragem normalizada ('t' inteiro), assumindo que o primeiro ponto na horizontal corresponda ao instante (Mx).
A escala horizontal corresponde a:
`t = Mx` => `1`; `t = (R*(Nx-1)+Mx)` => `Nx`
Na vertical, assume-se que
`x = 0` => `1`
`x = 1` => `Ny`
`col = ""` permite escolher o estilo do gráfico.
"""
function especplot(t, x, Mx, R, Nx, Ny; col="")
    plot((t .- Mx) * (Nx-1)/((R*(Nx-1))) .+ 1, Ny * x, col)
end

"""
    function especplot(t, x, Mx, R, Nx, Ny, fa; col="")
Desenha o gráfico de um conjunto de pontos `(t, x)` sobre um espectograma com `Nx` pontos na horizontal e `Ny` pontos na vertical, e frequência de amostragem `fa`, assumindo que o primeiro ponto na horizontal corresponda ao instante (Mx/fa).
A escala horizontal corresponde a:
`t = Mx/fa` => `1`; `t = (R*(Nx-1)+Mx)/fa` => `Nx`
Na vertical, assume-se que
`x = 0` => `1`
`x = fa/2` => `Ny`
`col = ""` permite escolher o estilo do gráfico.
"""
function especplot(t, x, Mx, R, Nx, Ny, fa; col="")
    plot((t .- Mx/fa) * (Nx-1)/((R*(Nx-1))/fa) .+ 1, Ny * x * 2 / fa, col)
end


"""
    function plota_stft(X, N; scale = "linear")
Desenha o espectrograma com as escalas de tempo e frequência corretas.
`X` deve ser um espectrograma, como calculado pelo pacote DSP.
`N` é o número de pontos da FFT.
Nesta opção o tempo é dado em amostras.
Parâmetro opcional:
`scale`: `linear` (escala linear), `dB` (escala em dB)
"""
function plota_stft(X::DSP.Periodograms.Spectrogram{Float64,AbstractFFTs.Frequencies{Float64}}, N; scale = "linear")
    Nyticks = 8 # número de marcações na vertical
    yrange = (mod(N,2)==0 ? N÷2+1 : (N+1)÷2)
    ytickpos = 0:yrange/Nyticks:yrange # posições das marcações na vertical
    yticklabels = ytickpos*freq(p)[end]/yrange
    Nxticks = 8
    xtickpos = 1:length(time(p))÷Nxticks:length(time(p))
    xticklabels =  (xtickpos .- 1) * fa*(time(p)[end]-time(p)[1])/(length(time(p))-1) .+ fa*time(p)[1]
    # Multiplico por fa para mostrar os resultados em número de amostras
    # Retirar a multiplicação por fa se o eixo horizontal for tempo mesmo
    if scale == "linear"
        ppow = power(p)
    else
        ppow = pow2db.(power(p))
    end
    ax = pcolormesh(ppow)
    yticks(round.(Int,ytickpos), yticklabels);
    xticks(xtickpos, round.(Int, xticklabels));
    return ax
end
"""
    function plota_stft(X, N, fa; scale = "linear")
Desenha o espectrograma com as escalas de tempo e frequência corretas.
`X` deve ser um espectrograma, como calculado pelo pacote DSP.
`N` é o número de pontos da FFT.
`fa` é a frequência de amostragem - neste caso, o tempo é dado em segundos, e a frequência, em Hz.
Parâmetro opcional:
`scale`: `linear` (escala linear), `dB` (escala em dB)
"""
function plota_stft(X::DSP.Periodograms.Spectrogram{Float64,AbstractFFTs.Frequencies{Float64}}, N, fa; scale = "linear")
    Nyticks = 8 # número de marcações na vertical
    yrange = (mod(N,2)==0 ? N÷2+1 : (N+1)÷2)
    ytickpos = (0:yrange/Nyticks:yrange)*fa/2 # posições das marcações na vertical
    yticklabels = ytickpos*freq(p)[end]/yrange
    Nxticks = 8
    xtickpos = 1:length(time(p))÷Nxticks:length(time(p))
    xticklabels =  (xtickpos .- 1) * (time(p)[end]-time(p)[1])/(length(time(p))-1) .+ time(p)[1]
    # Multiplico por fa para mostrar os resultados em número de amostras
    # Retirar a multiplicação por fa se o eixo horizontal for tempo mesmo
    if scale == "linear"
        ppow = power(p)
    else
        ppow = pow2db.(power(p))
    end
    ax = pcolormesh(ppow)
    yticks(round.(Int,ytickpos), yticklabels);
    xticks(xtickpos, xticklabels);
    return ax
end


"""
    function plota_stft(X, M, R; scale = "linear")
Desenha o espectrograma com as escalas de tempo e frequência corretas.
`X` deve ser uma STFT (matriz), como calculado pelo pacote DSP. Assume-se que só as frequências de 0 a π são fornecidas.
`M` é o comprimento da janela.
`R` é o passo no tempo.
Neste caso, o tempo é dado em amostras, e a frequência, de 0 a 1.
Parâmetro opcional:
`scale`: `linear` (escala linear), `dB` (escala em dB)
"""
function plota_stft(X, M, R; scale = "linear")
    N,Np = size(X) # N = número de pontos da FFT
    Nyticks = 8 # número de marcações na vertical
    ytickpos = (0:(N-1)/Nyticks:N-1) # posições das marcações na vertical
    yticklabels = ytickpos/(N-1)
    Nxticks = 8
    xtickpos = 1:Np÷Nxticks:Np
    xticklabels =  round.(Int, (xtickpos .- 1) * R .+ (M-1)/2)
    if scale == "linear"
        ppow = abs.(X)
    else
        ppow = pow2db.(abs.(X))
    end
    ax=pcolormesh(ppow);
    yticks(round.(Int,ytickpos), yticklabels);
    xticks(xtickpos, xticklabels);
    return ax
end


"""
    function plota_stft(X, M, R, fa; scale = "linear)
Desenha o espectrograma com as escalas de tempo e frequência corretas.
`X` deve ser uma STFT (matriz), como calculado pelo pacote DSP. Assume-se que só as frequências de 0 a π são fornecidas.
`M` é o comprimento da janela.
`R` é o passo no tempo.
`fa` é a frequência de amostragem - neste caso, o tempo é dado em segundos, e a frequência, em Hz.
Parâmetro opcional:
`scale`: `linear` (escala linear), `dB` (escala em dB)
"""
function plota_stft(X, M, R, fa; scale = "linear")
    N,Np = size(X) # N = número de pontos da FFT
    Nyticks = 8 # número de marcações na vertical
    ytickpos = (0:(N-1)/Nyticks:N-1) # posições das marcações na vertical
    yticklabels = (fa/2)*ytickpos/(N-1)
    Nxticks = 8
    xtickpos = 1:Np÷Nxticks:Np
    xticklabels =  (xtickpos .- 1) * R/fa .+ (M-1)/(2fa)
    if scale == "linear"
        ppow = abs.(X)
    else
        ppow = pow2db.(abs.(X))
    end
    ax=pcolormesh(ppow);
    yticks(round.(Int,ytickpos), yticklabels);
    xticks(xtickpos, xticklabels);
    return ax
end