#!/usr/bin/env python3
"""A synthetic signal experiment: 48 kHz -> 6 kHz with and without filtering.

Usage: python audio-aliasing-lab.py [output-directory]
No microphone, hardware, speech recognizer or listening test is performed.
Spectrum measurements use float arrays before PCM16 quantization.
"""

import importlib.metadata
import json
from pathlib import Path
import platform
import sys
import wave

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import resample_poly

FS_IN = 48000
FS_OUT = 6000
SECONDS = 2
TRIM_SECONDS = 0.1
AMPLITUDE = 0.4


def tone(hz, phase=0.0):
    n = np.arange(FS_IN * SECONDS, dtype=np.float64)
    return AMPLITUDE * np.sin(2 * np.pi * hz * n / FS_IN + phase)


def spectrum(samples, fs):
    trim = int(TRIM_SECONDS * fs)
    center = samples[trim:-trim]
    freqs = np.fft.rfftfreq(center.size, d=1 / fs)
    amplitudes = 2 * np.abs(np.fft.rfft(center)) / center.size
    amplitudes[0] /= 2
    if center.size % 2 == 0:
        amplitudes[-1] /= 2
    return freqs, amplitudes


def at_frequency(freqs, amplitudes, hz):
    index = int(np.argmin(np.abs(freqs - hz)))
    assert np.isclose(freqs[index], hz)
    return float(amplitudes[index])


def save_wav(path, samples, fs):
    assert np.max(np.abs(samples)) < 1, "The example would clip"
    pcm = np.rint(samples * 32767).astype("<i2")
    with wave.open(str(path), "wb") as output:
        output.setnchannels(1)
        output.setsampwidth(2)
        output.setframerate(fs)
        output.writeframes(pcm.tobytes())
    with wave.open(str(path), "rb") as saved:
        assert saved.getnchannels() == 1 and saved.getsampwidth() == 2
        assert saved.getframerate() == fs and saved.getnframes() == len(samples)
        reread = np.frombuffer(saved.readframes(saved.getnframes()), dtype="<i2")
    error = float(np.max(np.abs(reread.astype(np.float64) / 32767 - samples)))
    assert error <= 0.5 / 32767 + 1e-12
    return {"file": path.name, "sample_rate_hz": fs, "frames": len(samples),
            "seconds": len(samples) / fs, "channels": 1, "pcm_bits": 16,
            "max_pcm_round_trip_absolute_error": error}


def main():
    output = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("audio-aliasing-output")
    output.mkdir(parents=True, exist_ok=True)
    original = tone(900) + tone(3900)
    stride = original[::8]
    filtered = resample_poly(
        original, up=1, down=8, window=("kaiser", 5.0), padtype="constant",
    )
    assert len(stride) == len(filtered) == FS_OUT * SECONDS
    cases = [
        ("original", original, FS_IN, (900, 2100, 3900)),
        ("stride", stride, FS_OUT, (900, 2100)),
        ("filtered", filtered, FS_OUT, (900, 2100)),
    ]
    measurements, wavs = [], []
    spectra = []
    for name, samples, fs, frequencies in cases:
        freqs, amps = spectrum(samples, fs)
        measurements.append({
            "case": name, "sample_rate_hz": fs,
            "measurement_frames": len(samples) - 2 * int(TRIM_SECONDS * fs),
            "peak_amplitudes_at_hz": {
                str(hz): at_frequency(freqs, amps, hz) for hz in frequencies},
        })
        wavs.append(save_wav(output / f"{name}.wav", samples, fs))
        spectra.append((freqs, amps))
    assert np.isclose(measurements[0]["peak_amplitudes_at_hz"]["900"], 0.4)
    assert measurements[0]["peak_amplitudes_at_hz"]["2100"] < 1e-10
    assert np.isclose(measurements[1]["peak_amplitudes_at_hz"]["2100"], 0.4)
    assert measurements[2]["peak_amplitudes_at_hz"]["2100"] < 1e-4
    assert abs(measurements[2]["peak_amplitudes_at_hz"]["900"] - 0.4) < 0.001

    n_out = np.arange(FS_OUT * SECONDS, dtype=np.float64)
    high_only_stride = tone(3900)[::8]
    alias = -AMPLITUDE * np.sin(2 * np.pi * 2100 * n_out / FS_OUT)
    alias_error = float(np.max(np.abs(high_only_stride - alias)))
    assert alias_error < 1e-10
    boundary_sine_peak = float(np.max(np.abs(tone(3000)[::8])))
    boundary_cosine_peak = float(np.max(np.abs(tone(3000, np.pi / 2)[::8])))
    assert boundary_sine_peak < 1e-10
    assert np.isclose(boundary_cosine_peak, AMPLITUDE)

    fig, axes = plt.subplots(3, 1, figsize=(10, 8), sharex=True, sharey=True)
    labels = ["Input: 48,000 samples/s", "Take every 8th sample: 6,000 samples/s",
              "Low-pass + resample_poly: 6,000 samples/s"]
    colors = ["#176493", "#b44039", "#16744c"]
    for i, (axis, (freqs, amps), label, color) in enumerate(zip(axes, spectra, labels, colors)):
        axis.plot(freqs, amps, color=color, linewidth=1.2)
        axis.set_title(label, loc="left", fontsize=12)
        axis.set_ylabel("Peak amplitude")
        axis.set_ylim(-0.012, 0.46)
        axis.grid(alpha=0.2)
        if i:
            axis.axvspan(3000, 4500, color="#e7e7e7", alpha=0.8)
            axis.text(3700, 0.31, "Outside output\nNyquist range", ha="center", fontsize=10)
        for hz in ((900, 3900) if i == 0 else (900, 2100)):
            amp = at_frequency(freqs, amps, hz)
            axis.annotate(f"{hz} Hz", (hz, amp), xytext=(hz + 110, amp + 0.028), fontsize=10)
    residual = measurements[2]["peak_amplitudes_at_hz"]["2100"]
    axes[-1].text(2420, 0.17, f"2,100 Hz residual:\n{residual:.7f}", ha="center", fontsize=10)
    axes[-1].set_xlabel("Frequency (Hz)")
    axes[-1].set_xlim(0, 4500)
    fig.suptitle("Same two-tone input, different downsampling", fontsize=15)
    fig.tight_layout(rect=(0, 0, 1, 0.95))
    fig.savefig(output / "spectrum.png", dpi=180)
    plt.close(fig)

    report = {
        "experiment": "synthetic_two_tone_downsampling",
        "environment": {"python": platform.python_version(), "macos": platform.mac_ver()[0],
                        "machine": platform.machine(),
                        "packages": {name: importlib.metadata.version(name)
                                     for name in ("numpy", "scipy", "matplotlib")}},
        "input": {"sample_rate_hz": FS_IN, "seconds": SECONDS, "tones_hz": [900, 3900],
                  "peak_amplitude_per_tone": AMPLITUDE, "dtype": "float64"},
        "conversion": {"output_sample_rate_hz": FS_OUT, "up": 1, "down": 8,
                       "window": ["kaiser", 5.0], "padtype": "constant", "cval": 0},
        "measurement": {"trim_seconds_at_each_end": TRIM_SECONDS, "central_seconds": 1.8,
                        "method": "one-sided rFFT peak amplitude; DC and Nyquist not doubled",
                        "before_pcm_quantization": True},
        "measurements": measurements,
        "boundary_cases": {"3900_stride_equals_negative_2100_max_abs_error": alias_error,
                           "3000_zero_phase_stride_peak": boundary_sine_peak,
                           "3000_pi_over_2_phase_stride_peak": boundary_cosine_peak},
        "wav_files": wavs,
        "scope": "No microphones, hardware playback, subjective listening, speech or streaming tests.",
    }
    (output / "results.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8",
    )
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()