"""Small NumPy Griffin-Lim demo for https://ryel.io/blog/griffin-lim/."""

import argparse
import wave
from pathlib import Path

import numpy as np


def stft(audio, n_fft, win_length, hop_length):
    window = np.hanning(win_length + 1)[:-1]
    left = win_length // 2
    right = win_length - left
    padded = np.pad(audio, (left, right))
    extra = (-(len(padded) - win_length)) % hop_length
    padded = np.pad(padded, (0, extra))
    frames = np.lib.stride_tricks.sliding_window_view(padded, win_length)[::hop_length]
    return np.fft.rfft(frames * window, n=n_fft, axis=-1)


def istft(spectrum, length, n_fft, win_length, hop_length):
    window = np.hanning(win_length + 1)[:-1]
    frames = np.fft.irfft(spectrum, n=n_fft, axis=-1)[:, :win_length] * window
    padded_length = (len(frames) - 1) * hop_length + win_length
    audio = np.zeros(padded_length)
    weight = np.zeros(padded_length)

    for index, frame in enumerate(frames):
        start = index * hop_length
        audio[start:start + win_length] += frame
        weight[start:start + win_length] += window**2

    audio /= np.maximum(weight, 1e-12)
    left = win_length // 2
    return audio[left:left + length]


def griffin_lim(magnitude, length, n_fft, win_length, hop_length, iterations, seed=0):
    rng = np.random.default_rng(seed)
    phase = np.exp(2j * np.pi * rng.random(magnitude.shape))

    for _ in range(iterations):
        audio = istft(magnitude * phase, length, n_fft, win_length, hop_length)
        rebuilt = stft(audio, n_fft, win_length, hop_length)
        phase = rebuilt / np.maximum(np.abs(rebuilt), 1e-12)

    return istft(magnitude * phase, length, n_fft, win_length, hop_length)


def magnitude_error(audio, target, n_fft, win_length, hop_length):
    rebuilt = np.abs(stft(audio, n_fft, win_length, hop_length))
    return np.linalg.norm(rebuilt - target) / np.linalg.norm(target)


def demo_signal(sample_rate=24_000, duration=1.8):
    time = np.arange(round(sample_rate * duration)) / sample_rate
    fundamental = 135 + 8 * np.sin(2 * np.pi * 2 * time)
    phase = 2 * np.pi * np.cumsum(fundamental) / sample_rate
    first_formant = np.interp(time / duration, [0, 0.45, 0.75, 1], [700, 270, 320, 320])
    second_formant = np.interp(time / duration, [0, 0.45, 0.75, 1], [1100, 2200, 900, 900])
    audio = np.zeros_like(time)

    for harmonic in range(1, 41):
        frequency = harmonic * fundamental
        strength = (
            0.12
            + np.exp(-0.5 * ((frequency - first_formant) / 110) ** 2)
            + 0.7 * np.exp(-0.5 * ((frequency - second_formant) / 180) ** 2)
        ) / harmonic
        audio += strength * np.sin(harmonic * phase)

    fade = np.minimum(1, time / 0.08) * np.minimum(1, (duration - time) / 0.08)
    audio *= np.maximum(fade, 0)
    return 0.75 * audio / np.max(np.abs(audio))


def read_wav(path):
    with wave.open(str(path), "rb") as file:
        if file.getnchannels() != 1 or file.getsampwidth() != 2:
            raise ValueError("Input must be a mono, 16-bit PCM WAV file")
        sample_rate = file.getframerate()
        audio = np.frombuffer(file.readframes(file.getnframes()), dtype="<i2") / 32768.0
    return audio, sample_rate


def write_wav(path, audio, sample_rate):
    samples = np.round(np.clip(audio, -1, 1) * 32767).astype("<i2")
    with wave.open(str(path), "wb") as file:
        file.setnchannels(1)
        file.setsampwidth(2)
        file.setframerate(sample_rate)
        file.writeframes(samples.tobytes())


def main():
    parser = argparse.ArgumentParser(description="Hear Griffin-Lim recover phase")
    parser.add_argument("--input", type=Path, help="mono 16-bit PCM WAV; defaults to a synthetic vowel")
    parser.add_argument("--output", type=Path, default=Path(__file__).parent)
    args = parser.parse_args()

    audio, sample_rate = read_wav(args.input) if args.input else (demo_signal(), 24_000)
    n_fft = 2048 if sample_rate == 24_000 else 1 << (round(0.05 * sample_rate) - 1).bit_length()
    win_length = round(0.05 * sample_rate)
    hop_length = round(0.0125 * sample_rate)
    target = np.abs(stft(audio, n_fft, win_length, hop_length))

    args.output.mkdir(parents=True, exist_ok=True)
    write_wav(args.output / "source.wav", audio, sample_rate)

    for iterations in (0, 1, 8, 32):
        rebuilt = griffin_lim(target, len(audio), n_fft, win_length, hop_length, iterations)
        write_wav(args.output / f"iteration-{iterations}.wav", rebuilt, sample_rate)
        error = magnitude_error(rebuilt, target, n_fft, win_length, hop_length)
        print(f"{iterations:2d} iterations: relative magnitude error {error:.4f}")


if __name__ == "__main__":
    main()