#!/usr/bin/env python3
"""Generate Chapter 3 disc, tape-rate, envelope, and ring examples."""
from pathlib import Path
import wave
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, FancyArrowPatch, Rectangle

ROOT = Path(__file__).resolve().parents[3]
FIG = ROOT / "assets/figures/svg"
AUDIO = ROOT / "assets/audio/ch03"
FIG.mkdir(parents=True, exist_ok=True)
AUDIO.mkdir(parents=True, exist_ok=True)
FS = 48_000
plt.rcParams["svg.hashsalt"] = "contrapunk-ch03-foundations-v1"
RED = "#7f1d1d"
TEAL = "#0f6f70"
GOLD = "#b07d21"
INK = "#211d1a"


def save(name):
    plt.tight_layout()
    path = FIG / name
    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 write_wav(name, samples):
    samples = np.asarray(samples, dtype=np.float64)
    if not np.all(np.isfinite(samples)):
        raise ValueError(f"{name} contains a non-finite sample")
    if np.max(np.abs(samples), initial=0) > .180001:
        raise ValueError(f"{name} exceeds the 0.18 peak limit")
    pcm = np.round(np.clip(samples, -1, 1) * 32767).astype("<i2")
    with wave.open(str(AUDIO / name), "wb") as output:
        output.setparams((1, 2, FS, len(pcm), "NONE", "not compressed"))
        output.writeframes(pcm.tobytes())


def edge_fade(samples, seconds=.01):
    result = np.asarray(samples, dtype=np.float64).copy()
    frames = min(round(seconds * FS), len(result) // 2)
    if frames:
        ramp = np.linspace(0, 1, frames, endpoint=False)
        result[:frames] *= ramp
        result[-frames:] *= ramp[::-1]
    return result


def peak_scale(samples, peak=.16):
    samples = np.asarray(samples, dtype=np.float64)
    current = np.max(np.abs(samples), initial=0)
    return samples.copy() if current == 0 else samples * peak / current


def synthetic_strike(seconds=.25):
    time = np.arange(round(seconds * FS)) / FS
    signal = (np.sin(2*np.pi*317*time) + .55*np.sin(2*np.pi*701*time) + .25*np.sin(2*np.pi*1193*time)) * np.exp(-18*time)
    signal *= np.minimum(1, time/.002)
    return signal / np.max(np.abs(signal))


def rate_read(signal, rate):
    positions = np.arange(int(np.floor((len(signal) - 1) / rate)) + 1) * rate
    left = np.floor(positions).astype(int)
    right = np.minimum(left + 1, len(signal) - 1)
    fraction = positions - left
    return signal[left] * (1 - fraction) + signal[right] * fraction


def adsr(duration, attack, decay, sustain, gate_off, release):
    time = np.arange(round(duration * FS)) / FS
    envelope = np.zeros_like(time)
    attack_end = attack
    decay_end = attack + decay
    if attack > 0:
        mask = time < attack_end
        envelope[mask] = time[mask] / attack
    else:
        envelope[time < decay_end] = 1
    if decay > 0:
        mask = (time >= attack_end) & (time < decay_end)
        envelope[mask] = 1 - (1 - sustain) * (time[mask] - attack_end) / decay
    mask = (time >= decay_end) & (time < gate_off)
    envelope[mask] = sustain
    if release > 0:
        mask = (time >= gate_off) & (time < gate_off + release)
        envelope[mask] = sustain * (1 - (time[mask] - gate_off) / release)
    return np.clip(envelope, 0, 1)


# The same fixed trace read at three rates. Only read rate changes.
strike = synthetic_strike()
for rate, name in [(.5, "half"), (1, "normal"), (2, "double")]:
    write_wav(f"03-rate-{name}.wav", peak_scale(edge_fade(rate_read(strike, rate))))

# Project-authored reconstruction of ordinary and deliberately closed disc grooves.
fig, axes = plt.subplots(1, 3, figsize=(11, 3.8))
angle = np.linspace(0, 7*np.pi, 900)
radius = 1 - .025 * angle
axes[0].plot(radius*np.cos(angle), radius*np.sin(angle), color=INK, linewidth=1.6)
axes[0].add_patch(FancyArrowPatch((-.63, .45), (-.54, .35), arrowstyle="->", mutation_scale=12, color=RED))
axes[0].set_title("Ordinary groove")
axes[0].text(0, -1.2, "stylus advances inward", ha="center", color=TEAL)
axes[1].add_patch(Circle((0, 0), .72, fill=False, linewidth=2.2, color=RED))
axes[1].add_patch(Circle((0, 0), .48, fill=False, linewidth=.7, color="#a99d92"))
axes[1].plot([.72], [0], "o", color=INK)
axes[1].add_patch(FancyArrowPatch((.35, .62), (-.1, .71), connectionstyle="arc3,rad=.3", arrowstyle="->", mutation_scale=12, color=TEAL))
axes[1].set_title("Closed groove")
axes[1].text(0, -1.2, "stylus returns to the same circle", ha="center", color=TEAL)
for axis in axes[:2]:
    axis.set_aspect("equal")
    axis.set_xlim(-1.25, 1.25)
    axis.set_ylim(-1.35, 1.15)
    axis.axis("off")
for index in range(4):
    axes[2].add_patch(Rectangle((index, .3), .86, .42, facecolor=[RED, TEAL, GOLD, RED][index], alpha=.9))
    axes[2].text(index + .43, .51, "same\ntrace", ha="center", va="center", color="white", fontsize=9)
axes[2].set(xlim=(-.1, 4), ylim=(0, 1), title="Audible result")
axes[2].set_xlabel("time → repeated revolutions")
axes[2].set_yticks([])
axes[2].spines[["left", "right", "top"]].set_visible(False)
fig.suptitle("Disc-era repetition: the support determines the loop")
save("ch03-closed-groove-reconstruction.svg")

# Four fixed envelope hearings. Sustain is always a level; gate_off controls its duration.
envelope_specs = {
    "onset": (.35, .12, .65, 1.2, .2),
    "sustain": (.02, .1, .65, 1.65, .15),
    "accent": (.01, .22, .35, 1.2, .2),
    "release": (.01, .12, .6, .9, .9),
}
envelope_duration = 2.0
envelope_time = np.arange(round(envelope_duration * FS)) / FS
envelope_source = .16 * (np.sin(2*np.pi*220*envelope_time) + .22*np.sin(2*np.pi*440*envelope_time)) / 1.22
write_wav("03-envelope-source.wav", edge_fade(envelope_source))
fig, axes = plt.subplots(4, 1, figsize=(9, 7), sharex=True)
for axis, (name, (attack, decay_time, sustain, gate_off, release)) in zip(axes, envelope_specs.items()):
    envelope = adsr(envelope_duration, attack, decay_time, sustain, gate_off, release)
    write_wav(f"03-envelope-{name}.wav", envelope_source * envelope)
    plot_time = np.unique([0, attack, attack + decay_time, gate_off, gate_off + release, envelope_duration])
    plot_indices = np.minimum(np.round(plot_time * FS).astype(int), len(envelope) - 1)
    plot_envelope = envelope[plot_indices]
    axis.plot(plot_time, plot_envelope, color=TEAL, label="e(t)")
    axis.fill_between(plot_time, -plot_envelope, plot_envelope, color=RED, alpha=.16, label="output bounds")
    axis.axvline(gate_off, color=GOLD, linestyle="--", linewidth=1)
    axis.set_ylim(-1.05, 1.05)
    axis.set_ylabel(name)
    axis.grid(alpha=.18)
axes[0].set_title("One source multiplied by four envelope contours")
axes[-1].set_xlabel("Time (s); dashed line is gate off")
save("ch03-envelope-roles.svg")

# Carrier, modulator, their linear sum, and their ring product.
ring_time = np.arange(FS) / FS
fade = np.ones_like(ring_time)
fade_frames = round(.01 * FS)
fade[:fade_frames] = np.linspace(0, 1, fade_frames, endpoint=False)
fade[-fade_frames:] = np.linspace(1, 0, fade_frames, endpoint=False)
carrier = .14 * np.cos(2*np.pi*440*ring_time) * fade
modulator = .14 * np.cos(2*np.pi*110*ring_time) * fade
linear_sum = .07 * (np.cos(2*np.pi*440*ring_time) + np.cos(2*np.pi*110*ring_time)) * fade
ring_product = .14 * np.cos(2*np.pi*440*ring_time) * np.cos(2*np.pi*110*ring_time) * fade
for name, signal in [("carrier", carrier), ("modulator", modulator), ("linear-sum", linear_sum), ("product", ring_product)]:
    write_wav(f"03-ring-{name}.wav", signal)
ring_rows = [
    ("Carrier", carrier, [(440, 1)]),
    ("Modulator", modulator, [(110, 1)]),
    ("Add", linear_sum, [(110, .5), (440, .5)]),
    ("Multiply", ring_product, [(330, .5), (550, .5)]),
]
fig, axes = plt.subplots(4, 2, figsize=(10, 8), gridspec_kw={"width_ratios": [2.2, 1]})
window = ring_time < .04
for row, (label, signal, components) in enumerate(ring_rows):
    axes[row, 0].plot(ring_time[window] * 1000, signal[window], color=RED if label == "Multiply" else TEAL)
    axes[row, 0].set_ylabel(label)
    axes[row, 0].grid(alpha=.18)
    for frequency, magnitude in components:
        axes[row, 1].vlines(frequency, 0, magnitude, color=RED if label == "Multiply" else TEAL, linewidth=3)
        axes[row, 1].text(frequency, magnitude + .05, f"{frequency} Hz", ha="center", fontsize=9)
    axes[row, 1].set(xlim=(0, 650), ylim=(0, 1.18), yticks=[])
    axes[row, 1].grid(axis="x", alpha=.18)
axes[0, 0].set_title("First 40 ms")
axes[0, 1].set_title("Ideal one-sided components")
axes[-1, 0].set_xlabel("Time (ms)")
axes[-1, 1].set_xlabel("Frequency (Hz)")
fig.suptitle("Addition keeps 110 and 440 Hz; multiplication creates 330 and 550 Hz")
save("ch03-ring-step-by-step.svg")

print("Generated Chapter 3 disc, rate, envelope, and ring examples.")