using PyPlot
using NtToolBox
using WAV
# using Autoreload
# reload("NtToolBox")

n = 1024*16
s = 3 #number of sounds
p = 2 #number of micros

x = zeros(n,3)
x[:,1] = load_sound("NtToolbox/src/data/bird.wav",n)
x[:,2] = load_sound("NtToolbox/src/data/female.wav",n)
x[:,3] = load_sound("NtToolbox/src/data/male.wav",n);

x = x./repeat(std(x,1), outer=(n,1));

theta = Array(linspace(0, pi, s + 1)); theta = theta[1:3]
theta[1] = 0.2
M = vcat(cos(theta)', sin(theta)');

y = x*M';

figure(figsize = (10,10))

for i in 1:s
    subplot(s, 1, i)
    plot(x[:, i])
    xlim(0,n)
    title("Source #$i")
end

figure(figsize = (10,7))
    
for i in 1:p
    subplot(p, 1, i)
    plot(y[:, i])
    xlim(0,n)
    title("Micro #$i")
end

w = 128   #size of the window
q = Base.div(w,4);  #overlap of the window

X = complex(zeros(w,4*w+1,s))
Y = complex(zeros(w,4*w+1,p))

for i in 1:s
    X[:,:,i] = perform_stft(x[:,i],w,q,n)
    figure(figsize = (15,10))
    plot_spectrogram(X[:,:,i],"Source #$i")
end

#run -i nt_solutions/audio_2_separation/exo1
include("NtSolutions/audio_2_separation/exo1.jl")

## Insert your code here.

mf = size(Y)[1]
mt = size(Y)[2]
P = reshape(Y, (mt*mf,p))
P = vcat(real(P), imag(P));

npts = 6000;

sel = randperm(n)

sel = sel[1:npts]

figure(figsize = (7,5))
plot(y[sel,1], y[sel,2], ".", ms = 3)
xlim(-5,5)
ylim(-5,5)
title("Time domain");

include("NtSolutions/audio_2_separation/exo2.jl");

## Insert your code here.

nrow = size(P)[1]
Theta = zeros(nrow)
for i in 1:nrow
    Theta[i] = mod(atan2(P[i,2],P[i,1]),pi)
end

nbins = 100
h,t = plt[:hist](Theta,nbins)
h=h/sum(h)
clf()
bar(t[1:end-1], h, width = pi/nbins)
xlim(0,pi);

include("NtSolutions/audio_2_separation/exo3.jl");

## Insert your code here.

include("NtSolutions/audio_2_separation/exo4.jl")

## Insert your code here.

A = reshape(Y, (mt*mf,p));

C = abs(M1'*A');

tmp, I = compute_max(C,1)
I = reshape(I, (mf,mt));

T = .05
D = sqrt(sum(abs(Y).^2, 3))[:,:,1]
I = I.*(D .> T);

figure(figsize = (15,10))
imageplot(I[1:Base.div(mf,2),:])
imshow(I[1:Base.div(mf,2),:], cmap = get_cmap("jet"), interpolation = "nearest");

Proj = M1'*A'
Xr = complex(zeros(w,4*w+1,s))
for i in 1:s
    Xr[:,:,i] = reshape(Proj[i,:], (mf,mt)).*(I .== i)
end

xr = zeros(n,s)
for i in 1:s
    xr[:,i] = perform_stft(Xr[:,:,i], w, q, n)
end

figure(figsize = (10,10))

for i in 1:s
    subplot(s,1,i)
    plot(xr[:,i])
    xlim(0,n)
    title("Estimated source #$i")
end

i = 1
WAV.wavplay(x[:,i], 15000) # Supported back-ends : AudioQueue (MacOSX) and Pulse Audio (Linux, libpulse-simple). 
#There is not a native backend for Windows yet.

WAV.wavplay(xr[:,i], 15000)
<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>