#!/usr/bin/env python3
"""Generate Chapter 2 heterodyne/control figures and gesture-listening audio."""

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/ch02"
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",
    }
)
RED = "#7F1D1D"
TEAL = "#0F766E"
BLUE = "#1D4ED8"
PURPLE = "#7C3AED"


def normalize_rms(signal: np.ndarray, target: float = 0.17) -> 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: scaled heterodyne example. Frequencies are deliberately low enough to see.
t = np.linspace(0.0, 1.4, 6000, endpoint=False)
f_fixed = 10.0
f_variable = 11.5
fixed = np.cos(2 * np.pi * f_fixed * t)
variable = np.cos(2 * np.pi * f_variable * t)
mixed = fixed * variable
low_component = 0.5 * np.cos(2 * np.pi * (f_variable - f_fixed) * t)
high_component = 0.5 * np.cos(2 * np.pi * (f_variable + f_fixed) * t)

fig, axes = plt.subplots(3, 1, figsize=(7.2, 7.1), constrained_layout=True)
axes[0].plot(t, fixed, color=BLUE, linewidth=0.9, label="fixed oscillator: 10")
axes[0].plot(t, variable, color=TEAL, linewidth=0.9, alpha=0.85, label="variable oscillator: 11.5")
axes[0].set_xlim(0, 0.55)
axes[0].set_ylabel("amplitude")
axes[0].set_title("Two nearby oscillators")
axes[0].legend(loc="upper right", frameon=False, ncol=2)

axes[1].plot(t, mixed, color="#111827", linewidth=0.75, label="mixer output")
axes[1].plot(t, low_component, color=RED, linewidth=1.8, label="difference component: 1.5")
axes[1].set_xlim(0, 1.4)
axes[1].set_ylabel("amplitude")
axes[1].set_title("Multiplication creates a slowly varying difference component")
axes[1].legend(loc="upper right", frameon=False)

markerline, stemlines, baseline = axes[2].stem(
    [f_variable - f_fixed, f_variable + f_fixed], [0.5, 0.5], basefmt=" "
)
plt.setp(markerline, color=PURPLE, markersize=5)
plt.setp(stemlines, color=PURPLE, linewidth=1.5)
axes[2].axvspan(0, 5, color="#FEE2E2", alpha=0.8, label="low-pass output region")
axes[2].set_xlim(0, 25)
axes[2].set_ylim(0, 0.62)
axes[2].set_xlabel("frequency in scaled units")
axes[2].set_ylabel("relative amplitude")
axes[2].set_title("The mixer contains difference and sum frequencies")
axes[2].legend(loc="upper right", frameon=False)
fig.suptitle("Heterodyning: the audible pitch is a frequency difference", fontsize=12, fontweight="bold")
fig.savefig(SVG / "ch02-heterodyne-principle.svg", bbox_inches="tight")
plt.close(fig)

# Figure 2: idealized control chain. The capacitance curve is a teaching analogy, not a theremin calibration.
distance = np.linspace(0.12, 1.0, 500)
capacitance = 1.0 / distance
capacitance = (capacitance - capacitance.min()) / (capacitance.max() - capacitance.min())
downward_shift_magnitude = np.sqrt(capacitance + 0.03)
downward_shift_magnitude = (
    (downward_shift_magnitude - downward_shift_magnitude.min())
    / (downward_shift_magnitude.max() - downward_shift_magnitude.min())
)
pitch_octaves = np.log2(1.0 + 3.0 * downward_shift_magnitude)

fig, axes = plt.subplots(1, 3, figsize=(8.2, 3.0), constrained_layout=True)
axes[0].plot(distance, capacitance, color=TEAL, linewidth=2)
axes[0].invert_xaxis()
axes[0].set(xlabel="hand approaches antenna →", ylabel="relative capacitance", title="Distance → capacitance")
axes[1].plot(capacitance, downward_shift_magnitude, color=BLUE, linewidth=2)
axes[1].set(
    xlabel="relative capacitance",
    ylabel="magnitude of downward RF shift",
    title="Capacitance → RF shift",
)
axes[2].plot(downward_shift_magnitude, pitch_octaves, color=RED, linewidth=2)
axes[2].set(
    xlabel="relative audible difference frequency",
    ylabel="pitch in octaves",
    title="Difference → octaves",
)
fig.suptitle("Why a playable electronic instrument needs a calibrated gesture map", fontsize=11, fontweight="bold")
fig.savefig(SVG / "ch02-gesture-to-pitch.svg", bbox_inches="tight")
plt.close(fig)

# Listening file: steady tone, portamento, vibrato, and independent volume articulation.
sample_rate = 48_000
silence = np.zeros(int(sample_rate * 0.4))
segments = []

def expressive_tone(frequency: np.ndarray, amplitude: np.ndarray) -> np.ndarray:
    phase = 2 * np.pi * np.cumsum(frequency) / sample_rate
    tone = np.sin(phase) + 0.22 * np.sin(2 * phase) + 0.08 * np.sin(3 * phase)
    return normalize_rms(tone * amplitude)

# Steady reference.
duration = 2.0
time = np.arange(int(sample_rate * duration)) / sample_rate
segments += [fade(expressive_tone(np.full_like(time, 330.0), np.ones_like(time)), sample_rate), silence]

# Continuous glissando: interpolation in log-frequency gives equal octave motion.
duration = 3.0
time = np.arange(int(sample_rate * duration)) / sample_rate
frequency = 220.0 * np.power(660.0 / 220.0, time / duration)
segments += [fade(expressive_tone(frequency, np.ones_like(time)), sample_rate), silence]

# Vibrato, ±25 cents around A4.
duration = 3.0
time = np.arange(int(sample_rate * duration)) / sample_rate
cents = 25.0 * np.sin(2 * np.pi * 5.5 * time)
frequency = 440.0 * np.power(2.0, cents / 1200.0)
segments += [fade(expressive_tone(frequency, np.ones_like(time)), sample_rate), silence]

# Independent amplitude articulation while pitch stays steady.
duration = 3.0
time = np.arange(int(sample_rate * duration)) / sample_rate
amplitude = np.clip(0.5 + 0.48 * np.sin(2 * np.pi * 1.2 * time - np.pi / 2), 0.02, 1.0)
segments += [fade(expressive_tone(np.full_like(time, 392.0), amplitude), sample_rate)]

write_wav(AUDIO / "ch02-continuous-control-studies.wav", np.concatenate(segments), sample_rate)
print("Generated Chapter 2 figures and audio.")