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

import matplotlib.pyplot as plt
import numpy as np

plt.rcParams["svg.hashsalt"] = "contrapunk-ch05-formulas"
ROOT = Path(__file__).resolve().parents[3]
OUT = ROOT / "assets/figures/svg"
OUT.mkdir(parents=True, exist_ok=True)
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")


# Equation 5.1: sample a continuous-time model.
fs = 8_000
frequency = 1_000
n = np.arange(17)
t_samples = n / fs
t = np.linspace(0, 0.002, 800)
fig, ax = plt.subplots(figsize=(9, 4))
ax.plot(t * 1_000, np.sin(2 * np.pi * frequency * t), color="#8b8178", label="$x_c(t)$")
markerline, stemlines, _ = ax.stem(t_samples * 1_000, np.sin(2 * np.pi * frequency * t_samples), basefmt="k-")
plt.setp(stemlines, color=RED, linewidth=1.4)
plt.setp(markerline, color=RED, markerfacecolor=RED)
ax.set(xlabel="Time (ms)", ylabel="Amplitude", title="Eq. 5.1: 1 kHz sampled at 8 kHz", xlim=(0, 2), ylim=(-1.12, 1.12))
ax.legend()
ax.grid(alpha=.2)
save("eq-5-1-sampling.svg")

# Equation 5.2: index-to-time conversion.
indices = np.arange(5)
times_us = indices / 48_000 * 1_000_000
fig, ax = plt.subplots(figsize=(9, 3.8))
ax.hlines(0, 0, 90, color=INK, linewidth=1)
ax.scatter(times_us, np.zeros_like(times_us), color=TEAL, s=70, zorder=3)
for index, time_us in zip(indices, times_us):
    ax.annotate(f"n={index}\n{time_us:.3f} µs", (time_us, 0), xytext=(0, 16), textcoords="offset points", ha="center", fontsize=9)
ax.annotate("Δt = 20.833 µs", xy=(times_us[1], -.02), xytext=(times_us[0], -.24), arrowprops={"arrowstyle": "<->", "color": RED}, color=RED)
ax.set(xlabel="Time (µs)", title="Eq. 5.2: the first five sample instants at 48 kHz", xlim=(-3, 90), ylim=(-.35, .35), yticks=[])
save("eq-5-2-sample-spacing.svg")

# Equation 5.3: strict Nyquist boundary and phase ambiguity at equality.
fs = 8_000
n = np.arange(8)
fig, axes = plt.subplots(2, 1, figsize=(9, 5.5), sharex=True)
axes[0].stem(n, np.sin(2 * np.pi * 4_000 * n / fs), linefmt=RED, markerfmt="o", basefmt="k-")
axes[0].set(ylabel="Sine samples", ylim=(-1.15, 1.15), title="Eq. 5.3: at 8 kHz, 4 kHz is the unsafe equality boundary")
axes[1].stem(n, np.cos(2 * np.pi * 4_000 * n / fs), linefmt=TEAL, markerfmt="s", basefmt="k-")
axes[1].set(xlabel="Sample n", ylabel="Cosine samples", ylim=(-1.15, 1.15))
for ax in axes:
    ax.grid(alpha=.2)
fig.text(.99, .01, "3 kHz: condition met   ·   4 kHz: equality   ·   5 kHz: condition not met", ha="right", color=INK, fontsize=9)
save("eq-5-3-nyquist-boundary.svg")

# Equation 5.4: level count and quantization step.
peak = 1
bits = 3
levels = 2 ** bits
step = 2 * peak / levels
reconstruction = -peak + (np.arange(levels) + .5) * step
fig, ax = plt.subplots(figsize=(9, 4))
ax.hlines(reconstruction, -1, 1, color=TEAL, linewidth=2)
for value in reconstruction:
    ax.text(1.02, value, f"{value:g}", va="center", fontsize=8)
ax.annotate("Δ = 0.25", xy=(-.72, reconstruction[4]), xytext=(-.72, reconstruction[5]), arrowprops={"arrowstyle": "<->", "color": RED}, color=RED, ha="center", va="center")
ax.set(xlabel="Input range", ylabel="Reconstruction level", title="Eq. 5.4: B=3 gives 8 levels across A=±1", xlim=(-1.08, 1.16), ylim=(-1.05, 1.05), xticks=[-1, 0, 1])
ax.grid(axis="x", alpha=.2)
save("eq-5-4-quantization-step.svg")

# Equation 5.5: explicit mid-rise quantizer and bounded in-range error.
def quantize_mid_rise(values, bit_depth, full_scale=1.0):
    level_count = 2 ** bit_depth
    delta = 2 * full_scale / level_count
    code = np.floor((values + full_scale) / delta)
    code = np.clip(code, 0, level_count - 1)
    return -full_scale + (code + .5) * delta

x = np.linspace(-1.2, 1.2, 1_201)
q = quantize_mid_rise(x, bits)
error = x - q
fig, axes = plt.subplots(2, 1, figsize=(9, 6), sharex=True)
axes[0].plot(x, x, color="#8b8178", linestyle="--", label="ideal y=x")
axes[0].step(x, q, where="post", color=TEAL, label="$Q_B(x)$")
axes[0].axvspan(-1.2, -1, color=RED, alpha=.1)
axes[0].axvspan(1, 1.2, color=RED, alpha=.1)
axes[0].set(ylabel="Quantized output", title="Eq. 5.5: 3-bit uniform mid-rise teaching quantizer")
axes[0].legend(loc="upper left")
axes[1].plot(x, error, color=GOLD)
axes[1].axhline(step / 2, color=RED, linestyle="--", label="±Δ/2")
axes[1].axhline(-step / 2, color=RED, linestyle="--")
axes[1].set(xlabel="Input x", ylabel="Error x − Q(x)", ylim=(-.38, .38))
axes[1].legend(loc="upper left")
for ax in axes:
    ax.grid(alpha=.2)
save("eq-5-5-quantization-error.svg")

# Equation 5.6: buffer duration versus final sample instant.
rows = [(1, 48_000), (128, 48_000), (48_000, 48_000), (44_100, 44_100)]
fig, ax = plt.subplots(figsize=(9, 4.4))
labels = []
for row, (frames, rate) in enumerate(rows):
    duration_ms = frames / rate * 1_000
    last_ms = (frames - 1) / rate * 1_000
    ax.barh(row, duration_ms, color=TEAL, alpha=.7)
    ax.plot(last_ms, row, marker="|", color=RED, markersize=16, markeredgewidth=2)
    labels.append(f"N={frames:,}, Fs={rate:,} Hz")
    ax.text(duration_ms, row, f"  {duration_ms:.4g} ms", va="center", fontsize=9)
ax.set(xscale="symlog", xlabel="Buffer duration (ms, symmetric-log scale)", ylabel="Case", yticks=range(len(rows)), yticklabels=labels, title="Eq. 5.6: duration is N/Fs; red marks the last sample instant")
ax.grid(axis="x", alpha=.2)
save("eq-5-6-buffer-duration.svg")

# Equation 5.7: five contiguous render blocks.
fs = 48_000
block = 128
block_ms = block / fs * 1_000
starts = 100 + np.arange(5) * block_ms
fig, ax = plt.subplots(figsize=(9, 3.5))
for index, start in enumerate(starts):
    ax.broken_barh([(start, block_ms)], (.2, .6), facecolors=TEAL if index % 2 == 0 else GOLD)
    ax.text(start + block_ms / 2, .5, f"k={index}\n{start:.3f} ms", ha="center", va="center", color="white", fontsize=8)
ax.set(xlabel="Audio-clock time (ms)", yticks=[], ylim=(0, 1), xlim=(99.5, starts[-1] + block_ms + .5), title="Eq. 5.7: five 128-frame blocks at 48 kHz")
ax.grid(axis="x", alpha=.2)
save("eq-5-7-block-schedule.svg")

assert np.allclose(np.sin(2 * np.pi * 1_000 * np.arange(8) / 8_000), [0, np.sqrt(.5), 1, np.sqrt(.5), 0, -np.sqrt(.5), -1, -np.sqrt(.5)], atol=1e-12)
assert levels == 8 and step == .25
assert np.max(np.abs(error[(x >= -1) & (x < 1)])) <= step / 2 + 1e-12
assert np.isclose(block_ms, 8 / 3)
print("Generated Chapter 5 formula visuals.")