import soundbox
import numpy as np
import scipy.signal as sig
import matplotlib
import matplotlib.pyplot as plt
import sys

if len(sys.argv) != 3:
    print(f"""Utilisation: {sys.argv[0]} [source] [output]

Génère le sonagramme du fichier [source] dans le fichier [output].
Passer - comme [output] fait s’afficher le sonagramme dans une fenêtre.""")
    sys.exit(1)

source_file = sys.argv[1]
output_file = sys.argv[2]

# Calcul du STFT
signal = soundbox.load_signal(source_file)
freq, time, vecs = sig.stft(signal, soundbox.samp_rate, nperseg=soundbox.samp_rate * 0.5)
values = np.absolute(vecs)

# Génération du graphe
plt.rcParams.update({
    'figure.figsize': (10, 8),
    'figure.frameon': True,
    'font.size': 20,
    'font.family': 'Concourse T4',
})

fig, ax = plt.subplots(frameon=True)
ax.tick_params(axis='both', which='major', labelsize=12)

freq_filter = values.max(axis=1) / np.max(values) >= 0.01
x = np.arange(len(values))

ax.pcolormesh(
    time, freq[freq_filter], values[freq_filter],
    cmap='Greys',
    shading='gouraud')

# Configuration des axes
def time_format(value, pos):
    return f'{value:.0f} s'


def freq_format(value, pos):
    return f'{value:.0f} Hz'


ax.set_xlabel('Temps', labelpad=10)
ax.xaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(time_format))

ax.set_ylabel('Fréquence', labelpad=10)
ax.set_yscale('log', base=2)
ax.yaxis.set_major_locator(plt.MultipleLocator(100))
ax.yaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(freq_format))
ax.grid(alpha=.3)

# Rendu du résultat
if output_file == '-':
    plt.show()
else:
    plt.tight_layout()
    plt.savefig(output_file)