waveform.py

"""Waveform PNG with section markers — rebuilt to match v2 design spec.

Bar-style waveform on a dark panel, each bar coloured by the section it
falls inside, semantic section labels with a coloured top-border below.
The palette mirrors `design/README.md`:
    Intro = mint-deep `#2BB89E`
    Build = mint      `#5BE0C8`
    Core  = coral     `#FF6A3D`
    Outro = coral-bright `#FF8C5A`
For >4 sections, the bracketing rule is: first→Intro, last→Outro, the
middle ones cycle Build→Core→Build so the centre carries the most
energetic colour. <2 sections falls back to a single Core span.
"""
from __future__ import annotations

import os
import tempfile
from typing import Any

import numpy as np

from pipeline import Analysis


# Design-token colors — keep in sync with theme.py / design/README.md.
INK_BG_PANEL = "#121419"     # outer panel
INK_BG_CARD = "#0F1115"      # inside the waveform area
INK_BORDER = "#20242C"
INK_LABEL = "#5E6671"
INK_TIME = "#99A0AB"
INK_PLAYHEAD = "#F2EFE9"

CORAL = "#FF6A3D"
CORAL_BRIGHT = "#FF8C5A"
MINT = "#5BE0C8"
MINT_DEEP = "#2BB89E"
AMBER = "#FFC24B"

# Semantic ordering — first → Intro, last → Outro, centre alternates.
SECTION_COLORS_BY_ROLE = {
    "intro": MINT_DEEP,
    "build": MINT,
    "core": CORAL,
    "outro": CORAL_BRIGHT,
}


def _assign_roles(n: int) -> list[str]:
    """Pick a role per section index given a total count `n`. Always:
      - section 0 → intro
      - last section → outro
      - middle sections cycle build → core → build so the centre is coral.
    """
    if n <= 0:
        return []
    if n == 1:
        return ["core"]
    if n == 2:
        return ["intro", "outro"]
    if n == 3:
        return ["intro", "core", "outro"]
    roles = ["intro"]
    middle = n - 2
    for i in range(middle):
        # Pattern: build, core, build, core, … with core landing at the
        # middle so a 5-section piece reads Intro/Build/Core/Build/Outro.
        if (i + (1 if middle % 2 == 0 else 0)) % 2 == 0:
            roles.append("build")
        else:
            roles.append("core")
    # Force a single core at the centre when even-numbered middles default
    # would skip it. E.g. middle=2 → [build, build]; we want [build, core].
    if "core" not in roles:
        mid_idx = 1 + middle // 2
        roles[mid_idx] = "core"
    roles.append("outro")
    return roles


def render(a: Analysis) -> str | None:
    """Write a section-coloured bar waveform PNG. Returns path or None
    if rendering can't run (missing librosa / matplotlib)."""
    try:
        import librosa
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import matplotlib.patches as patches
    except Exception:
        return None

    try:
        y, sr = librosa.load(a.source_path, sr=22050, mono=True)
    except Exception:
        return None

    duration = len(y) / sr if sr else 1.0
    if duration <= 0:
        return None

    # Bar-style waveform: chunk into N evenly-spaced bins and take the
    # peak amplitude per bin. Density tuned so bars feel light at 15s
    # (~6 bars/s) but stay readable at 180s (~1 bar/s).
    n_bars = max(120, min(360, int(duration * 6)))
    samples_per_bar = max(1, len(y) // n_bars)
    bars = np.array([
        np.max(np.abs(y[i * samples_per_bar:(i + 1) * samples_per_bar]))
        for i in range(n_bars)
    ])
    # Normalise to a stable visual range; tiny clips don't go invisible.
    peak = bars.max() if bars.size else 1.0
    if peak > 0:
        bars = bars / peak
    bars = np.clip(bars, 0.02, 1.0)  # min height so silent regions still show a tick

    # Assign roles per section + map back to per-bar colour.
    sections: list[dict[str, Any]] = list(a.sections or [])
    roles = _assign_roles(len(sections))
    section_colors = [SECTION_COLORS_BY_ROLE[r] for r in roles]
    if not sections:
        bar_colors = [CORAL] * n_bars
    else:
        bar_colors = []
        for i in range(n_bars):
            t = (i + 0.5) / n_bars * duration
            picked = section_colors[0]
            for s, c in zip(sections, section_colors):
                if t >= s["start"] and t < s["end"]:
                    picked = c
                    break
                if t >= s["start"]:
                    picked = c
            bar_colors.append(picked)

    # Figure: 1200 × 240 with the waveform up top and a compact label
    # row below carrying section name + time range under coloured rules.
    fig_w_in, fig_h_in = 11.0, 2.6
    dpi = 110
    fig = plt.figure(figsize=(fig_w_in, fig_h_in), dpi=dpi,
                     facecolor=INK_BG_PANEL)
    # Main waveform axis (top ~75%); label strip axis (bottom ~25%).
    ax_wave = fig.add_axes([0.025, 0.30, 0.95, 0.62], facecolor=INK_BG_CARD)
    ax_labels = fig.add_axes([0.025, 0.05, 0.95, 0.22], facecolor=INK_BG_PANEL)

    # Bar widths in data coords — each spans 1 / n_bars of duration with
    # a small gap so the bars read individually rather than as a block.
    bar_w = duration / n_bars * 0.78
    centers = (np.arange(n_bars) + 0.5) / n_bars * duration
    # Mirror about y=0 to draw the classic symmetrical waveform.
    for c, x, color in zip(bars, centers, bar_colors):
        ax_wave.add_patch(patches.Rectangle(
            (x - bar_w / 2, -c), bar_w, 2 * c, color=color, linewidth=0,
        ))
    ax_wave.set_xlim(0, duration)
    ax_wave.set_ylim(-1.08, 1.08)
    ax_wave.set_xticks([])
    ax_wave.set_yticks([])
    for spine in ax_wave.spines.values():
        spine.set_edgecolor(INK_BORDER)
        spine.set_linewidth(0.8)

    # Section label strip — coloured top-border per section + name + time.
    ax_labels.set_xlim(0, duration)
    ax_labels.set_ylim(0, 1)
    ax_labels.set_xticks([])
    ax_labels.set_yticks([])
    for spine in ax_labels.spines.values():
        spine.set_visible(False)
    for s, color, role in zip(sections, section_colors, roles):
        start = float(s["start"])
        end = float(s["end"])
        # 2px colored top-border drawn as a thin filled rect.
        ax_labels.add_patch(patches.Rectangle(
            (start, 0.85), end - start, 0.06, color=color, linewidth=0,
        ))
        mid = (start + end) / 2.0
        # Role label (uppercase, mono feel via monospace family).
        ax_labels.text(
            mid, 0.55, role.upper(),
            ha="center", va="center",
            color=color, fontsize=8.5, fontweight="bold",
            family="monospace",
        )
        # Time range under the role.
        ax_labels.text(
            mid, 0.18, f"{start:0.1f}s – {end:0.1f}s",
            ha="center", va="center",
            color=INK_LABEL, fontsize=7.0, family="monospace",
        )

    out = os.path.join(a.workdir or tempfile.gettempdir(), "waveform.png")
    fig.savefig(out, facecolor=INK_BG_PANEL, edgecolor="none",
                bbox_inches="tight", pad_inches=0.18)
    plt.close(fig)
    return out