src/midnight_static/mixer.py

"""Midnight Static mixer.

`mix_crude_broadcast` is the Day 1 minimum stack (dialogue + bed + one SFX).
`mix_broadcast` is the production chain from AUDIO_PIPELINE.md: ducking, an
AM-radio band-pass with saturation, vinyl crackle, BS.1770-style loudness
normalization, and MP3 export. Everything is dependency-free numpy DSP so it
runs on the CPU Space without scipy.
"""

from __future__ import annotations

import math
import wave
from pathlib import Path

import numpy as np

SAMPLE_RATE = 24_000

# AM/telephone broadcast band and the integrated-loudness target for mastering.
AM_BAND_LOW_HZ = 280.0
AM_BAND_HIGH_HZ = 3400.0
TARGET_LUFS = -14.0  # AUDIO_PIPELINE.md master target
# Period AM voice chain (band-pass + saturation). On by default; demo with it ON.
AUTHENTIC_AM = True


def mix_crude_broadcast(
    dialogue_wav: Path,
    output_path: Path,
    *,
    sfx_layers: list[tuple[np.ndarray, float]] | None = None,
    bed: np.ndarray | None = None,
    opening_sting: np.ndarray | None = None,
    closing_sting: np.ndarray | None = None,
) -> Path:
    """Mix dialogue with a music bed and SFX into a mono WAV.

    Bed/SFX come from the matched asset library when provided (`bed`,
    `opening_sting`, `closing_sting`, `sfx_layers` of ``(audio, offset_seconds)``);
    otherwise a synthetic bed and hardcoded SFX are used (the station improvises).
    Only the bed/SFX *source* changes — the timing and ducking are unchanged.
    """
    dialogue = read_mono_wav(dialogue_wav)
    total_frames = len(dialogue) + int(SAMPLE_RATE * 2.0)
    bed_track = _bed_track(total_frames, bed, opening_sting, closing_sting, len(dialogue))
    sfx = _sfx_track(total_frames, sfx_layers, len(dialogue))

    dialogue_track = np.zeros(total_frames, dtype=np.float32)
    dialogue_offset = int(SAMPLE_RATE * 0.8)
    _overlay_in_place(dialogue_track, dialogue, dialogue_offset)

    # Manual ducking: lower the bed while dialogue is present.
    envelope = np.where(np.abs(dialogue_track) > 0.006, 0.22, 0.55).astype(np.float32)
    envelope = _smooth(envelope, int(SAMPLE_RATE * 0.12))
    mix = dialogue_track * 0.95 + bed_track * envelope + sfx * 0.5
    mix = _normalize_peak(mix, peak=0.92)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    write_mono_wav(output_path, mix)
    return output_path


def _tile_to(audio: np.ndarray, frames: int) -> np.ndarray:
    """Loop a clip to fill `frames` samples (for a bed shorter than the show)."""
    if len(audio) == 0:
        return np.zeros(frames, dtype=np.float32)
    reps = int(np.ceil(frames / len(audio)))
    return np.tile(audio.astype(np.float32), reps)[:frames]


def _bed_track(
    total_frames: int,
    bed: np.ndarray | None,
    opening_sting: np.ndarray | None,
    closing_sting: np.ndarray | None,
    dialogue_len: int,
) -> np.ndarray:
    """Music track from real library clips, or the synthetic bed when none given."""
    if bed is None and opening_sting is None and closing_sting is None:
        return _music_bed(total_frames)
    track = np.zeros(total_frames, dtype=np.float32)
    if bed is not None:
        track += _tile_to(bed, total_frames) * 0.6
    if opening_sting is not None:
        _overlay_in_place(track, opening_sting * 0.85, 0)
    if closing_sting is not None:
        _overlay_in_place(track, closing_sting * 0.85, min(total_frames - 1, dialogue_len))
    return _fade(track, fade_in=0.4, fade_out=1.2)


def _sfx_track(
    total_frames: int,
    sfx_layers: list[tuple[np.ndarray, float]] | None,
    dialogue_len: int,
) -> np.ndarray:
    """SFX track from matched clips at their offsets, or the synthetic hits."""
    track = np.zeros(total_frames, dtype=np.float32)
    if sfx_layers is None:
        _overlay_in_place(track, _static_burst(0.75), int(SAMPLE_RATE * 0.25))
        _overlay_in_place(track, _impact_sfx(0.85), max(0, int(dialogue_len * 0.52)))
        return track
    for audio, offset_seconds in sfx_layers:
        _overlay_in_place(track, np.asarray(audio, dtype=np.float32), max(0, int(SAMPLE_RATE * offset_seconds)))
    return track


def read_mono_wav(path: Path) -> np.ndarray:
    with wave.open(str(path), "rb") as wav_file:
        channels = wav_file.getnchannels()
        sample_rate = wav_file.getframerate()
        sample_width = wav_file.getsampwidth()
        if sample_rate != SAMPLE_RATE:
            raise ValueError(f"expected {SAMPLE_RATE} Hz WAV, got {sample_rate} Hz")
        if sample_width != 2:
            raise ValueError(f"expected 16-bit WAV, got sample width {sample_width}")
        raw = wav_file.readframes(wav_file.getnframes())
    data = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
    if channels == 1:
        return data
    if channels == 2:
        return data.reshape(-1, 2).mean(axis=1)
    raise ValueError(f"unsupported channel count: {channels}")


def write_mono_wav(path: Path, samples: np.ndarray) -> None:
    clipped = np.clip(samples, -1.0, 1.0)
    pcm = (clipped * 32767).astype("<i2")
    with wave.open(str(path), "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(SAMPLE_RATE)
        wav_file.writeframes(pcm.tobytes())


def _music_bed(frames: int) -> np.ndarray:
    seconds = np.arange(frames, dtype=np.float32) / SAMPLE_RATE
    fundamental = np.sin(2 * math.pi * 82.41 * seconds)
    fifth = np.sin(2 * math.pi * 123.47 * seconds)
    shimmer = np.sin(2 * math.pi * 246.94 * seconds) * 0.35
    bed = (fundamental * 0.5 + fifth * 0.35 + shimmer * 0.15).astype(np.float32)
    return _fade(bed * 0.18, fade_in=1.2, fade_out=1.8)


def _static_burst(seconds: float) -> np.ndarray:
    frames = int(SAMPLE_RATE * seconds)
    rng = np.random.default_rng(1948)
    noise = rng.uniform(-1.0, 1.0, frames).astype(np.float32)
    return _fade(noise * 0.25, fade_in=0.02, fade_out=0.45)


def _impact_sfx(seconds: float) -> np.ndarray:
    frames = int(SAMPLE_RATE * seconds)
    t = np.arange(frames, dtype=np.float32) / SAMPLE_RATE
    tone = np.sin(2 * math.pi * 58.0 * t) + np.sin(2 * math.pi * 116.0 * t) * 0.45
    decay = np.exp(-5.0 * t)
    return (tone * decay * 0.45).astype(np.float32)


def _overlay_in_place(base: np.ndarray, layer: np.ndarray, offset: int) -> None:
    if offset >= len(base):
        return
    end = min(len(base), offset + len(layer))
    base[offset:end] += layer[: end - offset]


def _fade(samples: np.ndarray, fade_in: float, fade_out: float) -> np.ndarray:
    out = samples.copy()
    in_frames = min(len(out), int(SAMPLE_RATE * fade_in))
    out_frames = min(len(out), int(SAMPLE_RATE * fade_out))
    if in_frames:
        out[:in_frames] *= np.linspace(0.0, 1.0, in_frames, dtype=np.float32)
    if out_frames:
        out[-out_frames:] *= np.linspace(1.0, 0.0, out_frames, dtype=np.float32)
    return out


def _smooth(values: np.ndarray, window: int) -> np.ndarray:
    if window <= 1:
        return values
    kernel = np.ones(window, dtype=np.float32) / window
    return np.convolve(values, kernel, mode="same").astype(np.float32)


def _normalize_peak(samples: np.ndarray, peak: float) -> np.ndarray:
    current = float(np.max(np.abs(samples))) if len(samples) else 0.0
    if current <= 0:
        return samples
    return samples * min(1.0, peak / current)


# --------------------------------------------------------------------------- #
# Production mixer chain (AUDIO_PIPELINE.md)
# --------------------------------------------------------------------------- #


def mix_broadcast(
    dialogue_wav: Path,
    output_path: Path,
    *,
    mp3: bool = True,
    target_lufs: float = TARGET_LUFS,
    authentic_am: bool = AUTHENTIC_AM,
    sfx_layers: list[tuple[np.ndarray, float]] | None = None,
    bed: np.ndarray | None = None,
    opening_sting: np.ndarray | None = None,
    closing_sting: np.ndarray | None = None,
) -> Path:
    """Master a dialogue WAV into a finished broadcast.

    Chain: bed + SFX with ducking under dialogue → AM band-pass + saturation →
    vinyl crackle → loudness normalization → MP3 (or WAV) export. Bed/SFX come
    from the matched asset library when provided (else synthetic). Returns the
    path actually written (``.mp3`` when the encoder is available, else ``.wav``).
    """
    dialogue = read_mono_wav(dialogue_wav)
    total_frames = len(dialogue) + int(SAMPLE_RATE * 2.0)

    dialogue_track = np.zeros(total_frames, dtype=np.float32)
    _overlay_in_place(dialogue_track, dialogue, int(SAMPLE_RATE * 0.8))

    bed_track = _bed_track(total_frames, bed, opening_sting, closing_sting, len(dialogue))
    sfx = _sfx_track(total_frames, sfx_layers, len(dialogue))

    # Duck bed and SFX while dialogue is present.
    duck = _ducking_envelope(dialogue_track)
    program = dialogue_track * 0.95 + bed_track * duck * 0.6 + sfx * 0.5 * np.clip(duck + 0.3, 0, 1)

    # AM radio character (AUTHENTIC_AM toggle), then analog noise floor.
    if authentic_am:
        program = apply_am_radio(program)
    program = program + vinyl_crackle(total_frames)

    # Master: loudness-normalize, then guard the true peak.
    program = normalize_loudness(program, target_lufs=target_lufs)
    program = _normalize_peak(program, peak=0.97)

    return export_audio(program, output_path, mp3=mp3)


def _ducking_envelope(dialogue_track: np.ndarray) -> np.ndarray:
    """Smooth gain that drops the bed/SFX under dialogue."""
    envelope = np.where(np.abs(dialogue_track) > 0.006, 0.28, 1.0).astype(np.float32)
    return _smooth(envelope, int(SAMPLE_RATE * 0.12))


def apply_am_radio(signal: np.ndarray, drive: float = 2.2) -> np.ndarray:
    """Band-pass to the AM voice band and soft-saturate for tube warmth."""
    banded = bandpass(signal, AM_BAND_LOW_HZ, AM_BAND_HIGH_HZ)
    if drive > 0:
        banded = np.tanh(banded * drive) / float(np.tanh(drive))
    return banded.astype(np.float32)


def bandpass(
    signal: np.ndarray,
    low_hz: float,
    high_hz: float,
    transition_hz: float = 120.0,
) -> np.ndarray:
    """Zero-phase FFT band-pass with raised-cosine edges (no scipy needed)."""
    n = len(signal)
    if n == 0:
        return signal
    spectrum = np.fft.rfft(signal)
    freqs = np.fft.rfftfreq(n, d=1.0 / SAMPLE_RATE)
    mask = _band_mask(freqs, low_hz, high_hz, transition_hz)
    return np.fft.irfft(spectrum * mask, n=n).astype(np.float32)


def _band_mask(
    freqs: np.ndarray, low_hz: float, high_hz: float, transition_hz: float
) -> np.ndarray:
    mask = np.ones_like(freqs, dtype=np.float32)
    mask[freqs < low_hz] = 0.0
    mask[freqs > high_hz] = 0.0
    if transition_hz > 0:
        # Raised-cosine ramps to suppress ringing at the band edges.
        lo = (freqs >= low_hz) & (freqs < low_hz + transition_hz)
        hi = (freqs <= high_hz) & (freqs > high_hz - transition_hz)
        mask[lo] = 0.5 - 0.5 * np.cos(np.pi * (freqs[lo] - low_hz) / transition_hz)
        mask[hi] = 0.5 - 0.5 * np.cos(np.pi * (high_hz - freqs[hi]) / transition_hz)
    return mask


def vinyl_crackle(frames: int, seed: int = 1955) -> np.ndarray:
    """Sparse pops plus a quiet hiss floor for an old-record character."""
    rng = np.random.default_rng(seed)
    out = (rng.standard_normal(frames).astype(np.float32)) * 0.0016  # hiss
    pop_count = max(1, int(frames / SAMPLE_RATE * 7))  # ~7 pops/sec
    positions = rng.integers(0, frames, size=pop_count)
    amplitudes = rng.uniform(0.05, 0.22, size=pop_count).astype(np.float32)
    decay = np.exp(-np.arange(64, dtype=np.float32) / 6.0)
    for pos, amp in zip(positions, amplitudes):
        end = min(frames, pos + decay.size)
        out[pos:end] += amp * decay[: end - pos] * rng.choice([-1.0, 1.0])
    return out


def normalize_loudness(signal: np.ndarray, target_lufs: float = TARGET_LUFS) -> np.ndarray:
    """Scale to a target integrated loudness (BS.1770-*style* approximation).

    Gated mean-square loudness over 400 ms blocks after a high-pass pre-filter;
    full K-weighting is omitted to stay dependency-free. This is NOT `pyloudnorm`
    / a certified LUFS meter — verify the real integrated loudness by measuring a
    rendered broadcast.
    """
    loudness = measure_loudness(signal)
    if loudness is None:
        return signal
    gain = 10.0 ** ((target_lufs - loudness) / 20.0)
    gain = float(np.clip(gain, 0.05, 20.0))
    return (signal * gain).astype(np.float32)


def measure_loudness(signal: np.ndarray) -> float | None:
    """Gated loudness in LUFS-ish units, or None for digital silence."""
    if len(signal) == 0:
        return None
    weighted = bandpass(signal, 60.0, min(AM_BAND_HIGH_HZ * 1.4, SAMPLE_RATE / 2 - 100), 80.0)
    block = int(SAMPLE_RATE * 0.4)
    if block <= 0 or len(weighted) < block:
        mean_square = float(np.mean(weighted**2))
    else:
        blocks = weighted[: len(weighted) // block * block].reshape(-1, block)
        powers = np.mean(blocks**2, axis=1)
        gate = powers > (10.0 ** (-70.0 / 10.0))  # absolute -70 LUFS gate
        kept = powers[gate]
        mean_square = float(np.mean(kept)) if kept.size else float(np.mean(powers))
    if mean_square <= 0:
        return None
    return -0.691 + 10.0 * math.log10(mean_square)


def export_audio(samples: np.ndarray, output_path: Path, *, mp3: bool = True) -> Path:
    """Write the master as MP3 when an encoder is available, else WAV."""
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    clipped = np.clip(samples, -1.0, 1.0).astype(np.float32)
    if mp3 and output_path.suffix.lower() == ".mp3":
        try:
            import soundfile as sf

            sf.write(str(output_path), clipped, SAMPLE_RATE, format="MP3")
            return output_path
        except Exception:
            # libsndfile without MPEG support: fall back to a sibling WAV.
            output_path = output_path.with_suffix(".wav")
    write_mono_wav(output_path, clipped)
    return output_path