import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import librosa
import librosa.display
# set LaTeX font
# ==============
nice_fonts = {
    # Use LaTeX to write all text
    "text.usetex": True,
    "font.family": "serif",
    # Use 10pt font in plots, to match 10pt font in document
    "axes.labelsize": 16,  # default 10
    "font.size": 16,  # default 10
    # Make the legend/label fonts a little smaller
    "legend.fontsize": 14,  # default 8
    "xtick.labelsize": 14,  # default 8
    "ytick.labelsize": 14,  # default 8
}
mpl.rcParams.update(nice_fonts, )
mpl.rcParams['axes.unicode_minus'] = False

# TODO LaTex
fan_normal = '../data/mimii/-6_dB_fan/id_04/normal/00000042.wav'
fan_anomaly = '../data/mimii/-6_dB_fan/id_04/abnormal/00000048.wav'

pump_normal = '../data/mimii/-6_dB_pump/id_04/normal/00000042.wav'
pump_anomaly = '../data/mimii/-6_dB_pump/id_04/abnormal/00000042.wav'

slider_normal = '../data/mimii/-6_dB_slider/id_04/normal/00000042.wav'
slider_anomaly = '../data/mimii/-6_dB_slider/id_04/abnormal/00000042.wav'

valve_normal = '../data/mimii/-6_dB_valve/id_04/normal/00000042.wav'
valve_anomaly = '../data/mimii/-6_dB_valve/id_04/abnormal/00000042.wav'

fig = plt.figure(figsize=(17, 5))
for i, (p, title) in enumerate(zip([fan_normal, pump_normal, slider_normal, valve_normal, fan_anomaly, pump_anomaly, slider_anomaly, valve_anomaly],
                          ['Fan', 'Pump', 'Slider', 'Valve', 'Fan', 'Pump', 'Slider', 'Valve'])):
    plt.subplot(2, 4, i+1)
    audio, sr = librosa.load(p, sr=None, mono=True)
    # n_fft=1024, hop_length=256, n_mels=80, center=False, power=2.0
    S = librosa.feature.melspectrogram(y=audio, sr=sr, n_fft=1024, hop_length=256, n_mels=80, center=False, power=2.0)
    S_dB = librosa.power_to_db(S, ref=np.max)
    librosa.display.specshow(S_dB, x_axis='s' if i == 4 else 'off', hop_length=256,
                             y_axis='mel' if i==4  else 'off', sr=16000, cmap='viridis')
    if i < 4:
        plt.title(title)
    else:
        plt.title(title + ' malfunction')

cbar_ax = fig.add_axes([0.94, 0.18, 0.015, 0.7])
cmap = mpl.cm.viridis
norm = mpl.colors.Normalize(vmin=0, vmax=-80)
cb1 = mpl.colorbar.ColorbarBase(cbar_ax, cmap=cmap,
                                norm=norm,
                                orientation='vertical',
                                format='%+2.0f dB')
plt.tight_layout()
fig.subplots_adjust(right=0.93)
plt.savefig('normal_vs_abnormal.png')