#!/usr/bin/env python3
"""Generate Chapter 1 harmonic figures and the level-matched listening file."""

from __future__ import annotations

import wave
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

ROOT = Path(__file__).resolve().parents[3]
SVG = ROOT / "assets/figures/svg"
AUDIO = ROOT / "assets/audio/ch01"
SVG.mkdir(parents=True, exist_ok=True)
AUDIO.mkdir(parents=True, exist_ok=True)

plt.rcParams.update(
    {
        "font.family": "DejaVu Sans",
        "font.size": 9,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.titleweight": "bold",
        "svg.fonttype": "none",
    }
)
COLORS = ["#7F1D1D", "#0F766E", "#1D4ED8", "#7C3AED", "#B45309"]


def harmonic_sum(phase: np.ndarray, amplitudes: list[float], phases: list[float] | None = None) -> np.ndarray:
    phases = phases or [0.0] * len(amplitudes)
    result = np.zeros_like(phase, dtype=float)
    for index, (amplitude, offset) in enumerate(zip(amplitudes, phases), start=1):
        result += amplitude * np.sin(2 * np.pi * index * phase + offset)
    return result


def normalize_rms(signal: np.ndarray, target: float = 0.18) -> np.ndarray:
    signal = signal - np.mean(signal)
    rms = np.sqrt(np.mean(signal**2))
    return signal * (target / rms) if rms else signal


def fade(signal: np.ndarray, sample_rate: int, seconds: float = 0.03) -> np.ndarray:
    frames = min(int(sample_rate * seconds), len(signal) // 2)
    ramp = np.linspace(0.0, 1.0, frames)
    signal = signal.copy()
    signal[:frames] *= ramp
    signal[-frames:] *= ramp[::-1]
    return signal


def write_wav(path: Path, signal: np.ndarray, sample_rate: int = 48_000) -> None:
    signal = np.clip(signal, -0.98, 0.98)
    pcm = np.round(signal * 32767).astype("<i2")
    with wave.open(str(path), "wb") as output:
        output.setnchannels(1)
        output.setsampwidth(2)
        output.setframerate(sample_rate)
        output.writeframes(pcm.tobytes())


# Figure 1: progressively combine a harmonic series.
phase = np.linspace(0.0, 2.0, 1600, endpoint=False)
amplitudes = [1.0, 0.50, 0.33, 0.25, 0.20]
fig, axes = plt.subplots(3, 1, figsize=(7.2, 6.5), sharex=True, constrained_layout=True)
for count, axis in zip([1, 3, 5], axes):
    partials = [amplitudes[k - 1] * np.sin(2 * np.pi * k * phase) for k in range(1, count + 1)]
    for k, partial in enumerate(partials, start=1):
        axis.plot(phase, partial, color=COLORS[(k - 1) % len(COLORS)], alpha=0.28, linewidth=0.8)
    total = np.sum(partials, axis=0)
    axis.plot(phase, total, color="#111827", linewidth=1.8, label="sum")
    axis.axhline(0, color="#6B7280", linewidth=0.5)
    axis.set_ylabel("amplitude")
    axis.set_title(f"Fundamental plus {count - 1} overtone{'s' if count != 2 else ''}")
axes[-1].set_xlabel("time in cycles of the fundamental")
fig.suptitle("A complex periodic tone built from harmonically related sinusoids", fontsize=12, fontweight="bold")
fig.savefig(SVG / "ch01-harmonic-construction.svg", bbox_inches="tight")
plt.close(fig)

# Figure 2: same pitch, different harmonic recipes.
recipes = {
    "Sine": [1.0] + [0.0] * 11,
    "Odd-only": [1.0 if k % 2 else 0.0 for k in range(1, 13)],
    "Saw-like": [1.0 / k for k in range(1, 13)],
    "Dark roll-off": [1.0 / (k * k) for k in range(1, 13)],
}
fig, axes = plt.subplots(len(recipes), 2, figsize=(7.2, 8.0), constrained_layout=True)
short_phase = np.linspace(0.0, 2.0, 1200, endpoint=False)
for row, (name, recipe) in enumerate(recipes.items()):
    waveform = normalize_rms(harmonic_sum(short_phase, recipe), 0.45)
    axes[row, 0].plot(short_phase, waveform, color=COLORS[row], linewidth=1.4)
    axes[row, 0].axhline(0, color="#6B7280", linewidth=0.5)
    axes[row, 0].set_ylabel(name)
    axes[row, 0].set_ylim(-1.05, 1.05)
    harmonics = np.arange(1, len(recipe) + 1)
    markerline, stemlines, baseline = axes[row, 1].stem(harmonics, recipe, basefmt=" ")
    plt.setp(markerline, color=COLORS[row], markersize=4)
    plt.setp(stemlines, color=COLORS[row], linewidth=1.2)
    axes[row, 1].set_xlim(0.5, 12.5)
    axes[row, 1].set_ylim(0.0, 1.05)
    axes[row, 1].set_ylabel("relative level")
axes[0, 0].set_title("two cycles at the same fundamental")
axes[0, 1].set_title("harmonic amplitudes")
axes[-1, 0].set_xlabel("time in cycles")
axes[-1, 1].set_xlabel("harmonic number")
fig.suptitle("Pitch can remain fixed while harmonic balance changes timbre", fontsize=12, fontweight="bold")
fig.savefig(SVG / "ch01-same-pitch-different-timbres.svg", bbox_inches="tight")
plt.close(fig)

# Listening file: four two-second, level-matched tones separated by silence.
sample_rate = 48_000
segment_seconds = 2.0
silence = np.zeros(int(sample_rate * 0.45))
t = np.arange(int(sample_rate * segment_seconds)) / sample_rate
fundamental = 220.0
segments = []
for recipe in recipes.values():
    signal = np.zeros_like(t)
    for harmonic, amplitude in enumerate(recipe, start=1):
        signal += amplitude * np.sin(2 * np.pi * harmonic * fundamental * t)
    signal = fade(normalize_rms(signal), sample_rate)
    segments.extend([signal, silence])
write_wav(AUDIO / "ch01-same-pitch-different-timbres.wav", np.concatenate(segments[:-1]), sample_rate)

print("Generated Chapter 1 figures and audio.")