import librosa
from matplotlib import pyplot as plt
import numpy as np
from models.utils import count_parameters
from models.ae import SubSpecCAE, AE, GlobalCAE
from cfg import ALL_DATASET_PATHS
import librosa.display
import librosa.feature
import seaborn as sns

purple = (126/255, 87/255, 194/255)
params_subspecae = count_parameters(SubSpecCAE())
params_ae = count_parameters(AE())
params_globalcae = count_parameters(GlobalCAE())
print(f'#Parameters SubSpecCAE: {params_subspecae}')
print(f'#Parameters AE: {params_ae}')
print(f'#SubSpecAe/#AE: {(params_subspecae/params_ae):.2f}')
print(f'#GlobalCAE: {params_globalcae}')
print(f'#GlobalCAE/#SubSpecCAE: {params_globalcae/params_subspecae}')


fig = plt.figure(figsize=(10, 5))
slider_normal = '../data/mimii/-6_dB_slider/id_04/normal/00000042.wav'
audio, sr = librosa.load(slider_normal)
waveplot = librosa.display.waveplot(audio, sr=sr)
plt.tight_layout()
plt.savefig('wavplot.png')

ids = range(0)
specs = []
i = 0
audios = list((ALL_DATASET_PATHS[1] / 'id_00' / 'normal').glob('*.wav'))
print(str(audios[80]))
audio, sr = librosa.load(str(audios[80]), sr=None)
mel_spec = librosa.feature.melspectrogram(audio, sr=sr, n_fft=1024, hop_length=256, n_mels=80, center=False, power=2.0)
mel_spec_db = librosa.power_to_db(mel_spec, ref=np.max)
librosa.display.specshow(mel_spec_db, hop_length=256, x_axis='s', y_axis='mel', cmap='viridis', sr=sr)
plt.savefig('notadjusted.png')
plt.clf()

centroids = []
for p in audios:
    audio, sr = librosa.load(str(p), sr=None)
    spectral_centroids = librosa.feature.spectral_centroid(audio, sr=sr)[0]
    centroids += spectral_centroids.tolist()

sns.distplot(centroids, hist=True, kde=True,
             color=purple,
             hist_kws={'edgecolor':'black'},
             kde_kws={'linewidth': 2})
plt.xlabel('Occurence counts')
plt.ylabel('Density')
plt.title('Spectral centroid distribution')
plt.tight_layout()
plt.savefig('hist.png')

def get_bands(centroids, n_mels=80):
    std = np.std(centroids)
    mu = np.mean(centroids)
    return int((mu-3*std)/n_mels), int((mu+3*std)/n_mels)

print(get_bands(centroids))