"""Extract prosody features from recorded audio.

Audio in, the two signals the conformer consumes out: interpolated f0 and
intensity at 200 Hz. This is the front of the pipeline the paper reports --
everything after it (conformer to mel, mel to words) is unchanged whatever
the audio was.

Ported from the reference implementation's `speech.py`, which in turn ports
the training preprocessor. The settings are not interchangeable with
sensible-looking alternatives: frame_length, nccf_thresh1, the zero-padding,
the six-frame drop and the two clamps all move the contour, and the
conformer only ever saw the combination below.
"""

from __future__ import annotations

import numpy as np

from .features import FEATURE_RATE_HZ, ProsodyFeatures

# Whisper's rate, and what the preprocessor ran at.
TARGET_SAMPLE_RATE = 16000

# pYAAPT settings from the training preprocessor's own config. Every one of
# these differed from the obvious default, and the conformer only ever saw
# contours produced with these.
FRAME_LENGTH_MS = 20.0
FRAME_SPACE_MS = 5.0          # 200 Hz
TDA_FRAME_LENGTH_MS = 25.0
NCCF_THRESH1 = 0.25
# f0_min/f0_max are deliberately not passed: the preprocessor leaves yaapt's
# own defaults (60/400 Hz) in place. Passing 50/500 widened the search and
# changed which frames came back voiced.

# The preprocessor drops the first six frames, which the zero-padding below
# makes spurious.
LEADING_FRAMES_DROPPED = 6
# Clamped after interpolation, as the preprocessor does.
F0_CEILING_HZ = 500.0
INTENSITY_CLAMP_DB = 20.0

# The conformer was exported with a ten-second window, so anything past this
# is truncated rather than transcribed.
MAX_AUDIO_S = 10.0


def resample(audio: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
    """Linear resample, which is what the reference implementation uses."""
    audio = np.asarray(audio, dtype=np.float64)
    if orig_sr == target_sr or audio.size == 0:
        return audio
    duration = audio.size / orig_sr
    target_n = max(1, int(round(duration * target_sr)))
    source_t = np.arange(audio.size) / orig_sr
    target_t = np.arange(target_n) / target_sr
    return np.interp(target_t, source_t, audio)


def extract_features(
    audio: np.ndarray, sr: int = TARGET_SAMPLE_RATE
) -> ProsodyFeatures:
    """Audio -> the f0 and intensity the conformer consumes.

    Returns raw f0_hz and intensity_db; call `.normalised()` for the model
    input.

    The audio is used as given. Nothing here cuts a quiet lead-in or tail:
    dataset clips arrive cut to the utterance already, and re-cutting them
    on this pipeline's own intensity estimate is a step the published runs
    did not perform, so it would make the WER incomparable with theirs. A
    live microphone does need that, and the caller that has one -- the
    drawing demo -- trims before calling.
    """
    import parselmouth
    from amfm_decompy import basic_tools, pYAAPT

    from ._yaapt_patch import patched_interpolate

    audio = np.asarray(audio, dtype=np.float64)
    if audio.ndim > 1:
        audio = audio.squeeze()
    if audio.size == 0:
        return ProsodyFeatures(np.zeros(0), np.zeros(0))

    if sr != TARGET_SAMPLE_RATE:
        audio = resample(audio, sr, TARGET_SAMPLE_RATE)
        sr = TARGET_SAMPLE_RATE

    # The preprocessor patches pYAAPT's interpolation before every call and
    # reads samp_interp off the result rather than interpolating the voiced
    # values afterwards. Patch the class, not a stub, so samp_interp is
    # produced the same way.
    pYAAPT.PitchObj.interpolate = patched_interpolate

    # Zero-pad by half a frame each side before tracking. This shifts every
    # frame boundary, so without it the whole contour sits offset against
    # what training saw.
    pad = int(FRAME_LENGTH_MS / 1000 * sr) // 2
    padded = np.pad(audio, (pad, pad), "constant", constant_values=0)

    signal = basic_tools.SignalObj(padded, sr)
    pitch = pYAAPT.yaapt(
        signal,
        frame_length=FRAME_LENGTH_MS,
        frame_space=FRAME_SPACE_MS,
        nccf_thresh1=NCCF_THRESH1,
        tda_frame_length=TDA_FRAME_LENGTH_MS,
    )
    f0 = np.asarray(pitch.samp_interp, dtype=np.float64)
    f0 = f0[LEADING_FRAMES_DROPPED:]
    f0[f0 > F0_CEILING_HZ] = 0.0
    f0[f0 < 0.0] = 0.0

    # No minimum_pitch: the preprocessor passes only time_step, and
    # parselmouth's default pitch floor (100 Hz) sets the analysis window.
    # Passing 50 Hz instead doubles that window and smooths the envelope.
    sound = parselmouth.Sound(audio, sampling_frequency=sr, start_time=0)
    intensity = sound.to_intensity(time_step=1.0 / FEATURE_RATE_HZ)
    intensity_db = np.asarray(intensity.values.T.flatten(), dtype=np.float64)

    frames = min(len(f0), len(intensity_db))
    f0 = f0[:frames]
    intensity_db = intensity_db[:frames]
    intensity_db[intensity_db < INTENSITY_CLAMP_DB] = INTENSITY_CLAMP_DB

    return ProsodyFeatures(f0_hz=f0, intensity_db=intensity_db)