import numpy as np
import pandas as pd
import scipy.io.wavfile as wav
import sklearn as skl
import matplotlib.pyplot as plt
import librosa as lb
from IPython.display import Audio as ipd
from IPython.core.display import display
import random
import cv2 as cv
import time

audio_data_path = 'audio_data/genres_original'
#file path is genre/genre.000{n}.wav where n is a 2 digit number from 00 to 99
ex_file_path =  audio_data_path + '/jazz/jazz.00000.wav'
#returns a tuple (sample rate of file, numpy array of audio file)
tp = wav.read(ex_file_path)
print(tp)
print(tp[1].shape)
print('\n')
# get the audio variable and 'rate' so that you can play it
audio, rate = lb.load(ex_file_path)
# display audio player so that you can listen to the audio
ipd(audio, rate=rate)

# this cell will take about 5-10 minutes to run initially

# all genres that we are working with
GENRES = ['blues', 'classical', 'country', 'disco', 'hiphop', 
          'jazz', 'metal', 'pop', 'reggae', 'rock']

# an arbitrary cutoff value so that all of our song data is the same size
CUTOFF = 660000

# the following are parallel arrays
# file names of the song
filenames = []
# the sample rate of each song
songrates = []
#the numpy nd array encoding of each song
songs = []

for i, genre in enumerate(GENRES): # for each genre
    for j in range(100): #for each song 
        # get the song file address
        idx = str(j) if j > 9 else ('0' + str(j));
        file_path = audio_data_path + '/'+ genre +'/'+ genre +'.000'+ idx + '.wav'
        # print(file_path)
        try: #try to load the array into 
            temp = wav.read(file_path)
            filenames.append(file_path)
            songrates.append(temp[0])
            songs.append(temp[1][0:CUTOFF])
        except ValueError:
            # skip this iteration if error
            # value error means that the audio file is corrupted
            idx = str(j) if j > 9 else ('0' + str(j));
            file_path = './'+ genre +'/'+ genre +'.000'+ idx + '.wav'
            print('audio file '+ file_path +' is corrupted')

songs = np.asarray(songs)
songs[0]

# save the invalid address just in case
# jazz 00000 is index 600, jazz 00054 is index 654 
INVALID_SONG_INDEX = 654

songs.shape

songrates = np.asarray(songrates)
songrates[songrates != 22050]

SONGRATE = 22050

## function to play a song given an index or by nd array depending on input
# usage 
#   play(song=songs[101], rate=songrates[101])
# or
#   play(idx=101)

def play(idx=0, song=None, rate=0):
    if (song is None or rate <= 0):
        audio, rate = lb.load(filenames[idx])
        display(ipd(audio, rate=songrates[idx]))
    else:
        wav.write('audio_data/_outputs/temp.wav', rate, song)
        audio, rate = lb.load('audio_data/_outputs/temp.wav')
        display(ipd(audio, rate=rate))

# function to get the genre of an index of song array
# makes use of the parallel arrays
# usage
#   index, name = get_genre(0)

def get_genre(idx):
    directories = filenames[idx].split('/')
    return directories[3]

mean_song = np.mean(songs, axis=0)
mean_rate = int(np.mean(songrates))

# its white noise 🤪
play(song=mean_song, rate=SONGRATE)

import numpy as np
import scipy.io as sio
import matplotlib
import matplotlib.pyplot as plt
from numpy.matlib import repmat
from sklearn.preprocessing import normalize

def eigsort(V, eigvals):
    lohival = np.sort(eigvals)
    lohiindex = np.argsort(eigvals)
    lambd = np.flip(lohival)
    index = np.flip(lohiindex)
    Dsort = np.diag(lambd)
    M = np.size(lambd)
    Vsort = np.zeros((M, M))
    for i in range(M):
        Vsort[:,i] = V[:,index[i]]
    return Vsort, Dsort

def normc(Mat):
    return normalize(Mat, norm='l2', axis=0)

Z = np.subtract(songs, mean_song)
N = songs.shape[0]

# Trying to reduce RAM usage, not storing the covariance
# eigvals, eigvects = np.linalg.eig(covariance) 
eigvals, eigvects = np.linalg.eig((1/N) * Z @ Z.T) 
eigvects, eigvals = eigsort(eigvects, eigvals)

eigvects.shape

eigvals.shape

plt.plot(range(len(eigvals)), np.diagonal(eigvals), '-o')
plt.title("Plot of Eigenvalues \n")

eigvals_diagonal = np.diagonal(eigvals);
differences = []
for i in range(len(eigvals_diagonal) - 1):
    differences.append(eigvals_diagonal[i] - eigvals_diagonal[i + 1])

#  814 or 948
ax = plt.plot(range(50), differences[0:50], '-o')
plt.title("Differences in Eigenvalue [i] Compared to [i+1] \n")

eigvects.shape

U = Z.T @ eigvects

U = normc(U)
U.shape

# mapped and reduced song matrix
C = U.T @ (Z.T)

C.shape

# reconstruct reduced song (Choose C[i], C5[i], C2[i])
reduced_song = (U @ C[2]).T + mean_song
reduced_song.shape

play(song=(Z[2]), rate=SONGRATE)

play(song=reduced_song, rate=SONGRATE)

plt.plot(C[0], C[1], 'o')
plt.title("Song Data in the First Two Principal Component Directions \n")

# this code exists to check outliers, 

# # trying to remove outliers
pca_mean = np.median(C, axis=1)# np.mean(C, axis = 1)

# #compute each point's distance from the mean
dists = []
dists_outliers = []
indecies = []
i = 0
for pt in C.T:
    dist = np.linalg.norm(pt - pca_mean) 
    if(1e18 < dist):
        indecies.append(i)
    else:
        dists.append(dist)
    i+=1
    dists_outliers.append(dist)

print("outliers: ")
print(indecies)

# this code exists to check outliers, but they were removed with hard code in 
# the Preprocessing step earlier

plt.boxplot(dists_outliers)
plt.title("Distribution of Each Point's Distance from the Median \n")

# this code exists to check outliers, but they were removed with hard code in 
# the Preprocessing step earlier

plt.boxplot(dists)
plt.title("Distribution of Each Point's Distance from the Median -- No Outliers \n")

FEATURE_PATH = 'audio_data/features_30_sec.csv'

df = pd.read_csv(FEATURE_PATH, sep=',', header=0)
songs = songs.astype('float64')

mfcc = lb.feature.mfcc(y=songs[0], sr=SONGRATE)
MFCC_SAMPLE_AMT = 1290
mfcc.shape

song_from_mfcc = lb.feature.inverse.mfcc_to_audio(mfcc)

print(song_from_mfcc.shape)
plt.plot(range(len(song_from_mfcc)), song_from_mfcc)

print(songs[0].shape)
plt.plot(range(len(songs[0])), songs[0])

play(song=songs[0], rate=SONGRATE)

play(song=song_from_mfcc, rate=SONGRATE)

def mfcc_to_timeseries(mfcc):

    if(len(mfcc)%2 != 0):
        print('Not given matching means and variances')
  
    counter = 0
    means = []
    vars = []
    for m in mfcc:
        if(counter%2 == 0):
            means.append(m)
        else:
            vars.append(m)
        counter+=1

    arr = []
    for i in range(len(means)):
        arr.append( np.random.normal(loc=means[i], scale=np.sqrt(vars[i]), size=MFCC_SAMPLE_AMT).tolist() )

    timeseries = lb.feature.inverse.mfcc_to_audio(arr)
                                          
    return timeseries

df.head()

features = df.drop(['filename', 'length', 'label'], axis = 1).values

song_from_distributions = mfcc_to_timeseries(features[0][17:57])

play(song=songs[0], rate=SONGRATE)

play(song=song_from_mfcc, rate=SONGRATE)

play(song=song_from_distributions, rate=SONGRATE)

features = df.drop(['filename', 'length', 'label'], axis=1).values

pca = skl.decomposition.PCA(n_components = 10)
pca.fit(features)

pca.explained_variance_ratio_

plt.plot(range(len(pca.explained_variance_ratio_)), pca.explained_variance_ratio_, '-o')

pca = skl.decomposition.PCA(n_components = 2)

reduced_features = pca.fit_transform(features)

reduced_features.shape

plt.plot(reduced_features.T[0], reduced_features.T[1], 'o')
plt.xlabel("PCA Dimesnion 1")
plt.ylabel("PCA Dimension 2")
plt.title("Our Feature Data in 2 Principal Component Dimensions \n")

reconstructed_features = pca.inverse_transform(reduced_features)

reconstructed_features.shape

def calcSqDistances(X, Kmus):
    N, D = X.shape
    K = Kmus.shape[0]
    X_tile = np.tile(X, (K,1,1))
    X_tile = np.transpose(X_tile, axes=[1,0,2])
    differences = X_tile - Kmus
    sqDmat = np.linalg.norm(differences, axis=-1)
    return sqDmat

def determineRnk(sqDmat):
    N, K = sqDmat.shape
    Rnk = np.zeros((N, K))

    for i in range(0,N):
        row = sqDmat[i]
        Rnk[i, np.argmin(row)] = 1
    
    return Rnk

def recalcMus(X, Rnk):
    sum = (1/np.sum(Rnk.T, axis=1))
    return (sum*np.dot(Rnk.T, X).T).T

def runKMeans(K, X):
    N, D = X.shape

    Kmus = np.zeros((K, D))

    rand_inds = np.random.permutation(N)
    Kmus = X[rand_inds[0:K],:]

    maxiters = 1000

    for iter in range(maxiters):

        sqDmat = calcSqDistances(X, Kmus)

        Rnk = determineRnk(sqDmat)

        KmusOld = Kmus

        Kmus = recalcMus(X, Rnk)

        if np.sum(np.abs(KmusOld.reshape((-1, 1)) - Kmus.reshape((-1, 1)))) < 1e-6:
            print("converged on iteration "+ str(iter))
            # print(Kmus)
            break

    return Kmus, Rnk

K_CLUSTERS = 10

Kmus, Rnk = runKMeans(K_CLUSTERS, reduced_features)

spectro_labels = []
for x in range(1000):
    index = np.argwhere(Rnk[x])[0][0]
    spectro_labels.append(index)
spectro_labels = np.asarray(spectro_labels)

df['kmeans_label'] = spectro_labels

df_plot_kmeans = df.groupby(['label', 'kmeans_label']).size().reset_index().pivot(columns='label', index='kmeans_label', values=0)

ax = df_plot_kmeans.plot(kind='bar', stacked=True)

plt.legend(title='Genre', bbox_to_anchor=(1.05, 1), loc='upper center')
ax.figure.set_size_inches(18,10)
plt.ylabel("Number of Songs")
plt.xlabel("KMeans Cluster Group No.")
plt.title("KMeans Song Groupings Colored by Genre")

kmeans = skl.cluster.KMeans(n_clusters=K_CLUSTERS)
kmeans.fit(features)

df["kmeans_label_sklearn_features"] = kmeans.labels_

df_plot_kmeans = df.groupby(['label', 'kmeans_label_sklearn_features']).size().reset_index().pivot(columns='label', index='kmeans_label_sklearn_features', values=0)

ax = df_plot_kmeans.plot(kind='bar', stacked=True)

plt.legend(title='Genre', bbox_to_anchor=(1.05, 1), loc='upper center')
ax.figure.set_size_inches(18,8)

plt.ylabel("Amount of Songs")
plt.xlabel("KMeans Cluster Group No.")
plt.title("KMeans Song Groupings Colored by Genre -- SKLearn")

spectogram_data_path = "audio_data/images_original"

images = []

GENRES = ['blues', 'classical', 'country', 'disco', 'hiphop', 
          'jazz', 'metal', 'pop', 'reggae', 'rock']

for i, genre in enumerate(GENRES): # for each genre
    for j in range(100): #for each song 
        # get the song file address
        idx = str(j) if j > 9 else ('0' + str(j));
        file_path = spectogram_data_path + '/'+ genre +'/'+ genre +'000'+ idx + '.png'
        # print(file_path)
        try: #try to load the array into 
            temp = cv.cvtColor(cv.imread(file_path), cv.COLOR_BGR2GRAY)[35:-35, 54:-42]
            images.append(temp)

        except Exception:
            # skip this iteration if error
            # value error means that the audio file is corrupted
            idx = str(j) if j > 9 else ('0' + str(j));
            file_path = './'+ genre +'/'+ genre +'.000'+ idx + '.png'
            # print('audio file '+ file_path +' is corrupted')

images = np.array(images)
images = images.reshape((999, -1))/255
images.shape

from sklearn.preprocessing import normalize

def eigsort(V, eigvals):
    
    # Sort the eigenvalues from largest to smallest. Store the sorted
    # eigenvalues in the column vector lambd.
    lohival = np.sort(eigvals)
    lohiindex = np.argsort(eigvals)
    lambd = np.flip(lohival)
    index = np.flip(lohiindex)
    Dsort = np.diag(lambd)
    
    # Sort eigenvectors to correspond to the ordered eigenvalues. Store sorted
    # eigenvectors as columns of the matrix vsort.
    M = np.size(lambd)
    Vsort = np.zeros((M, M))
    for i in range(M):
        Vsort[:,i] = V[:,index[i]]
    return Vsort, Dsort

def normc(Mat):
    return normalize(Mat, norm='l2', axis=0)

image_mean = np.mean(images, 0).reshape((1, 73248))
print(image_mean.shape)

plt.imshow(image_mean.reshape([218,336]))
plt.title("Spectogram Mean")
plt.imsave("image_mean.png", image_mean.reshape([218,336]))

images_minus_mean = images - image_mean
print(images_minus_mean.shape)

w, v = np.linalg.eig(images_minus_mean @ images_minus_mean.T)
V, W = eigsort(v, w)
print(V.shape)
U = images_minus_mean.T @ V
print(U.shape)

plt.scatter(U[0], U[1])
plt.xlabel("Principal Component 1")
plt.ylabel("Principal Component 2")
plt.title("Spectogram data reduced to 2 principal components")

U = normc(U)

plt.imshow(U[:, 0].reshape([218,336]))

compoenent_space = U.T @ images_minus_mean.T

compoenent_space.shape

Z = (U @ compoenent_space)
Z = Z.T + image_mean

#Reconstuction of images using all 999 components

Z_100 = (U[:, :100] @ compoenent_space[:100, :])
Z_100 = Z_100.T + image_mean

#Reconstruction of images usin first 100 components 

plt.imshow((Z[12]).reshape([218,336]))
plt.title("Image using all components")

plt.imshow((Z_100[12]).reshape([218,336]))
plt.title("Image using 100 components")

plt.imshow(images[12].reshape([218,336]))
plt.title("Original Image")

S = np.diagonal(W)
plt.plot(np.arange(100),S[:100])
plt.xlabel("Principal Component")
plt.ylabel("Variance")

'''
Vectorized calcSqDistances
'''
def calcSqDistances(X, Kmus):
    N, D = X.shape
    K = Kmus.shape[0]
    X_tile = np.tile(X, (K,1,1))
    X_tile = np.transpose(X_tile, axes=[1,0,2])
    differences = X_tile - Kmus
    sqDmat = np.linalg.norm(differences, axis=-1)

    return sqDmat

def determineRnk(sqDmat):
    N, K = sqDmat.shape
    Rnk = np.zeros((N, K))

    for i in range(0,N):
        row = sqDmat[i]
        Rnk[i, np.argmin(row)] = 1
    
    return Rnk

def recalcMus(X, Rnk):
    sum = (1/np.sum(Rnk.T, axis=1))
    return (sum*np.dot(Rnk.T, X).T).T

def runKMeans(K, X):
    #determine and store data set information
    N, D = X.shape

    #allocate space for the K mu vectors
    Kmus = np.zeros((K, D))

    #initialize cluster centers by randomly picking points from the data
    rand_inds = np.random.permutation(N)
    Kmus = X[rand_inds[0:K],:]

    #specify the maximum number of iterations to allow
    maxiters = 1000

    for iter in range(maxiters):
        #assign each data vector to closest mu vector as per Bishop (9.2)
        #do this by first calculating a squared distance matrix where the n,k entry
        #contains the squared distance from the nth data vector to the kth mu vector

        #sqDmat will be an N-by-K matrix with the n,k entry as specfied above
        sqDmat = calcSqDistances(X, Kmus)

        #given the matrix of squared distances, determine the closest cluster
        #center for each data vector

        #R is the "responsibility" matrix
        #R will be an N-by-K matrix of binary values whose n,k entry is set as
        #per Bishop (9.2)
        #Specifically, the n,k entry is 1 if point n is closest to cluster k,
        #and is 0 otherwise
        Rnk = determineRnk(sqDmat)

        KmusOld = Kmus

        #recalculate mu values based on cluster assignments as per Bishop (9.4)
        Kmus = recalcMus(X, Rnk)


        #check to see if the cluster centers have converged.  If so, break.
        if np.sum(np.abs(KmusOld.reshape((-1, 1)) - Kmus.reshape((-1, 1)))) < 1e-6:
            print(iter)
            print(Kmus)
            break

    return Kmus, Rnk

Kmus, Rnk = runKMeans(10, Z_100)

index_list = list(range(10))
number_list = [0,0,0,0,0,0,0,0,0,0]
spectro_labels = []
for x in range(999):
    index = np.argwhere(Rnk[x])[0][0]
    spectro_labels.append(index)
    number_list[index] += 1

number_list

spectro_df = pd.DataFrame(data={'labels': df['label'].drop(INVALID_SONG_INDEX), 'kmeans_labels': spectro_labels})

df_plot_kmeans_spectro = spectro_df.groupby(['labels', 'kmeans_labels']).size().reset_index().pivot(columns='labels', index='kmeans_labels', values=0)

ax = df_plot_kmeans_spectro.plot(kind='bar', stacked=True)

plt.legend(title='Genre', bbox_to_anchor=(1.05, 1), loc='upper center')
ax.figure.set_size_inches(18,10)
plt.ylabel("Number of Songs")
plt.xlabel("KMeans Cluster Group No.")
plt.title("KMeans Song Groupings Colored by Genre")