using PyPlot

# Bayesian NMF package
using BNMF

# Read test file
using WAV
x, fs = wavread("arayuru.wav")
x = x[:] # monoral
println(length(x)/fs, " sec.")

# Short-Time Fourier Transform
include("stft.jl")

# Compute spectrogram
framelen = 2048
X = stft(x, framelen=framelen, hopsize=256, window=hanning(framelen))
X = abs(X[:,1:framelen/2])

figure(figsize=(16, 6), dpi=80, facecolor="w", edgecolor="k")
imshow(log(X)', origin="lower", aspect="auto")
colorbar()

p = GaPNMF(X, K=100, a=0.1, b=0.1, alpha=1.0, smoothness=100)
typeof(p)

tic()
fit!(p, epochs=100, verbose=true)
toc()

# Get active K-components
good = goodk(p)

goodspec = p.Eh[good,:]

# Visualize K-active spectrum
figure(figsize=(16, 6), dpi=80, facecolor="w", edgecolor="k")

for k=1:size(goodspec,1)
    plot(goodspec[k,:][:], label=string("#", k))
end
legend()

# Reconstruction by goodk
Y = xbar(p, good)
figure(figsize=(16, 6), dpi=80, facecolor="w", edgecolor="k")
imshow(log(Y)', origin="lower", aspect="auto", interpolation="nearest")
colorbar()

# E[templates]
figure(figsize=(16, 12), dpi=80, facecolor="w", edgecolor="k")
for k=1:length(good)
    subplot(length(good), 1, k)
    plot(p.Eh[good[k],:][:])
    xlim(0,1024)
end

# E[activations]
figure(figsize=(16, 12), dpi=80, facecolor="w", edgecolor="k")
for k=1:length(good)
    subplot(length(good), 1, k)
    plot(p.Ew[:,good[k]][:])
    xlim(0, 229)
end

# Only length(good) components are active
figure(figsize=(16, 6), dpi=80, facecolor="w", edgecolor="k")
bar(good, p.Et[good])