#!/usr/bin/env python3
"""Generate deterministic Chapter 4 modular studies and mono PCM16 WAVs."""
from pathlib import Path
import wave

import matplotlib.pyplot as plt
import numpy as np

plt.rcParams["svg.hashsalt"] = "contrapunk-ch04-studies"
ROOT = Path(__file__).resolve().parents[3]
FIG = ROOT / "assets/figures/svg"
AUDIO = ROOT / "assets/audio/ch04"
FIG.mkdir(parents=True, exist_ok=True)
AUDIO.mkdir(parents=True, exist_ok=True)
FS = 48_000
TAU = 2 * np.pi
RED = "#7f1d1d"
TEAL = "#0f6f70"
GOLD = "#a56a12"
INK = "#211b17"


def save_svg(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 edge_fade(samples, milliseconds=5):
    result = np.asarray(samples, dtype=float).copy()
    frames = min(round(milliseconds * FS / 1_000), len(result) // 2)
    if frames:
        fade = np.linspace(0, 1, frames, endpoint=False)
        result[:frames] *= fade
        result[-frames:] *= fade[::-1]
    return result


def write_wav(name, samples, gain=1.0, fade=True):
    values = edge_fade(samples) if fade else np.asarray(samples, dtype=float)
    values = np.clip(values * gain, -1, 1)
    pcm = np.rint(values * 32_767).astype("<i2")
    with wave.open(str(AUDIO / name), "wb") as output:
        output.setparams((1, 2, FS, 0, "NONE", "not compressed"))
        output.writeframes(pcm.tobytes())
    return values


def phase_from_frequency(frequency):
    frequency = np.asarray(frequency, dtype=float)
    return np.mod(np.cumsum(np.r_[0.0, frequency[:-1]]) / FS, 1)


def saw(frequency):
    phase = phase_from_frequency(frequency)
    return 2 * phase - 1, phase


def sine(frequency):
    return np.sin(TAU * phase_from_frequency(frequency))


def triangle(frequency):
    phase = phase_from_frequency(frequency)
    return 1 - 4 * np.abs(phase - 0.5)


def adsr(gate, attack, decay, sustain, release):
    output = np.zeros(len(gate))
    level = 0.0
    stage = "idle"
    origin = 0.0
    elapsed = 0
    prior = False
    settings = {
        "attack": (1.0, attack, "decay"),
        "decay": (sustain, decay, "sustain"),
        "release": (0.0, release, "idle"),
    }
    for index, on in enumerate(gate):
        if on and not prior:
            stage, origin, elapsed = "attack", level, 0
        elif prior and not on:
            stage, origin, elapsed = "release", level, 0
        prior = bool(on)
        if stage == "sustain":
            level = sustain
        elif stage == "idle":
            level = 0.0
        else:
            while stage in settings:
                target, seconds, following = settings[stage]
                frames = round(seconds * FS)
                if frames == 0:
                    level, origin, elapsed, stage = target, target, 0, following
                    continue
                amount = min(elapsed / frames, 1)
                level = origin + (target - origin) * amount
                elapsed += 1
                if elapsed >= frames:
                    origin, elapsed, stage = target, 0, following
                break
        output[index] = level
    return output


def one_pole(source, cutoff):
    output = np.zeros_like(source, dtype=float)
    previous = 0.0
    for index, sample in enumerate(source):
        frequency = cutoff[index] if np.ndim(cutoff) else cutoff
        frequency = np.clip(frequency, np.finfo(float).eps, 0.45 * FS)
        coefficient = np.exp(-TAU * frequency / FS)
        previous = (1 - coefficient) * sample + coefficient * previous
        output[index] = previous
    return output


def spectrum(source):
    count = 16_384
    window = np.hanning(count)
    if len(source) == count:
        data = source
    else:
        center = len(source) // 2
        data = source[center - count // 2:center + count // 2]
    magnitude = np.abs(np.fft.rfft(data * window)) / (window.sum() / 2)
    return np.fft.rfftfreq(count, 1 / FS), 20 * np.log10(np.maximum(magnitude, 1e-5))


# Study 1: One Voltage, Two Octaves.
count = 2 * FS
time = np.arange(count) / FS
control = (time >= 1).astype(float)
frequency = 220 * 2**control
source, phase = saw(frequency)
vco_audio = write_wav("04-vco-cv.wav", 0.2 * source)
figure, axes = plt.subplots(3, 2, figsize=(11, 6), sharex="col")
for column, start in enumerate((0, FS)):
    excerpt = slice(start, start + round(0.02 * FS))
    local_time = (np.arange(excerpt.start, excerpt.stop) - start) / FS * 1_000
    axes[0, column].plot(local_time, control[excerpt], color=GOLD)
    axes[1, column].plot(local_time, phase[excerpt], color=TEAL)
    axes[2, column].plot(local_time, source[excerpt], color=RED)
    axes[0, column].set_title(f"{frequency[start]:.0f} Hz, v={control[start]:.0f}")
    axes[0, column].set_ylim(-0.1, 1.1)
    axes[1, column].set_ylim(-0.05, 1.05)
    axes[2, column].set_ylim(-1.1, 1.1)
    axes[2, column].set_xlabel("Time from section start (ms)")
for axis, label in zip(axes[:, 0], ("Pitch control (octaves)", "Phase (cycles)", "Saw amplitude")):
    axis.set_ylabel(label)
save_svg("ch04-vco-cv.svg")

# Study 2: per-sample, held, and interpolated vibrato control.
segment = FS
silence = round(0.01 * FS)
base_time = np.arange(segment) / FS
modulator = np.sin(TAU * 5 * base_time)
reference_frequency = 440 * 2 ** ((3 / 12) * modulator)
held_frequency = np.repeat(reference_frequency[::480], 480)[:segment]
anchors = np.arange(0, segment + 480, 480)
anchor_time = anchors / FS
anchor_frequency = 440 * 2 ** ((3 / 12) * np.sin(TAU * 5 * anchor_time))
interpolated_frequency = np.interp(np.arange(segment), anchors, anchor_frequency)
versions = [reference_frequency, held_frequency, interpolated_frequency]
control_audio = np.zeros(3 * segment + 2 * silence)
for index, version in enumerate(versions):
    start = index * (segment + silence)
    control_audio[start:start + segment] = edge_fade(0.15 * sine(version))
write_wav("04-audio-vs-control-rate.wav", control_audio, fade=False)
figure, (frequency_axis, spectrum_axis) = plt.subplots(2, 1, figsize=(10, 7))
excerpt = slice(round(0.9 * FS), FS)
labels = ("per sample", "100 Hz hold", "100 Hz linear")
for version, label, color in zip(versions, labels, (INK, RED, TEAL)):
    frequency_axis.plot(base_time[excerpt], version[excerpt], color=color, label=label)
frequency_axis.set(xlabel="Time (s)", ylabel="Instantaneous frequency (Hz)")
frequency_axis.legend()
for version, label, color in zip(versions, labels, (INK, RED, TEAL)):
    hertz, decibels = spectrum(sine(version))
    mask = hertz < 3_000
    spectrum_axis.plot(hertz[mask], decibels[mask], color=color, label=label)
spectrum_axis.set(xlabel="Frequency (Hz)", ylabel="Magnitude (dBFS)", ylim=(-100, 10))
spectrum_axis.legend()
save_svg("ch04-audio-vs-control-rate.svg")

# Study 3: ADSR-controlled VCA.
count = 2 * FS
time = np.arange(count) / FS
gate = (time >= 0.1) & (time < 1.2)
envelope = adsr(gate, 0.1, 0.2, 0.55, 0.35)
carrier, _ = saw(np.full(count, 220.0))
adsr_output = 0.2 * envelope * carrier
write_wav("04-adsr-vca.wav", adsr_output)
figure, axes = plt.subplots(3, 1, figsize=(10, 6), sharex=True)
axes[0].plot(time, gate, color=GOLD)
axes[1].plot(time, envelope, color=TEAL)
axes[2].plot(time, adsr_output, color=RED)
axes[2].set_xlabel("Time (s)")
for axis, label in zip(axes, ("Gate", "Envelope and gain", "Output amplitude")):
    axis.set_ylabel(label)
save_svg("ch04-adsr-vca.svg")

# Study 4: subtractive sweep through a one-pole teaching filter.
count = 4 * FS
time = np.arange(count) / FS
source, _ = saw(np.full(count, 110.0))
cutoff = 8_000 * (200 / 8_000) ** (time / 4)
filtered = 0.2 * one_pole(source, cutoff)
write_wav("04-subtractive-sweep.wav", filtered)
figure, (cutoff_axis, spectrum_axis) = plt.subplots(2, 1, figsize=(10, 7))
cutoff_axis.semilogy(time, cutoff, color=TEAL)
cutoff_axis.set(xlabel="Time (s)", ylabel="Cutoff (Hz)")
for center, label, color in ((0.5, "0.5 s", RED), (3.5, "3.5 s", TEAL)):
    excerpt = filtered[round(center * FS) - 8_192:round(center * FS) + 8_192]
    hertz, decibels = spectrum(excerpt)
    mask = hertz < 10_000
    spectrum_axis.plot(hertz[mask], decibels[mask], color=color, label=label)
spectrum_axis.set(xlabel="Frequency (Hz)", ylabel="Magnitude (dBFS)", ylim=(-100, 5))
spectrum_axis.legend()
save_svg("ch04-subtractive-sweep.svg")

# Study 5: Clock Garden sequence.
notes = np.array([48, 55, 60, 63, 60, 55, 51, 55] * 2)
step_frames = round(0.25 * FS)
count = len(notes) * step_frames
time = np.arange(count) / FS
step_index = (np.arange(count) // step_frames) % len(notes)
gate = np.arange(count) % step_frames < round(0.8 * step_frames)
frequency = 440 * 2 ** ((notes[step_index] - 69) / 12)
envelope = adsr(gate, 0.005, 0.06, 0.6, 0.04)
source = triangle(frequency)
sequence_output = 0.22 * envelope * one_pole(source, 600 + 2_400 * envelope)
write_wav("04-sequence.wav", sequence_output)
figure, axes = plt.subplots(4, 1, figsize=(11, 7), sharex=True)
axes[0].step(time, notes[step_index] - 48, where="post", color=GOLD)
axes[1].plot(time, gate, color=INK)
axes[2].plot(time, envelope, color=TEAL)
axes[3].plot(time, sequence_output, color=RED)
axes[3].set_xlabel("Time (s)")
for axis, label in zip(axes, ("Pitch (semitones)", "Gate", "Envelope", "Output")):
    axis.set_ylabel(label)
save_svg("ch04-sequence.svg")

# Study 6: Almost Singing, bounded delayed feedback through a one-pole filter.
count = 3 * FS
time = np.arange(count) / FS
feedback_output = np.zeros(count)
rms = np.zeros(count)
for section, gain_value in enumerate((0.0, 0.7, 0.95)):
    previous_filter = 0.0
    previous_output = 0.0
    coefficient = np.exp(-TAU * 900 / FS)
    start = section * FS
    for local_index in range(FS):
        source = 0.5 if local_index == 0 else 0.0
        processor_input = source + gain_value * np.tanh(previous_output)
        previous_filter = (1 - coefficient) * processor_input + coefficient * previous_filter
        feedback_output[start + local_index] = previous_filter
        previous_output = previous_filter
for start in range(0, count, 480):
    rms[start:start + 480] = np.sqrt(np.mean(feedback_output[start:start + 480] ** 2))
feedback_audio = write_wav("04-feedback.wav", feedback_output, gain=10 ** (-12 / 20), fade=False)
figure, (waveform_axis, rms_axis) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)
waveform_axis.plot(time, feedback_output, color=RED)
rms_axis.plot(time, rms, color=TEAL)
waveform_axis.set_ylabel("Output amplitude")
rms_axis.set(xlabel="Time (s)", ylabel="10 ms RMS")
for axis in (waveform_axis, rms_axis):
    for boundary in (1, 2):
        axis.axvline(boundary, color=INK, linestyle="--", alpha=0.5)
save_svg("ch04-feedback.svg")

test_gate = np.arange(round(0.08 * FS)) / FS < 0.05
test_envelope = adsr(test_gate, 0.005, 0.01, 0.5, 0.01)
assert test_envelope[0] == 0 and test_envelope[round(0.005 * FS)] == 1
assert test_envelope[round(0.015 * FS) - 1] > 0.5 and test_envelope[round(0.015 * FS)] == 0.5
assert test_envelope[round(0.05 * FS)] == test_envelope[round(0.05 * FS) - 1]
assert test_envelope[round(0.06 * FS) - 1] > 0 and test_envelope[round(0.06 * FS)] == 0
assert not np.isclose(interpolated_frequency[-1], anchor_frequency[-2])
assert np.isclose(frequency[0], 440 * 2 ** ((48 - 69) / 12))
assert np.isclose(phase[0], 0) and np.isclose(phase[1], 220 / FS)
assert np.isclose(envelope.max(), 1)
assert len(control_audio) == 144_960
assert np.all(np.isfinite(filtered)) and np.all(np.isfinite(feedback_output))
assert max(np.max(np.abs(values)) for values in (vco_audio, control_audio, adsr_output, filtered, sequence_output, feedback_audio)) <= 0.22
print("Generated Chapter 4 figures and project-authored PCM audio.")