using Random, LinearAlgebra, Statistics
using TotalLeastSquares, DSP, LTVModels, ControlSystems, LPVSpectral, AlphaStableDistributions, Plots
gr()
T        = 2^14                                 # Number of samples
fs       = 16_000                               # Sample frequency
t        = range(0,step=1/fs, length=T)         # time vector
chirp_f1 = LinRange(2000, 2500, T)              # frequency of chirp
chirp_f2 = 5000 .+ 200sin.(2pi/T .* (1:T))
y        = sin.(2pi .* chirp_f1 .* t )          # clean signal
y      .+= sin.(2pi .* chirp_f2 .* t )       
plot(y[1:100], lab="Noise-free signal")

Random.seed!(0)
e  = rand(AlphaSubGaussian(n=T))
yn = y + e
plot([y yn][1:500,:], lab=["Clean signal" "Noisy signal"], alpha=0.5)

eye(n) = Matrix{Float64}(I(n))
n  = 6          # Order of the LTV model. 6 poles means we can estimate 3 frequencies.
R1 = eye(n)
R2 = [1e5]      # This parameter contrrols the amount of smoothing
P0 = 1e4R1
D  = 2;         # Can be set to 2 for added filter momentum

@time yf = mapwindows(Windows2(yn,1:length(yn),length(yn)÷20)) do (yn,t)
    yf = lowrankfilter(yn,400,lag=8)
end
@time model = KalmanAR(yf,R1,R2,P0,extend=true, printfit=false, D=D);

function rootspectrogram(model, fs)
    roots = map(1:length(model)) do i
        sys = tf(1,[1;-reverse(model.θ[:,i])],1)        # create a transfer function from the coefficients
        p = sort(pole(sys), by=imag, rev=true)[1:end÷2] # calculate poles and throw away those with negative imag
        p = fs/(2pi) .* angle.(p)       # The frequency of a pole is approx its angle times the sample rate
        sort(replace(p, fs÷2=>0))
    end
    S = reduce(hcat,roots)'
    
end

RS = rootspectrogram(model,fs)
S  = spectrogram(y, 2^10, fs=fs, window=hanning)
pt = range(S.time[1],stop=S.time[end], length=size(RS,1))
f1 = plot(S, yscale=:identity, title="Spectrogram clean")

S  = spectrogram(yn, 2^10, fs=fs, window=hanning)
f2 = plot(S, yscale=:identity, title="Spectrogram noisy")
plot!(repeat(pt,1,n÷2)[1:100:end,:], RS[1:100:end,:], lab="", l=(:blue, 0.5))

plot(f1,f2)