#
# File: tone_generator.py
# Author: Alex Spataru <https://github.com/alex-spataru>
# Date: 2023-Oct-12
#
# Description:
#   Python script to generate WAV "beep" tones with different 
#   musical notes & octaves, includes a simple ADSR envelope 
#   and a fade in / fade out filter to avoid blowing up 
#   your innocent laptop speakers.
#

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

# Directory where the generated WAV files will be saved
OUTPUT_DIR = 'beeps'

# List of musical notes for which beeps will be generated
NOTES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]

def generate_tone(frequency, duration_ms=50, sample_rate=44100):
    """
    Generate a tone with ADSR envelope

    @param frequency Frequency of the tone in Hz
    @param duration_ms Duration of the tone in milliseconds
    @param sample_rate Sample rate for generating the tone (default is 44.1 kHz)
    @return NumPy array containing the tone samples
    """

    # Generate sine wave for the whole audio
    t = np.linspace(0, duration_ms / 1000.0, int(sample_rate * (duration_ms / 1000.0)), endpoint=False)
    wave = np.sin(2 * np.pi * frequency * t)

    # Apply ADSR envelope
    total_samples = len(wave)
    attack_samples = int(total_samples * 0.1)
    decay_samples = int(total_samples * 0.1)
    sustain_samples = int(total_samples * 0.5)
    release_samples = total_samples - (attack_samples + decay_samples + sustain_samples)

    # Construct ADSR envelope
    attack = np.linspace(0, 1, attack_samples)
    decay = np.linspace(1, 0.7, decay_samples)
    sustain = np.ones(sustain_samples) * 0.7
    release = np.linspace(0.7, 0, release_samples)

    # Apply the envelope to the sine wave
    envelope = np.concatenate([attack, decay, sustain, release])
    return wave * envelope

def calculate_frequency(note, octave):
    """
    Calculate the frequency for a given note and octave

    @param note The note for which to calculate the frequency (e.g., "A", "C#", etc.)
    @param octave The octave of the note
    @return The frequency of the note in Hz
    """

    reference_note = "A"
    reference_frequency = 440.0
    reference_octave = 4

    note_number = NOTES.index(note)
    reference_note_number = NOTES.index(reference_note)

    frequency = reference_frequency * 2 ** (octave - reference_octave + (note_number - reference_note_number) / 12)
    return frequency

# Generate WAV files for each note in each octave
if __name__ == '__main__':
    for octave in range(1, 5):
        os.makedirs(OUTPUT_DIR, exist_ok=True)
        for note in NOTES:
            frequency = calculate_frequency(note, octave)
            tone = generate_tone(frequency)
            wavfile.write(f"{OUTPUT_DIR}/{note}_{octave}.wav", 44100, np.int16(tone * 32767))