# This code is supporting material for the book
# Building Machine Learning Systems with Python
# by Willi Richert and Luis Pedro Coelho
# published by PACKT Publishing
#
# It is made available under the MIT License

import sys
import os
import glob

import numpy as np
import scipy
import scipy.io.wavfile

from utils import GENRE_DIR, CHART_DIR

import matplotlib.pyplot as plt
from matplotlib.ticker import EngFormatter


def write_fft(fft_features, fn):
    """
    Write the FFT features to separate files to speed up processing.
    """
    base_fn, ext = os.path.splitext(fn)
    data_fn = base_fn + ".fft"

    np.save(data_fn, fft_features)
    print("Written "%data_fn)


def create_fft(fn):
    sample_rate, X = scipy.io.wavfile.read(fn)

    fft_features = abs(scipy.fft(X)[:1000])
    write_fft(fft_features, fn)


def read_fft(genre_list, base_dir=GENRE_DIR):
    X = []
    y = []
    for label, genre in enumerate(genre_list):
        genre_dir = os.path.join(base_dir, genre, "*.fft.npy")
        file_list = glob.glob(genre_dir)
        assert(file_list), genre_dir
        for fn in file_list:
            fft_features = np.load(fn)

            X.append(fft_features[:2000])
            y.append(label)

    return np.array(X), np.array(y)


def plot_wav_fft(wav_filename, desc=None):
    plt.clf()
    plt.figure(num=None, figsize=(6, 4))
    sample_rate, X = scipy.io.wavfile.read(wav_filename)
    spectrum = np.fft.fft(X)
    freq = np.fft.fftfreq(len(X), 1.0 / sample_rate)

    plt.subplot(211)
    num_samples = 200.0
    plt.xlim(0, num_samples / sample_rate)
    plt.xlabel("time [s]")
    plt.title(desc or wav_filename)
    plt.plot(np.arange(num_samples) / sample_rate, X[:num_samples])
    plt.grid(True)

    plt.subplot(212)
    plt.xlim(0, 5000)
    plt.xlabel("frequency [Hz]")
    plt.xticks(np.arange(5) * 1000)
    if desc:
        desc = desc.strip()
        fft_desc = desc[0].lower() + desc[1:]
    else:
        fft_desc = wav_filename
    plt.title("FFT of %s" % fft_desc)
    plt.plot(freq, abs(spectrum), linewidth=5)
    plt.grid(True)

    plt.tight_layout()

    rel_filename = os.path.split(wav_filename)[1]
    plt.savefig("%s_wav_fft.png" % os.path.splitext(rel_filename)[0],
                bbox_inches='tight')

    plt.show()


def plot_wav_fft_demo():
    plot_wav_fft("sine_a.wav", "400Hz sine wave")
    plot_wav_fft("sine_b.wav", "3,000Hz sine wave")
    plot_wav_fft("sine_mix.wav", "Mixed sine wave")


def plot_specgram(ax, fn):
    sample_rate, X = scipy.io.wavfile.read(fn)
    ax.specgram(X, Fs=sample_rate, xextent=(0, 30))


def plot_specgrams(base_dir=CHART_DIR):
    """
    Plot a bunch of spectrograms of wav files in different genres
    """
    plt.clf()
    genres = ["classical", "jazz", "country", "pop", "rock", "metal"]
    num_files = 3
    f, axes = plt.subplots(len(genres), num_files)

    for genre_idx, genre in enumerate(genres):
        for idx, fn in enumerate(glob.glob(os.path.join(GENRE_DIR, genre, "*.wav"))):
            if idx == num_files:
                break
            axis = axes[genre_idx, idx]
            axis.yaxis.set_major_formatter(EngFormatter())
            axis.set_title("%s song %i" % (genre, idx + 1))
            plot_specgram(axis, fn)

    specgram_file = os.path.join(base_dir, "Spectrogram_Genres.png")
    plt.savefig(specgram_file, bbox_inches="tight")

    plt.show()


if __name__ == "__main__":
    # for fn in glob.glob(os.path.join(sys.argv[1], "*.wav")):
    #    create_fft(fn)

    # plot_decomp()

    if len(sys.argv) > 1:
        plot_wav_fft(sys.argv[1], desc="some sample song")
    else:
        plot_wav_fft_demo()

    plot_specgrams()