#!/usr/bin/env python3
"""Generate Chapter 5 figures and project-authored PCM listening studies."""
from pathlib import Path
import wave

import matplotlib.pyplot as plt
import numpy as np

plt.rcParams["svg.hashsalt"] = "contrapunk-ch05-study"
ROOT = Path(__file__).resolve().parents[3]
FIG = ROOT / "assets/figures/svg"
AUDIO = ROOT / "assets/audio/ch05"
FIG.mkdir(parents=True, exist_ok=True)
AUDIO.mkdir(parents=True, exist_ok=True)
FS = 48_000
RED = "#7f1d1d"
TEAL = "#0f6f70"


def save(name):
    path = FIG / name
    plt.tight_layout()
    plt.savefig(path, format="svg", metadata={"Date": None})
    plt.close()
    path.write_text("\n".join(line.rstrip() for line in path.read_text().splitlines()) + "\n")


def quantize_mid_rise(values, bits, full_scale=1.0):
    levels = 2 ** bits
    step = 2 * full_scale / levels
    codes = np.clip(np.floor((values + full_scale) / step), 0, levels - 1)
    return -full_scale + (codes + .5) * step


def tone(frequency, seconds, bits):
    frames = round(seconds * FS)
    time = np.arange(frames) / FS
    raw = .7 * np.sin(2 * np.pi * frequency * time)
    quantized = quantize_mid_rise(raw, bits)
    fade_frames = max(1, round(.006 * FS))
    envelope = np.ones(frames)
    envelope[:fade_frames] = np.linspace(0, 1, fade_frames, endpoint=False)
    envelope[-fade_frames:] = np.linspace(1, 0, fade_frames, endpoint=True)
    return .18 * quantized * envelope


# The opening sequence: exactly two cycles of 1 kHz at 8 kHz.
n = np.arange(16)
values = np.sin(2 * np.pi * 1_000 * n / 8_000)
fig, ax = plt.subplots(figsize=(9, 4))
ax.stem(n, values, linefmt=RED, markerfmt="o", basefmt="k-")
for index, value in zip(n, values):
    ax.text(index, value + (.11 if value >= 0 else -.16), f"{value:.3f}", ha="center", fontsize=7)
ax.set(xlabel="Sample n", ylabel="Value", ylim=(-1.25, 1.25), title="Sixteen numbers become two cycles of a 1 kHz tone at 8 kHz")
ax.grid(alpha=.2)
save("ch05-number-sequence.svg")

# One waveform under three quantizers.
time = np.arange(round(.012 * FS)) / FS
reference = .7 * np.sin(2 * np.pi * 440 * time)
fig, axes = plt.subplots(3, 1, figsize=(9, 7), sharex=True)
for ax, bits in zip(axes, [3, 5, 8]):
    ax.plot(time * 1_000, reference, color="#9b9188", linewidth=1, label="reference")
    ax.step(time * 1_000, quantize_mid_rise(reference, bits), where="post", color=TEAL, linewidth=1.2, label=f"{bits}-bit")
    ax.set(ylabel="Amplitude", ylim=(-.9, .9), title=f"{bits} bits: {2 ** bits} levels")
    ax.grid(alpha=.2)
    ax.legend(loc="upper right")
axes[-1].set_xlabel("Time (ms)")
save("ch05-quantized-waveforms.svg")

# Listening sequence: same 440 Hz tone at 3, 5, 8, and 16 bits.
silence = np.zeros(round(.22 * FS))
clips = []
for bits in [3, 5, 8, 16]:
    clips.extend([tone(440, .7, bits), silence])

# Original eight-note study. It copies no melody or recording.
melody = [220, 275, 330, 440, 330, 275, 247.5, 220]
for bits in [4, 12]:
    phrase = np.concatenate([tone(frequency, .24, bits) for frequency in melody])
    clips.extend([phrase, silence])

output = np.concatenate(clips)
with wave.open(str(AUDIO / "ch05-pcm-studies.wav"), "wb") as wav_file:
    wav_file.setparams((1, 2, FS, len(output), "NONE", "not compressed"))
    wav_file.writeframes((np.clip(output, -1, 1) * 32767).astype("<i2").tobytes())

assert np.allclose(values[:8], [0, np.sqrt(.5), 1, np.sqrt(.5), 0, -np.sqrt(.5), -1, -np.sqrt(.5)], atol=1e-12)
assert len(np.unique(quantize_mid_rise(np.linspace(-1, 1, 1000, endpoint=False), 3))) == 8
assert np.max(np.abs(output)) < .19
print("Generated Chapter 5 figures and project-authored PCM audio.")