from __future__ import print_function

from pyAudioAnalysis import audioBasicIO
import matplotlib.pyplot as plt
from pyAudioAnalysis import ShortTermFeatures as sF
from pyAudioAnalysis import MidTermFeatures as aF
from pyAudioAnalysis import audioTrainTest as aT
from pyAudioAnalysis import audioSegmentation as aS
from pyAudioAnalysis import audioVisualization as aV
import scipy.io.wavfile as wavfile
import matplotlib.patches

name = "umbraid2"
sbp = 0.4285714286

inputFile = "./data/"+name+".wav"

#[Fs, x] = audioBasicIO.read_audio_file(inputFile)
#F, f_names = ShortTermFeatures.feature_extraction(x[:, 0], Fs, 0.4166666*Fs, 0.5*0.416666*Fs)
#F, f_names = ShortTermFeatures.chroma_features(x[:, 0], Fs, 64)

#audioAnalysis.thumbnailWrapper(inputFile, 3.3333 * 4)

[fs, x] = audioBasicIO.read_audio_file(inputFile)
thumbnailWrapperSize = sbp*32

#plt.figure(figsize=(11,10))

st_window = sbp*8
st_step = sbp*8

[A1, A2, B1, B2, Smatrix] = aS.music_thumbnailing(x[:, 0], fs, st_window, st_step,
                                                  thumbnailWrapperSize)

# write thumbnailWrappers to WAV files:
if inputFile.endswith(".wav"):
    thumbnailWrapperFileName1 = inputFile.replace(".wav", "_thumb1.wav")
    thumbnailWrapperFileName2 = inputFile.replace(".wav", "_thumb2.wav")
if inputFile.endswith(".mp3"):
    thumbnailWrapperFileName1 = inputFile.replace(".mp3", "_thumb1.mp3")
    thumbnailWrapperFileName2 = inputFile.replace(".mp3", "_thumb2.mp3")
wavfile.write(thumbnailWrapperFileName1, fs, x[int(fs * A1):int(fs * A2)])
wavfile.write(thumbnailWrapperFileName2, fs, x[int(fs * B1):int(fs * B2)])
print("1st thumbnailWrapper (stored in file {0:s}): {1:4.1f}sec" \
      " -- {2:4.1f}sec".format(thumbnailWrapperFileName1, A1, A2))
print("2nd thumbnailWrapper (stored in file {0:s}): {1:4.1f}sec" \
      " -- {2:4.1f}sec".format(thumbnailWrapperFileName2, B1, B2))


# Plot self-similarity matrix:
fig = plt.figure(figsize=(11,10))

ax = fig.add_subplot(111, aspect="auto")
plt.imshow(Smatrix)
# Plot best-similarity diagonal:
Xcenter = (A1 / st_step + A2 / st_step) / 2.0
Ycenter = (B1 / st_step + B2 / st_step) / 2.0

e1 = matplotlib.patches.Ellipse((Ycenter, Xcenter),
                                thumbnailWrapperSize * 1.4, 3, angle=45,
                                linewidth=3, fill=False)
ax.add_patch(e1)

plt.plot([B1/ st_step, Smatrix.shape[0]], [A1/ st_step, A1/ st_step], color="k",
         linestyle="--", linewidth=2)
plt.plot([B2/ st_step, Smatrix.shape[0]], [A2/ st_step, A2/ st_step], color="k",
         linestyle="--", linewidth=2)
plt.plot([B1/ st_step, B1/ st_step], [A1/ st_step, Smatrix.shape[0]], color="k",
         linestyle="--", linewidth=2)
plt.plot([B2/ st_step, B2/ st_step], [A2/ st_step, Smatrix.shape[0]], color="k",
         linestyle="--", linewidth=2)

plt.xlim([0, Smatrix.shape[0]])
plt.ylim([Smatrix.shape[1], 0])

ax.yaxis.set_label_position("right")
ax.yaxis.tick_right()

plt.xlabel("frame no")
plt.ylabel("frame no")
plt.title("Self-similarity matrix")

plt.imsave('./plots/'+name+'_self_sim_16_raw.png', Smatrix)
plt.savefig("./plots/"+name+"_self_sim_16.png")
plt.show()

#plt.savefig("./plots/born_self_sim.png")

"""
plt.figure(figsize=(19.2,10.8))

c_names, c_features = ShortTermFeatures.chroma_features(x[:, 0], Fs, 0.5*0.4166666*Fs)
#    plt.plot(F[i, :])
#    plt.xlabel('Frame no')
#    plt.ylabel(f_names[i])
"""
#plt.tight_layout()
#plt.show()