#!/usr/bin/env python3
"""Generate deterministic visuals for Chapter 4 equations."""
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

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


def save(name):
    path = OUT / 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 base(title):
    figure, axis = plt.subplots(figsize=(8.6, 3.8))
    axis.set_title(title)
    axis.grid(alpha=0.2)
    return figure, axis


# Equation 4.1: phase-accumulator VCO.
count = 960
time = np.arange(count) / FS
frequency = np.where(time < 0.01, 220.0, 440.0)
phase = np.zeros(count)
for index in range(count - 1):
    phase[index + 1] = (phase[index] + frequency[index] / FS) % 1
waveform = 2 * phase - 1
figure, axes = plt.subplots(2, 1, figsize=(8.6, 5), sharex=True)
axes[0].plot(time * 1_000, phase, color=TEAL)
axes[1].plot(time * 1_000, waveform, color=RED)
axes[0].axvline(10, color=GOLD, linestyle="--")
axes[1].axvline(10, color=GOLD, linestyle="--")
axes[0].set(ylabel="Phase (cycles)", title="Eq. 4.1: phase accumulation at 220 Hz, then 440 Hz")
axes[1].set(xlabel="Time (ms)", ylabel="Saw output")
for axis in axes:
    axis.grid(alpha=0.2)
save("eq-4-1-phase-vco.svg")

# Equation 4.2: exponential pitch control.
control = np.linspace(-2, 2, 401)
pitch = 220 * 2**control
figure, axis = base("Eq. 4.2: exponential pitch control from a 220 Hz base")
axis.semilogy(control, pitch, color=TEAL)
axis.scatter([-1, 0, 1], [110, 220, 440], color=RED, zorder=3)
axis.set(xlabel="Control v (octaves)", ylabel="Frequency (Hz)")
save("eq-4-2-exponential-pitch.svg")

# Equation 4.3: VCA multiplication.
count = round(0.03 * FS)
time = np.arange(count) / FS
source = np.sin(TAU * 220 * time)
gain = np.linspace(0, 1, count)
output = gain * source
figure, axis = base("Eq. 4.3: a VCA multiplies a 220 Hz sine by gain 0 to 1")
axis.plot(time * 1_000, source, color=INK, alpha=0.35, label="input x[n]")
axis.plot(time * 1_000, gain, color=GOLD, label="gain g[n]")
axis.plot(time * 1_000, output, color=TEAL, label="output y[n]")
axis.set(xlabel="Time (ms)", ylabel="Normalized amplitude")
axis.legend()
save("eq-4-3-vca.svg")

# Equation 4.4: one-pole low-pass step responses.
count = 500
source = np.ones(count)
figure, axis = base("Eq. 4.4: one-pole step responses at 48 kHz")
responses = {}
for cutoff, color in zip((200, 900, 4_000), (GOLD, TEAL, RED)):
    coefficient = np.exp(-TAU * cutoff / FS)
    output = np.zeros(count)
    previous = 0.0
    for index, sample in enumerate(source):
        previous = (1 - coefficient) * sample + coefficient * previous
        output[index] = previous
    responses[cutoff] = output
    axis.plot(np.arange(count) / FS * 1_000, output, color=color, label=f"{cutoff:,} Hz")
axis.set(xlabel="Time (ms)", ylabel="Output")
axis.legend()
save("eq-4-4-one-pole.svg")

# Equation 4.5: exact linear envelope segment.
time = np.linspace(0, 1.2, 301)
start = 0.25
target = 0.9
duration = 0.8
envelope = start + (target - start) * np.clip(time / duration, 0, 1)
figure, axis = base("Eq. 4.5: linear segment from 0.25 to 0.9 in 0.8 seconds")
axis.plot(time, envelope, color=TEAL)
axis.scatter([0, duration], [start, target], color=RED, zorder=3)
axis.set(xlabel="Time (s)", ylabel="Envelope level", ylim=(0, 1))
save("eq-4-5-envelope-segment.svg")

# Equation 4.6: control-rate zero-order hold.
ratio = 8
control = np.array([0.1, 0.8, 0.35, 0.65])
held = np.repeat(control, ratio)
figure, axis = base("Eq. 4.6: four control values held for eight audio samples each")
axis.step(np.arange(len(held)), held, where="post", color=TEAL)
axis.scatter(np.arange(0, len(held), ratio), control, color=RED, zorder=3)
axis.set(xlabel="Audio sample n", ylabel="Held control", ylim=(0, 1))
save("eq-4-6-control-hold.svg")

# Equation 4.7: delayed bounded feedback with an identity processor y[n]=u[n].
count = 240
source = np.zeros(count)
source[0] = 0.5
figure, axis = base("Eq. 4.7: delayed bounded feedback, identity-processor case")
feedback_outputs = {}
for gain_value, color in zip((0.0, 0.7, 0.95), (INK, GOLD, RED)):
    output = np.zeros(count)
    for index in range(count):
        delayed = output[index - 1] if index else 0.0
        output[index] = source[index] + gain_value * np.tanh(delayed)
    feedback_outputs[gain_value] = output
    axis.plot(np.arange(count), output, color=color, label=f"k={gain_value:g}")
axis.set(xlabel="Sample n", ylabel="Output y[n]")
axis.legend()
save("eq-4-7-bounded-feedback.svg")

assert phase[0] == 0 and np.isclose(phase[1], 220 / FS)
assert np.isclose(220 * 2**-1, 110) and np.isclose(220 * 2, 440)
assert np.isclose(output[-1], feedback_outputs[0.95][-1])
assert np.isclose(gain[0], 0) and np.isclose(gain[-1], 1)
assert responses[4_000][20] > responses[900][20] > responses[200][20]
assert np.isclose(envelope[0], start) and np.isclose(envelope[-1], target)
assert np.array_equal(held[:ratio], np.full(ratio, control[0]))
assert all(np.isfinite(values).all() and np.max(np.abs(values)) <= 1.5 for values in feedback_outputs.values())
print("Generated Chapter 4 formula visuals.")