"""Turn drawn strokes into the prosody features a speech model expects.

The transcription models in `finetune_whisper` were trained on two signals
extracted from real speech at 200 Hz -- interpolated F0 and intensity -- which
is exactly what a stylus produces: vertical position is pitch, pressure is
loudness. This module resamples strokes onto that 200 Hz grid and applies the
same normalisation the training pipeline used, so a drawing can be fed to the
model as if it were speech.

The constants here are not free choices. They mirror the training pipeline's
own FEATURE_SCALES and the checkpoint's `model_config.json`; if those change,
these must change with them. This module is a verbatim port of the reference
implementation's `features.py`, minus code paths the demo never takes, so a
drawing is turned into model input exactly as in the paper.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from .f0_interp import interpolate_f0

# Prosody features are extracted at 200 Hz (time_step 0.005 s) and the mel
# target runs at 100 Hz; the model's length_ratio=0.5 encodes that 2:1.
FEATURE_RATE_HZ = 200.0

# Physical-range normalisation, copied from FeaturesToMelDataset.FEATURE_SCALES.
F0_SCALE = 500.0  # Hz, human pitch range ~50-500
INTENSITY_SCALE = 100.0  # dB SPL, typical range 0-100

# Pen pressure is a 0..1 ratio; speech intensity is dB SPL. These bounds are
# measured, not chosen: across the extracted templates (2600 frames of real
# speech from the training distribution) intensity runs p5=51.8, p50=68.6,
# p95=81.0 dB. Mapping 0..1 onto 50..85 dB covers that spread and puts the
# speech median at pressure 0.53, so a natural mid-pressure stroke produces
# typical speech loudness.
#
# The earlier 45-80 dB guess put the median at 0.67 and clipped the top: a
# comfortable touch read ~6 dB quieter than real speech, and the loudest
# frames (84.5 dB) could not be reached at all.
#
# Re-derive with: uv run python scripts/calibrate_intensity.py
INTENSITY_FLOOR_DB = 50.0
INTENSITY_CEIL_DB = 85.0

# Where speech actually lives. These are measured from the extracted
# templates (2600 frames of real speech from the training distribution)
# rather than taken from textbook adult F0 ranges: p1 is MIDI 38.7 and p99.5
# is MIDI 64.9, so real contours reach well outside the "average adult"
# 85-255 Hz band.
#
# An earlier F2-C#4 (41-61) surface clipped 4.8% of frames, and the clipped
# frames are the pitch extremes that carry the most information. B1-G4 clips
# none of the measured 2600.
#
# Re-derive with: uv run python scripts/trace_fidelity.py (which reports
# clipping) after rebuilding templates.json.
SPEECH_MIDI_LOW = 35  # B1,  61.7 Hz
SPEECH_MIDI_HIGH = 67  # G4, 392.0 Hz

# The widest surface the model can be drawn on at all, as opposed to the
# range real speech occupies. F0_SCALE is 500 Hz, which is MIDI 71 -- above
# it the normalised input exceeds 1.0, a value training never contained.
# Downward there is no such wall: a low pitch simply normalises small.
#
# Drawing out here is out of distribution and will transcribe worse, but it
# is a usable surface rather than a broken one, and a wide sweep is sometimes
# what the drawing wants.
# Extended downward rather than upward. The ceiling is a hard wall -- MIDI 71
# is 494 Hz and F0_SCALE is 500, so there is almost nothing above it -- while
# downward a low pitch simply normalises small. Logged drawings sit at a p99
# of MIDI 58, so headroom above 60 is canvas nobody uses; the room worth
# having is underneath.
DRAWABLE_MIDI_LOW = 12  # C0,  16.4 Hz
DRAWABLE_MIDI_HIGH = 71  # B4, 493.9 Hz


def midi_to_hz(midi: float | np.ndarray) -> np.ndarray:
    return 440.0 * 2.0 ** ((np.asarray(midi, dtype=np.float64) - 69.0) / 12.0)


@dataclass
class ProsodyFeatures:
    """f0 and intensity on the model's 200 Hz grid.

    `raw` holds physical units (Hz, dB) for inspection; `normalised` is what
    the model consumes.
    """

    f0_hz: np.ndarray
    intensity_db: np.ndarray

    @property
    def duration_s(self) -> float:
        return len(self.f0_hz) / FEATURE_RATE_HZ

    @property
    def num_frames(self) -> int:
        return len(self.f0_hz)

    def normalised(self) -> np.ndarray:
        """[T, 2] float32 array scaled exactly as in training."""
        f0 = np.asarray(self.f0_hz, dtype=np.float32) / F0_SCALE
        intensity = np.asarray(self.intensity_db, dtype=np.float32) / INTENSITY_SCALE
        return np.stack([f0, intensity], axis=-1).astype(np.float32)

def pressure_to_db(
    pressure: np.ndarray,
    floor_db: float = INTENSITY_FLOOR_DB,
    ceil_db: float = INTENSITY_CEIL_DB,
) -> np.ndarray:
    """Map 0..1 stylus pressure onto a speech-like dB SPL range."""
    p = np.clip(np.asarray(pressure, dtype=np.float64), 0.0, 1.0)
    return floor_db + (ceil_db - floor_db) * p


def resample_stroke(
    times_s: np.ndarray,
    values: np.ndarray,
    rate_hz: float = FEATURE_RATE_HZ,
) -> tuple[np.ndarray, np.ndarray]:
    """Resample irregular pen samples onto a uniform grid.

    Tablet events arrive at an irregular rate, so the raw samples cannot be
    used directly: the model assumes a fixed 200 Hz step. Returns the grid
    times and the linearly interpolated values.
    """
    times_s = np.asarray(times_s, dtype=np.float64)
    values = np.asarray(values, dtype=np.float64)
    if times_s.size == 0:
        return np.zeros(0), np.zeros(0)
    if times_s.size == 1:
        return times_s.copy(), values.copy()

    # Pen events are timestamped by arrival and can repeat or (rarely) go
    # backwards; np.interp needs a strictly increasing x.
    order = np.argsort(times_s, kind="stable")
    times_s, values = times_s[order], values[order]
    keep = np.concatenate(([True], np.diff(times_s) > 0))
    times_s, values = times_s[keep], values[keep]
    if times_s.size == 1:
        return times_s.copy(), values.copy()

    duration = times_s[-1] - times_s[0]
    n = max(2, int(round(duration * rate_hz)) + 1)
    grid = times_s[0] + np.arange(n) / rate_hz
    # Clamp so the last grid point never extrapolates past the final sample.
    grid = np.minimum(grid, times_s[-1])
    return grid, np.interp(grid, times_s, values)


def smooth(values: np.ndarray, window_frames: int) -> np.ndarray:
    """Moving-average smoothing with edge padding.

    Drawn lines are geometrically smooth but can still have corners that no
    speech F0 contour would contain. A short smoothing window keeps the
    contour closer to what the model saw in training.
    """
    if window_frames <= 1 or values.size == 0:
        return np.asarray(values, dtype=np.float64)
    window_frames = min(window_frames, values.size)
    kernel = np.ones(window_frames) / window_frames
    pad = window_frames // 2
    padded = np.pad(np.asarray(values, dtype=np.float64), pad, mode="edge")
    out = np.convolve(padded, kernel, mode="same")
    return out[pad : pad + values.size]


def features_from_stroke(
    times_s: np.ndarray,
    midi: np.ndarray,
    pressure: np.ndarray,
    smooth_ms: float = 25.0,
    rate_hz: float = FEATURE_RATE_HZ,
) -> ProsodyFeatures:
    """Build 200 Hz prosody features from one pen stroke.

    `times_s` are event timestamps in seconds, `midi` the pitch under the pen
    (fractional MIDI note), `pressure` the stylus pressure in 0..1.

    Pitch is read as a literal note and converted to Hz, so the drawing
    surface must already sit in the range speech occupies -- see
    SPEECH_MIDI_LOW/HIGH and the roll bounds the app is built with.
    """
    times_s = np.asarray(times_s, dtype=np.float64)
    midi = np.asarray(midi, dtype=np.float64)
    pressure = np.asarray(pressure, dtype=np.float64)
    if not (times_s.size == midi.size == pressure.size):
        raise ValueError("times, midi and pressure must be the same length")
    if times_s.size == 0:
        return ProsodyFeatures(np.zeros(0, np.float64), np.zeros(0, np.float64))

    grid, midi_grid = resample_stroke(times_s, midi, rate_hz)
    _, pressure_grid = resample_stroke(times_s, pressure, rate_hz)
    # resample_stroke derives its grid from the data, so both calls agree in
    # length; guard anyway rather than silently mis-aligning the channels.
    n = min(len(midi_grid), len(pressure_grid))
    midi_grid, pressure_grid = midi_grid[:n], pressure_grid[:n]

    f0_raw = midi_to_hz(midi_grid)

    window = max(1, int(round(smooth_ms / 1000.0 * rate_hz)))
    f0_hz = smooth(f0_raw, window)
    intensity_db = smooth(pressure_to_db(pressure_grid), window)

    return ProsodyFeatures(f0_hz=f0_hz, intensity_db=intensity_db)


# Quiet frames around a phrase, so it neither starts nor stops dead.
#
# Measured, not guessed: across the templates the quiet before and after the
# speech has a median of 0 ms and a maximum of 80 ms -- these dataset clips
# are trimmed tight to the utterance, so training really does start at
# speech. An earlier 250 ms on each side was 3x the observed maximum and
# contributed more frames than the drawing itself on a short phrase.
#
# 40 ms keeps the onset from being a step discontinuity without displacing
# content; it sits inside the observed range rather than beyond it.
LEAD_SILENCE_S = 0.04
TAIL_SILENCE_S = 0.04


def pad_with_silence(
    features: ProsodyFeatures,
    lead_s: float = LEAD_SILENCE_S,
    tail_s: float = TAIL_SILENCE_S,
    rate_hz: float = FEATURE_RATE_HZ,
) -> ProsodyFeatures:
    """Surround a phrase with quiet frames so it neither starts nor stops dead.

    Pitch is held at the adjacent value rather than zeroed, matching how
    f0_interp behaves through unvoiced regions, and intensity sits at the
    floor.
    """
    if features.num_frames == 0:
        return features

    lead = max(0, int(round(lead_s * rate_hz)))
    tail = max(0, int(round(tail_s * rate_hz)))
    if not lead and not tail:
        return features

    f0_parts, db_parts = [], []
    if lead:
        f0_parts.append(np.full(lead, features.f0_hz[0]))
        db_parts.append(np.full(lead, INTENSITY_FLOOR_DB))
    f0_parts.append(features.f0_hz)
    db_parts.append(features.intensity_db)
    if tail:
        f0_parts.append(np.full(tail, features.f0_hz[-1]))
        db_parts.append(np.full(tail, INTENSITY_FLOOR_DB))

    return ProsodyFeatures(
        f0_hz=np.concatenate(f0_parts),
        intensity_db=np.concatenate(db_parts),
    )


def concatenate(
    features: list[ProsodyFeatures],
    gap_s: float | list[float] = 0.12,
    rate_hz: float = FEATURE_RATE_HZ,
    lead_s: float = LEAD_SILENCE_S,
    tail_s: float = TAIL_SILENCE_S,
) -> ProsodyFeatures:
    """Join strokes into one utterance, with a quiet gap between them.

    Each stroke is a "note"; the gaps are what give the model something it can
    read as separate syllables rather than one continuous tone.

    `gap_s` may be a single duration or one per gap. Per-gap is what the
    canvas passes: the pen's actual lift time is recorded, so there is no
    reason to substitute a constant for it. A fixed gap also compounds -- at
    120 ms across the 3-8 strokes a phrase typically has, the inserted silence
    can outweigh the drawing itself.

    Across a gap the pitch is **interpolated by pYAAPT's own routine**, not
    held and not linearly ramped. The conformer was trained on f0_interp,
    which median filters and fits a shape-preserving PCHIP spline through the
    unvoiced regions; holding the previous value gives a flat shelf then a
    step, and a straight line gives a shape speech never has. Gap frames are
    marked unvoiced (f0 = 0) and handed to interpolate_f0, so they are filled
    the same way training filled them. Intensity drops to the floor, which is
    what makes the gap read as a pause.
    """
    usable = [f for f in features if f.num_frames > 0]
    if not usable:
        return ProsodyFeatures(np.zeros(0, np.float64), np.zeros(0, np.float64))

    if isinstance(gap_s, (int, float)):
        gaps = [float(gap_s)] * max(0, len(usable) - 1)
    else:
        gaps = [float(g) for g in gap_s]
        if len(gaps) < len(usable) - 1:
            gaps += [0.12] * (len(usable) - 1 - len(gaps))

    f0_parts: list[np.ndarray] = []
    db_parts: list[np.ndarray] = []

    for i, feat in enumerate(usable):
        gap_frames = max(0, int(round(gaps[i - 1] * rate_hz))) if i > 0 else 0
        if i > 0 and gap_frames:
            # Zero marks the gap as unvoiced; interpolate_f0 fills it below,
            # exactly as yaapt filled unvoiced regions in training.
            f0_parts.append(np.zeros(gap_frames))
            db_parts.append(np.full(gap_frames, INTENSITY_FLOOR_DB))
        f0_parts.append(np.asarray(feat.f0_hz, dtype=np.float64))
        db_parts.append(np.asarray(feat.intensity_db, dtype=np.float64))

    joined = ProsodyFeatures(
        f0_hz=interpolate_f0(np.concatenate(f0_parts)),
        intensity_db=np.concatenate(db_parts),
    )
    return pad_with_silence(joined, lead_s, tail_s, rate_hz)