from __future__ import division, print_function
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline
import scipy.io.wavfile

sig = np.loadtxt("sound1.dat")
# sound1 = np.asarray((2**16)*sig/(max(sig)-min(sig)), np.int16)
sound1 = sig
scipy.io.wavfile.write("sound1_orig.wav", 8192, sound1)

sig = np.loadtxt("sound2.dat")
sound2 = sig
# sound2 = np.asarray((2**16)*sig/(max(sig)-min(sig)), np.int16)
scipy.io.wavfile.write("sound2_orig.wav", 8192, sound2)

sounds = np.array([sound1, sound2])

"""
Explanation: Machine Intelligence II (week 4) - Team MensaNord

Nikolai Zaki
Alexander Moore
Johannes Rieke
Georg Hoelger
Oliver Atanaszov

Exercise 1
A
End of explanation
"""


A = np.random.random((2, 2))
A_inv = np.linalg.inv(A)
xsounds = np.dot(A, sounds)
# xsounds = mixed sounds
scipy.io.wavfile.write("sound1_mixed.wav", 8192, xsounds[0])
scipy.io.wavfile.write("sound2_mixed.wav", 8192, xsounds[1])

"""
Explanation: B
End of explanation
"""


neworder = np.random.permutation(np.arange(18000)) 
pxsounds = xsounds[:, np.asarray(neworder)]

# pxsounds = permutated mixed sounds
scipy.io.wavfile.write("sound1_mixed_perm.wav", 8192, pxsounds[0])
scipy.io.wavfile.write("sound2_mixed_perm.wav", 8192, pxsounds[1])

"""
Explanation: C
End of explanation
"""


correlation = np.cov(sounds, pxsounds) / np.std(sounds) / np.std(pxsounds)
plt.imshow(correlation, interpolation='none')
plt.title('Correlations of original and mixed sounds',y=1.05)
plt.colorbar();

"""
Explanation: D
End of explanation
"""


# cpxsounds = centered permutated mixed sounds
cpxsounds = (pxsounds.T - np.mean(pxsounds, 1)).T           

# cxsounds = centered mixed sounds
cxsounds = xsounds - np.mean(xsounds,1)[:,np.newaxis]

"""
Explanation: E
End of explanation
"""


def ddf_by_df(inp):
    return 1-2.0/(1+np.exp(-inp))

W = np.random.random((2, 2))
print("Goal:")
print(A_inv)
print("W start")
print(W)
k = 0
W_store=np.zeros((18,2,2))
eta_0 = .8
for t in range(18000):
    eta = eta_0/(t+1)
    W_inv = np.linalg.inv(W)
    x = cxsounds[:,t]
    gradient = W_inv.T + np.dot(ddf_by_df(np.dot(W,x)).reshape(2,1),x.reshape(1,2))#ddf_by_df(np.dot(W,x))[:,np.newaxis]*x
    W += eta*gradient
    if t%1000==0:
        W_store[k]=W
        k+=1
print("W end")
print(W)
W_reg_store = W_store
W_regular = W

"""
Explanation: Exercise 2
A
End of explanation
"""


W = np.random.random((2, 2))
print("Goal:")
print(A_inv)
print("W start")
print(W)
eta_0 = .35
k = 0
W_store=np.zeros((18,2,2))
for t in range(18000):
    x = cxsounds[:,t]
    gradient = np.dot(np.eye(2) + np.dot(ddf_by_df(np.dot(W,x)).reshape(2,1),np.dot(W,x).reshape(1,2)),W)#np.dot(ddf_by_df(np.dot(W,x)),np.dot(W,x)[:,np.newaxis]),W)
    eta = eta_0/(t+1)
    W += eta*gradient
    if t%1000==0:
        W_store[k]=W
        k+=1
print("W end")
print(W)
W_nat_store = W_store
W_natural = W

natural_unmixed = np.dot(W_natural, cxsounds)
regular_unmixed = np.dot(W_regular, cxsounds)

fig, axs = plt.subplots(5,2, figsize = (12,8))
for i in range(2):
    
    axs[0][i].plot(sounds[i])    
    axs[1][i].plot(cxsounds[i])
    axs[2][i].plot(cpxsounds[i])
    axs[3][i].plot(natural_unmixed[i])
    axs[4][i].plot(regular_unmixed[i])
    axs[0][i].set_title('original{} '.format(i+1))
    axs[1][i].set_title('mixed{} '.format(i+1))
    axs[2][i].set_title('permutated{} '.format(i+1))
    axs[3][i].set_title('natural{} '.format(i+1))
    axs[4][i].set_title('regular{} '.format(i+1))
    
plt.setp(axs, xlabel = 'time', ylabel = 'amplitude');
fig.tight_layout()
plt.show()

correlation = np.cov(sounds, natural_unmixed) / np.std(sounds) / np.std(natural_unmixed)
plt.imshow(correlation, interpolation='none')
plt.title('Correlation of original and unmixed with natural gradient',y=1.05)
plt.colorbar();

correlation = np.cov(sounds, regular_unmixed) /(np.std(sounds)*np.std(regular_unmixed))
plt.imshow(correlation, interpolation='none')
plt.title('Correlation of original and unmixed with regular gradient',y=1.05)
plt.colorbar();

W_reg_norms = np.linalg.norm(W_reg_store,axis=(2,1))
W_nat_norms = np.linalg.norm(W_nat_store,axis=(2,1))

fig, axs = plt.subplots(1,2,figsize=(12,4))
axs[0].plot(W_reg_norms)
axs[0].set_title('Norms of $W$ in regular gradient descent')
axs[1].plot(W_nat_norms)
axs[1].set_title('Norms of $W$ in natural gradient descent')

white1 = (sound1-np.mean(sound1))/np.std(sound1)
white2 = (sound2-np.mean(sound1))/np.std(sound2)
white = np.vstack((white1,white2))
mixed_white = np.dot(A,white)

def ddf_by_df(inp):
    return 1-2.0/(1+np.exp(-inp))

W = np.random.random((2, 2))
print("Goal:")
print(A_inv)
print("W start")
print(W)
k = 0
wW_store=np.zeros((18,2,2))
eta_0 = .8
for t in range(18000):
    eta = eta_0/(t+1)
    W_inv = np.linalg.inv(W)
    x = mixed_white[:,t]
    gradient = W_inv.T + np.dot(ddf_by_df(np.dot(W,x)).reshape(2,1),x.reshape(1,2))#ddf_by_df(np.dot(W,x))[:,np.newaxis]*x
    W += eta*gradient
    if t%1000==0:
        wW_store[k]=W
        k+=1
print("W end")
print(W)
wW_reg_store = wW_store
wW_regular = W

W = np.random.random((2, 2))
print("Goal:")
print(A_inv)
print("W start")
print(W)
eta_0 = .35
k = 0
wW_store=np.zeros((18,2,2))
for t in range(18000):
    x = mixed_white[:,t]
    gradient = np.dot(np.eye(2) + np.dot(ddf_by_df(np.dot(W,x)).reshape(2,1),np.dot(W,x).reshape(1,2)),W)#np.dot(ddf_by_df(np.dot(W,x)),np.dot(W,x)[:,np.newaxis]),W)
    eta = eta_0/(t+1)
    W += eta*gradient
    if t%1000==0:
        wW_store[k]=W
        k+=1
print("W end")
print(W)
wW_nat_store = wW_store
wW_natural = W

wW_reg_norms = np.linalg.norm(wW_reg_store,axis=(2,1))
wW_nat_norms = np.linalg.norm(wW_nat_store,axis=(2,1))

fig, axs = plt.subplots(1,2,figsize=(12,4))
axs[0].plot(wW_reg_norms)
axs[0].set_title('Norms of $W$ in regular gradient descent on whitened data',y=1.05);
axs[1].plot(wW_nat_norms)
axs[1].set_title('Norms of $W$ in natural gradient descent on whitened data',y=1.05);

"""
Explanation: B
End of explanation
"""


