"""Mono 16-bit PCM WAV output and the level mapping from pascals to full scale."""

from __future__ import annotations

import wave
from typing import Literal

import numpy as np

Level = Literal["peak", "loudness", "absolute"]

KNEE = 0.9
PEAK_TARGET = 10 ** (-1 / 20)  # -1 dBFS
LOUDNESS_TARGET = 10 ** (-20 / 20)  # -20 dBFS RMS
LOUDNESS_SKIP_S = 0.05


def level_gain(samples: np.ndarray, fs: float, level: Level = "peak", gain_db: float = 0.0) -> float:
    """Gain that maps a pressure signal [Pa] to full scale: ``"peak"`` normalizes
    the peak to -1 dBFS, ``"loudness"`` the RMS after the first 50 ms to -20 dBFS,
    and ``"absolute"`` maps 1 Pa to 0 dBFS shifted by ``gain_db``."""
    if level == "absolute":
        return 10 ** (gain_db / 20)
    if level == "peak":
        m = float(np.abs(samples).max()) if len(samples) else 0.0
        return PEAK_TARGET / m if m > 0 else 1.0
    if level == "loudness":
        start = min(len(samples), round(LOUDNESS_SKIP_S * fs))
        tail = samples[start:]
        rms = float(np.sqrt((tail.astype(float) ** 2).mean())) if len(tail) else 0.0
        return LOUDNESS_TARGET / rms if rms > 0 else 1.0
    raise ValueError(f"unknown level {level!r}")


def apply_level(samples: np.ndarray, gain: float) -> np.ndarray:
    """Applies the gain and a soft clipper: below 0.9 of full scale the identity,
    above it the excess is squashed with a tanh so that the output stays within
    (-1, 1) with a continuous slope at the knee."""
    y = samples.astype(np.float64) * gain
    a = np.abs(y)
    room = 1 - KNEE
    clipped = KNEE + room * np.tanh((a - KNEE) / room)
    return np.where(a <= KNEE, y, np.sign(y) * clipped).astype(np.float32)


def write_wav(path, samples: np.ndarray, fs: int) -> None:
    """Writes samples in [-1, 1] as a 16-bit PCM mono WAV file."""
    pcm = np.round(np.clip(np.nan_to_num(samples), -1, 1) * 32767).astype("<i2")
    with wave.open(str(path), "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(int(round(fs)))
        w.writeframes(pcm.tobytes())


def peak_dbfs(samples: np.ndarray) -> float:
    m = float(np.abs(samples).max()) if len(samples) else 0.0
    return 20 * np.log10(m) if m > 0 else -np.inf