# SPDX-License-Identifier: CC-BY-NC-SA-4.0
# Copyright 2026 TeraSpace. See LICENSE.md for encoder attribution.
"""TeraAnaly: one frozen speech encoder and independently callable task heads.

Local model execution needs only Python's standard library and PyTorch.
Transformers enables AutoModel(..., trust_remote_code=True); huggingface_hub
is used only when downloading a component from a Hub repository.
"""
import json
import math
import mmap
import re
import struct
import sys
import unicodedata
from pathlib import Path
from typing import Any, Sequence

import torch
import torch.nn.functional as F
from torch import Tensor, nn

SAMPLE_RATE = 16000
HOP_SAMPLES = 320
FRAME_SECONDS = 0.02
WINDOW_SAMPLES = 1024
DEFAULT_MIN_PAUSE_SECONDS = 0.10
SHARED_LAYERS = (4, 8, 16, 24)
STRESS_VOWEL_ORDER = "аеёиоуыэюя"
RUSSIAN_VOWELS = frozenset(STRESS_VOWEL_ORDER)
RUSSIAN_WORD = re.compile(r"[А-Яа-яЁё][А-Яа-яЁё\u0301]*")


def _resample(audio: Tensor, source_rate: int, target_rate: int = 16000) -> Tensor:
    if source_rate == target_rate:
        return audio
    # Windowed sinc low-pass reconstruction, in bounded chunks.
    length = math.ceil(audio.numel() * target_rate / source_rate)
    radius = 32
    taps = torch.arange(-radius + 1, radius + 1, device=audio.device)
    cutoff = min(1.0, target_rate / source_rate)
    output = []
    for start in range(0, length, 16384):
        positions = torch.arange(start, min(start + 16384, length), device=audio.device)
        positions = positions.to(torch.float64) * source_rate / target_rate
        indices = positions.floor().long()[:, None] + taps[None, :]
        offsets = positions[:, None] - indices
        window = (0.5 + 0.5 * torch.cos(math.pi * offsets / radius)).clamp_min(0)
        weights = torch.sinc(offsets * cutoff) * cutoff * window
        weights *= (offsets.abs() <= radius).to(weights.dtype)
        weights /= weights.sum(dim=1, keepdim=True)
        values = audio[indices.clamp(0, audio.numel() - 1)]
        output.append((values * weights.to(values.dtype)).sum(dim=1))
    return torch.cat(output)


def read_wav(path: str | Path) -> Tensor:
    """Read PCM or IEEE float WAV, average channels, and return 16 kHz audio."""
    data = memoryview(Path(path).read_bytes())
    if len(data) < 12 or data[:4].tobytes() != b"RIFF" or data[8:12].tobytes() != b"WAVE":
        raise ValueError("Expected a RIFF WAV file")
    fmt = raw = None
    offset = 12
    while offset + 8 <= len(data):
        kind = data[offset : offset + 4].tobytes()
        size = struct.unpack_from("<I", data, offset + 4)[0]
        chunk = data[offset + 8 : offset + 8 + size]
        if len(chunk) != size:
            raise ValueError("Truncated WAV chunk")
        if kind == b"fmt ":
            fmt = chunk
        elif kind == b"data":
            raw = chunk
        offset += 8 + size + (size & 1)
    if fmt is None or raw is None or len(fmt) < 16:
        raise ValueError("WAV file has no fmt or data chunk")
    encoding, channels, rate, _, block_align, bits = struct.unpack_from("<HHIIHH", fmt)
    if encoding == 0xFFFE and len(fmt) >= 40:
        encoding = struct.unpack_from("<H", fmt, 24)[0]
    width = bits // 8
    if channels < 1 or rate < 1 or bits % 8 or block_align != channels * width:
        raise ValueError("Invalid WAV channel, sample rate, or bit depth")
    if encoding == 3 and width == 4:
        samples = torch.frombuffer(bytearray(raw), dtype=torch.float32).clone()
    elif encoding != 1 or width not in (1, 2, 3, 4):
        raise ValueError("Expected PCM or float32 WAV")
    elif width == 1:
        samples = torch.tensor(list(raw), dtype=torch.float32).sub_(128).div_(128)
    elif width == 3:
        triples = torch.tensor(list(raw), dtype=torch.int32).reshape(-1, 3)
        values = triples[:, 0] | (triples[:, 1] << 8) | (triples[:, 2] << 16)
        values = (values ^ 0x800000) - 0x800000
        samples = values.to(torch.float32).div_(8388608)
    else:
        dtype = torch.int16 if width == 2 else torch.int32
        samples = torch.frombuffer(bytearray(raw), dtype=dtype).to(torch.float32)
        samples /= float(1 << (8 * width - 1))
    if samples.numel() % channels:
        raise ValueError("WAV data has an incomplete frame")
    samples = samples.reshape(-1, channels).mean(dim=1)
    if samples.numel() < 400:
        raise ValueError("WAV audio is too short for feature extraction")
    return _resample(samples, rate)


def _read_safetensors(path: Path) -> dict[str, Tensor]:
    """Map dense F32/F16/BF16 tensors without a safetensors dependency."""
    if sys.byteorder != "little":
        raise ValueError("Safetensors reader requires a little-endian host")
    with path.open("rb") as stream:
        raw_size = stream.read(8)
        if len(raw_size) != 8:
            raise ValueError("Truncated safetensors header")
        size = struct.unpack("<Q", raw_size)[0]
        if not 2 <= size <= min(path.stat().st_size - 8, 10_000_000):
            raise ValueError("Invalid safetensors header size")
        header = json.loads(stream.read(size))
        mapped = mmap.mmap(stream.fileno(), 0, access=mmap.ACCESS_COPY)
    dtypes = {"F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16}
    state = {}
    for name, info in header.items():
        if name == "__metadata__":
            continue
        dtype = dtypes.get(info["dtype"])
        if dtype is None:
            raise ValueError(f"Unsupported safetensors dtype for {name}: {info['dtype']}")
        start, end = info["data_offsets"]
        shape = info["shape"]
        count = math.prod(shape)
        width = torch.empty((), dtype=dtype).element_size()
        if not 0 <= start <= end <= len(mapped) - 8 - size or end - start != count * width:
            raise ValueError(f"Invalid safetensors offsets for {name}")
        state[name] = torch.frombuffer(mapped, dtype=dtype, count=count, offset=8 + size + start).reshape(shape)
    return state


def _frame_audio(waveform: Tensor, count: int) -> Tensor:
    """Return windows centered halfway into each alignment frame."""
    left = WINDOW_SAMPLES // 2 - HOP_SAMPLES // 2
    right = WINDOW_SAMPLES // 2
    mode = "reflect" if waveform.numel() > right else "replicate"
    padded = F.pad(waveform[None, None], (left, right), mode=mode)[0, 0]
    windows = padded.unfold(0, WINDOW_SAMPLES, HOP_SAMPLES)
    if count > windows.shape[0]:
        raise ValueError(
            f"Alignment requires {count} frames, but the audio provides only {windows.shape[0]}"
        )
    return windows[:count]


def _yin_pitch(
    windows: Tensor,
    *,
    min_hz: float,
    max_hz: float,
    threshold: float,
    silence_rms: float,
) -> tuple[Tensor, Tensor]:
    """Return Hz (zero for unvoiced) and voiced mask for one batch of windows."""
    length = windows.shape[1]
    min_lag = max(2, math.floor(SAMPLE_RATE / max_hz))
    max_lag = min(length // 2, math.ceil(SAMPLE_RATE / min_hz))
    if max_lag <= min_lag + 1:
        raise ValueError("Pitch range does not fit the analysis window")

    centered = windows - windows.mean(dim=1, keepdim=True)
    rms = windows.square().mean(dim=1).sqrt()
    # YIN is scale invariant; normalize before numerical stabilization.
    scale = centered.square().mean(dim=1, keepdim=True).sqrt()
    centered = centered / scale.clamp_min(torch.finfo(centered.dtype).tiny)
    spectrum = torch.fft.rfft(centered, n=length * 2, dim=1)
    correlation = torch.fft.irfft(spectrum.abs().square(), n=length * 2, dim=1)
    lags = torch.arange(max_lag + 1, device=windows.device)
    powers = F.pad(centered.square().cumsum(dim=1), (1, 0))
    difference = (
        powers[:, length - lags]
        + powers[:, length, None]
        - powers[:, lags]
        - 2 * correlation[:, : max_lag + 1]
    ).clamp_min_(0)
    cumulative = difference[:, 1:].cumsum(dim=1)
    cmndf = torch.ones_like(difference)
    cmndf[:, 1:] = difference[:, 1:] * lags[1:] / cumulative.clamp_min(1e-12)

    candidates = cmndf[:, min_lag : max_lag + 1]
    valleys = (candidates < threshold) & (candidates <= cmndf[:, min_lag - 1 : max_lag])
    valleys &= candidates <= F.pad(candidates[:, 1:], (0, 1), value=float("inf"))
    has_valley = valleys.any(dim=1)
    first_valley = valleys.to(torch.int32).argmax(dim=1)
    best = candidates.argmin(dim=1)
    lag = torch.where(has_valley, first_valley, best) + min_lag
    selected = cmndf.gather(1, lag[:, None])[:, 0]
    voiced = (selected < 0.25) & (rms >= silence_rms)

    before = cmndf.gather(1, (lag - 1)[:, None])[:, 0]
    after = cmndf.gather(1, (lag + 1).clamp(max=max_lag)[:, None])[:, 0]
    denominator = before - 2 * selected + after
    safe_denominator = torch.where(
        denominator.abs() > 1e-8, denominator, torch.ones_like(denominator)
    )
    delta = torch.where(
        (denominator.abs() > 1e-8) & (lag < max_lag),
        0.5 * (before - after) / safe_denominator,
        torch.zeros_like(selected),
    ).clamp(-1, 1)
    hz = SAMPLE_RATE / (lag.to(torch.float32) + delta)
    return torch.where(voiced, hz, torch.zeros_like(hz)), voiced


def _pause_statistics(waveform: Tensor, *, min_pause_seconds: float = DEFAULT_MIN_PAUSE_SECONDS, pause_threshold_db: float = -35.0) -> dict[str, Any]:
    """Find quiet internal intervals on a 20 ms RMS grid; retain edge silence."""
    if not math.isfinite(min_pause_seconds) or min_pause_seconds < FRAME_SECONDS:
        raise ValueError("Minimum pause duration must be finite and at least 20 ms")
    if not math.isfinite(pause_threshold_db) or pause_threshold_db >= 0:
        raise ValueError("Pause threshold must be a finite negative dB offset")
    count = math.ceil(waveform.numel() / HOP_SAMPLES)
    padded = F.pad(waveform, (0, count * HOP_SAMPLES - waveform.numel()))
    rms = padded.reshape(count, HOP_SAMPLES).square().mean(dim=1).sqrt()
    # Avoid understating RMS in the last partial frame.
    last_samples = waveform.numel() - (count - 1) * HOP_SAMPLES
    rms[-1] *= math.sqrt(HOP_SAMPLES / last_samples)
    cutoff = float(rms.max()) * 10 ** (pause_threshold_db / 20)
    quiet = (rms <= cutoff).tolist()
    pauses = []
    leading = trailing = None
    position = 0
    while position < count:
        if not quiet[position]:
            position += 1
            continue
        start = position
        while position < count and quiet[position]:
            position += 1
        begin = start * FRAME_SECONDS
        end = min(position * FRAME_SECONDS, waveform.numel() / SAMPLE_RATE)
        span = {"start": round(begin, 6), "end": round(end, 6), "duration_seconds": round(end - begin, 6)}
        if start == 0:
            leading = span
        if position == count:
            trailing = span
        if start > 0 and position < count and end - begin + 1e-9 >= min_pause_seconds:
            pauses.append(span)
    return {"pauses": pauses, "leading_silence": leading, "trailing_silence": trailing, "average_pause_duration_seconds": math.fsum(pause["duration_seconds"] for pause in pauses) / len(pauses) if pauses else 0.0, "pause_count": len(pauses), "pause_method": "20 ms RMS below peak-relative threshold, no absolute floor; quiet internal intervals only; leading/trailing silence excluded", "pause_threshold_db": pause_threshold_db, "pause_rms_cutoff": cutoff, "min_pause_seconds": min_pause_seconds}


def w2vbert_fbank(waveform: Tensor, *, spectral_statistics: bool = False, raw_features: bool = False):
    """Kaldi-style 80-bin log mel, unbiased CMVN, stacked 10 ms frames.

    Matches SeamlessM4TFeatureExtractor for a single 16 kHz recording, including
    its complex64 FFT rounding and padding value of 1. No NumPy/torchaudio.
    """
    if waveform.ndim != 1 or waveform.numel() < 560 or not torch.isfinite(waveform).all():
        raise ValueError("w2v-BERT needs finite mono 16 kHz audio of at least 35 ms")
    audio = (waveform.detach().cpu().float() * 32768).double()
    windows = audio.unfold(0, 400, 160).clone()
    windows -= windows.mean(1, keepdim=True)
    windows[:, 1:] -= 0.97 * windows[:, :-1].clone()
    windows[:, 0] *= 0.03
    window = torch.hann_window(400, periodic=False, dtype=torch.float64).pow(0.85)
    spectrum = torch.fft.rfft(windows * window, n=512).to(torch.complex64).to(torch.complex128).abs().square()
    frequencies = 1127 * torch.log1p(torch.arange(257, dtype=torch.float64) * (16000 / 512) / 700)
    centers = torch.linspace(1127 * math.log1p(20 / 700), 1127 * math.log1p(8000 / 700), 82, dtype=torch.float64)
    lower = (frequencies[:, None] - centers[:-2]) / (centers[1:-1] - centers[:-2])
    upper = (centers[2:] - frequencies[:, None]) / (centers[2:] - centers[1:-1])
    filters = torch.minimum(lower, upper).clamp_min(0)
    features = (spectrum @ filters).clamp_min(1.192092955078125e-7).log()
    if spectral_statistics:
        mean = features.mean(0)
        spectral = torch.cat((mean-mean.mean(),features.std(0,correction=0))).float()
    if raw_features:
        raw = (features-features.mean()).float()
    features = (features - features.mean(0)) / (features.var(0, correction=1) + 1e-7).sqrt()
    valid = features.shape[0] // 2
    if features.shape[0] % 2:
        features = F.pad(features, (0, 0, 0, 1), value=1)
        if raw_features:raw = F.pad(raw,(0,0,0,1))
    result = (features.reshape(1, -1, 160).float(), valid)
    if spectral_statistics:result=(*result,spectral)
    if raw_features:result=(*result,raw.reshape(-1,160)[:valid])
    return result


class _BertFFN(nn.Module):
    def __init__(self, config: dict[str, Any], activation: str = "silu") -> None:
        super().__init__()
        self.intermediate_dense = nn.Linear(config["hidden_size"], config["intermediate_size"])
        self.output_dense = nn.Linear(config["intermediate_size"], config["hidden_size"])
        self.activation = activation

    def forward(self, x: Tensor) -> Tensor:
        x = self.intermediate_dense(x)
        return self.output_dense(F.relu(x) if self.activation == "relu" else F.silu(x))


class _BertAttention(nn.Module):
    def __init__(self, config: dict[str, Any], relative: bool = True) -> None:
        super().__init__()
        hidden = config["hidden_size"]
        self.heads = config["num_attention_heads"]
        self.width = hidden // self.heads
        for name in ("linear_q", "linear_k", "linear_v", "linear_out"):
            setattr(self, name, nn.Linear(hidden, hidden))
        self.left = config["left_max_position_embeddings"]
        self.right = config["right_max_position_embeddings"]
        self.distance_embedding = nn.Embedding(self.left + self.right + 1, self.width) if relative else None

    def forward(self, x: Tensor, valid: int | Tensor) -> Tensor:
        batch, length, hidden = x.shape
        def project(name):
            return getattr(self, name)(x).reshape(batch, length, self.heads, self.width).transpose(1, 2)
        q, k, v = (project(name) for name in ("linear_q", "linear_k", "linear_v"))
        bias = None
        if self.distance_embedding is not None:
            positions = torch.arange(length, device=x.device)
            distances = (positions[None] - positions[:, None]).clamp(-self.left, self.right) + self.left
            # Compute Q dot each of the 73 relative vectors once, then gather.
            relative = q @ self.distance_embedding.weight.T
            bias = relative.gather(-1, distances[None, None].expand(batch, self.heads, -1, -1)) / math.sqrt(self.width)
        if isinstance(valid, Tensor) or valid < length:
            if bias is None:
                bias = x.new_zeros(1, 1, 1, length)
            limits = valid[:, None, None, None] if isinstance(valid, Tensor) else valid
            bias = bias.masked_fill(torch.arange(length, device=x.device)[None, None, None] >= limits, float("-inf"))
        attended = F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.0)
        return self.linear_out(attended.transpose(1, 2).reshape(batch, length, hidden))


class _BertConv(nn.Module):
    def __init__(self, config: dict[str, Any]) -> None:
        super().__init__()
        hidden, kernel = config["hidden_size"], config["conv_depthwise_kernel_size"]
        self.layer_norm = nn.LayerNorm(hidden, eps=config["layer_norm_eps"])
        self.pointwise_conv1 = nn.Conv1d(hidden, 2 * hidden, 1, bias=False)
        self.depthwise_conv = nn.Conv1d(hidden, hidden, kernel, groups=hidden, bias=False)
        self.depthwise_layer_norm = nn.LayerNorm(hidden, eps=config["layer_norm_eps"])
        self.pointwise_conv2 = nn.Conv1d(hidden, hidden, 1, bias=False)

    def forward(self, x: Tensor, valid: int | Tensor) -> Tensor:
        x = self.layer_norm(x)
        if isinstance(valid, Tensor) or valid < x.shape[1]:
            limits = valid[:, None, None] if isinstance(valid, Tensor) else valid
            x = x.masked_fill(torch.arange(x.shape[1], device=x.device)[None, :, None] >= limits, 0)
        x = F.glu(self.pointwise_conv1(x.transpose(1, 2)), dim=1)
        x = self.depthwise_conv(F.pad(x, (self.depthwise_conv.kernel_size[0] - 1, 0)))
        x = F.silu(self.depthwise_layer_norm(x.transpose(1, 2)))
        return self.pointwise_conv2(x.transpose(1, 2)).transpose(1, 2)


class _BertLayer(nn.Module):
    def __init__(self, config: dict[str, Any]) -> None:
        super().__init__()
        for name in ("ffn1_layer_norm", "self_attn_layer_norm", "ffn2_layer_norm", "final_layer_norm"):
            setattr(self, name, nn.LayerNorm(config["hidden_size"], eps=config["layer_norm_eps"]))
        self.ffn1 = _BertFFN(config)
        self.ffn2 = _BertFFN(config)
        self.self_attn = _BertAttention(config)
        self.conv_module = _BertConv(config)

    def forward(self, x: Tensor, valid: int | Tensor) -> Tensor:
        x = x + 0.5 * self.ffn1(self.ffn1_layer_norm(x))
        x = x + self.self_attn(self.self_attn_layer_norm(x), valid)
        x = x + self.conv_module(x, valid)
        return self.final_layer_norm(x + 0.5 * self.ffn2(self.ffn2_layer_norm(x)))


def ctc_forced_spans(log_probs: Tensor, tokens: Sequence[int], blank: int) -> list[tuple[int, int]]:
    """Viterbi on CTC's blank/character graph, including repeated characters."""
    if log_probs.ndim != 2 or not tokens or not torch.isfinite(log_probs).all():
        raise ValueError("Expected finite CTC frame log probabilities and a nonempty target")
    if any(type(token) is not int or not 0 <= token < log_probs.shape[1] or token == blank for token in tokens):
        raise ValueError("Invalid CTC target token")
    needed = len(tokens) + sum(a == b for a, b in zip(tokens, tokens[1:]))
    if log_probs.shape[0] < needed:
        raise ValueError("Audio is too short for this transcript's CTC path")
    emissions = torch.full((len(tokens) * 2 + 1,), blank, dtype=torch.long)
    emissions[1::2] = torch.tensor(tokens)
    probabilities = log_probs.detach().float().cpu()[:, emissions]
    skip = torch.zeros(len(emissions), dtype=torch.bool)
    skip[2:] = (emissions[2:] != blank) & (emissions[2:] != emissions[:-2])
    scores = torch.full((len(emissions),), float("-inf"))
    scores[:2] = probabilities[0, :2]
    predecessors = torch.empty(probabilities.shape, dtype=torch.int8)
    predecessors[0] = 0
    for frame in range(1, len(probabilities)):
        previous1 = F.pad(scores[:-1], (1, 0), value=float("-inf"))
        previous2 = F.pad(scores[:-2], (2, 0), value=float("-inf")).masked_fill(~skip, float("-inf"))
        best, transition = torch.stack((scores, previous1, previous2)).max(0)
        scores = best + probabilities[frame]
        predecessors[frame] = transition.to(torch.int8)
    state = len(emissions) - 2 + scores[-2:].argmax().item()
    if not math.isfinite(scores[state].item()):
        raise ValueError("No valid forced CTC alignment")
    path = [state]
    for frame in range(len(probabilities) - 1, 0, -1):
        state -= predecessors[frame, state].item()
        path.append(state)
    path.reverse()
    positions = [[] for _ in tokens]
    for frame, state in enumerate(path):
        if state % 2:
            positions[state // 2].append(frame)
    if any(not frames for frames in positions):
        raise ValueError("CTC path omitted a target character")
    return [(frames[0], frames[-1] + 1) for frames in positions]


def wav2vec_vowel_features(alignment: dict[str, Any], hidden: Tensor, *, hidden_size: int = 1024) -> tuple[Tensor, list[dict[str, Any]]]:
    """Pool frozen audio embeddings; do not use energy or word spelling as inputs.

    Each row is vowel mean/std, mean minus the word's vowel mean, vowel identity,
    and three duration measurements, on the native 20 ms encoder grid.
    """
    text = alignment["normalized_text"]
    characters = alignment["characters"]
    if hidden.ndim != 2 or hidden.shape[1] != hidden_size or len(characters) != len(text):
        raise ValueError("Invalid alignment or hidden features")
    if not torch.isfinite(hidden).all():
        raise ValueError("Non-finite wav2vec features")
    rows, words = [], []
    hidden = hidden.float()
    for match in RUSSIAN_WORD.finditer(text):
        vowels = [i for i in range(match.start(), match.end()) if text[i].lower() in RUSSIAN_VOWELS]
        if not vowels:
            continue
        means, deviations, durations = [], [], []
        for index in vowels:
            item = characters[index]
            if item.get("char") != text[index]:
                raise ValueError("Character alignment does not match text")
            start = round(item["start"] / FRAME_SECONDS)
            end = round(item["end"] / FRAME_SECONDS)
            if not 0 <= start < end <= hidden.shape[0]:
                raise ValueError(f"Hidden features do not cover vowel {index}")
            span = hidden[start:end]
            means.append(span.mean(0))
            deviations.append(span.std(0, correction=0))
            durations.append(end - start)
        mean = torch.stack(means)
        word_mean = mean.mean(0)
        word_frames = (characters[match.end() - 1]["end"] - characters[match.start()]["start"]) / FRAME_SECONDS
        row_start = len(rows)
        for index, avg, std, duration in zip(vowels, means, deviations, durations):
            identity = hidden.new_zeros(len(STRESS_VOWEL_ORDER))
            identity[STRESS_VOWEL_ORDER.index(text[index].lower())] = 1
            timing = hidden.new_tensor([
                math.log(duration), duration / max(word_frames, 1), duration / (sum(durations) / len(durations)),
            ])
            rows.append(torch.cat((avg, std, avg - word_mean, identity, timing)))
        words.append({"word": match.group(), "start_index": match.start(), "end_index": match.end(),
                      "vowels": vowels, "row_start": row_start, "row_end": len(rows)})
    dimension = hidden_size * 3 + len(STRESS_VOWEL_ORDER) + 3
    return (torch.stack(rows) if rows else hidden.new_empty((0, dimension))), words


def shared_recording_features(taps: dict[int, Tensor], layers: Sequence[int] = SHARED_LAYERS,
                              statistics: dict[str, Any] | None = None) -> Tensor:
    values = [value for layer in layers for value in (taps[layer].float().mean(0), taps[layer].float().std(0, correction=0))]
    if statistics is not None:
        # Voice and prosody cues are computed once alongside the encoder.
        pitch = [statistics.get(name) for name in ("average_pitch_hz", "pitch_p10_hz", "pitch_median_hz", "pitch_p90_hz")]
        pitch = [math.log2(value) if value is not None and value > 0 else 0.0 for value in pitch]
        duration = statistics["duration_seconds"]
        pause_fraction = sum(item["duration_seconds"] for item in statistics.get("pauses", [])) / max(duration, 1e-6)
        extra = pitch + [(statistics.get("pitch_range_semitones") or 0.0) / 12,
                         statistics.get("voiced_fraction", 0.0), math.log(max(duration, 0.01)), pause_fraction]
        values.append(values[0].new_tensor(extra))
    return torch.cat(values)


def recording_statistics(waveform, frame_count, device):
    windows = _frame_audio(waveform.to(device), frame_count)
    peak_rms = windows.square().mean(1).sqrt().max().item()
    gate = peak_rms * 0.05 if peak_rms > 0 else float('inf')
    pitch, voiced = [], []
    for chunk in windows.split(512):
        p, v = _yin_pitch(chunk, min_hz=50, max_hz=800, threshold=0.15, silence_rms=gate)
        pitch.append(p.cpu())
        voiced.append(v.cpu())
    pitch, voiced = torch.cat(pitch), torch.cat(voiced)
    values = pitch[voiced]
    quantiles = torch.quantile(values, torch.tensor([0.1, 0.5, 0.9])) if values.numel() else None
    return {
        'average_pitch_hz': float(values.mean()) if values.numel() else None,
        'pitch_p10_hz': float(quantiles[0]) if quantiles is not None else None,
        'pitch_median_hz': float(quantiles[1]) if quantiles is not None else None,
        'pitch_p90_hz': float(quantiles[2]) if quantiles is not None else None,
        'pitch_range_semitones': 12 * math.log2(float(quantiles[2] / quantiles[0])) if quantiles is not None else None,
        'voiced_fraction': float(voiced.float().mean()),
        'duration_seconds': waveform.numel() / SAMPLE_RATE,
        'pauses': _pause_statistics(waveform)['pauses'],
    }


def align(text, logits, vocab, blank, frame_count):
    text = ' '.join(unicodedata.normalize('NFKC', text).split())
    indices, tokens = [], []
    for i, char in enumerate(text):
        token = '|' if char == ' ' else char.lower()
        if token == 'ё' and token not in vocab:
            token = 'е'
        if token in vocab and vocab[token] != blank:
            indices.append(i)
            tokens.append(vocab[token])
        elif char.isalpha() or char.isdigit():
            raise ValueError(f'Unsupported character {char!r}; spell out numbers and use Russian text')
    if not tokens:
        raise ValueError('Text has no supported spoken characters')
    spans = ctc_forced_spans(logits.log_softmax(-1), tokens, blank)
    starts = [start for start, _ in spans]
    ends = [min(end, frame_count) for _, end in spans]
    boundaries = [starts[0]] + [(ends[i] + starts[i+1]) // 2 for i in range(len(spans)-1)] + [ends[-1]]
    mapped = {index: (boundaries[i], boundaries[i+1]) for i, index in enumerate(indices)}
    characters, position = [], boundaries[0]
    for i, char in enumerate(text):
        start, end = mapped.get(i, (position, position))
        characters.append({'char': char, 'start': start * FRAME_SECONDS, 'end': end * FRAME_SECONDS})
        position = end
    return {'normalized_text': text, 'characters': characters}


# Transformers is optional for local PyTorch use, and supplies AutoConfig on HF.
try:
    from transformers import PretrainedConfig as _ConfigBase
except ImportError:
    class _ConfigBase:
        def __init__(self, **kwargs):
            self.__dict__.update(kwargs)

        def to_dict(self):
            return dict(self.__dict__)


class TeraAnalyConfig(_ConfigBase):
    model_type = 'tera_analy'

    def __init__(self, encoder_config=None, heads=None, ctc=None, vocab=None,
                 global_layers=None, use_prosody_features=True, labels=None,
                 source=None, **kwargs):
        super().__init__(**kwargs)
        self.encoder_config = encoder_config or {}
        self.heads = heads or {}
        self.ctc = ctc or {}
        self.vocab = vocab or {}
        self.global_layers = global_layers or [4, 8, 16, 24]
        self.use_prosody_features = use_prosody_features
        self.labels = labels or {}
        self.source = source or {}


def _resolve_file(source, name, **options):
    directory = Path(source).expanduser()
    if directory.is_dir():
        path = directory / name
        if not path.is_file():
            raise FileNotFoundError(path)
        return path
    try:
        from huggingface_hub import hf_hub_download
    except ImportError as error:
        raise ImportError('Install huggingface_hub to download models, or pass a local model folder') from error
    allowed = {key: options[key] for key in ('revision', 'cache_dir', 'token', 'local_files_only', 'force_download')
               if key in options}
    return Path(hf_hub_download(str(source), name, **allowed))


def _load_config(source, config, options):
    options = dict(options)
    if config is None:
        path = _resolve_file(source, 'config.json', **options)
        config = TeraAnalyConfig(**json.loads(path.read_text(encoding='utf-8')))
        if not Path(source).expanduser().is_dir() and re.fullmatch('[0-9a-f]{40}', path.parent.name):
            options['revision'] = path.parent.name
    elif isinstance(config, dict):
        config = TeraAnalyConfig(**config)
    commit = getattr(config, '_commit_hash', None)
    if commit:
        options['revision'] = commit
    dtype = options.get('dtype', options.get('torch_dtype'))
    if dtype not in (None, 'auto', 'float32', torch.float32):
        raise ValueError('This release uses float32 inference; omit dtype or use torch.float32')
    return config, options


def _load_state(module, source, filename, options):
    module.load_state_dict(_read_safetensors(_resolve_file(source, filename, **options)), strict=True, assign=True)


class _Component(nn.Module):
    filename = None

    @classmethod
    def from_pretrained(cls, source, *, config=None, device='cpu', **options):
        """Load only this component's safetensors (stress also loads CTC)."""
        config, options = _load_config(source, config, options)
        with torch.device('meta'):
            module = cls(config)
        module._load_weights(source, options)
        module.to(device).eval()
        if isinstance(module, TeraEncoder):
            module.requires_grad_(False)
        return module

    def _load_weights(self, source, options):
        _load_state(self, source, self.filename, options)

    @property
    def device(self):
        return next(self.parameters()).device


def _is_batch(value):
    return isinstance(value, (list, tuple))


def _audio_items(waveform, sample_rate):
    batched = _is_batch(waveform) or isinstance(waveform, Tensor) and waveform.ndim == 2
    items = list(waveform) if batched else [waveform]
    rates = list(sample_rate) if _is_batch(sample_rate) else [sample_rate]*len(items)
    if len(rates) != len(items) or any(not isinstance(rate, int) or rate <= 0 for rate in rates):
        raise ValueError('Supply one positive integer sample rate per waveform, or one rate for all')
    if any(not isinstance(item, Tensor) or item.ndim != 1 for item in items):
        raise ValueError('Expected a mono waveform, a list of mono waveforms, or a [batch, samples] tensor')
    return items, rates, batched


def _matching_texts(states, text):
    if _is_batch(states):
        if not _is_batch(text) or len(text) != len(states):
            raise ValueError('Supply one matching transcript per batch item')
        texts = list(text)
    else:
        texts = [text]
    if any(not isinstance(item, str) or not item.strip() for item in texts):
        raise ValueError('Text must contain the matching audio transcript')
    return texts


class TeraEncoder(_Component):
    """Frozen 24-layer Conformer; returns states for any number of task heads."""
    filename = 'encoder.safetensors'

    def __init__(self, config):
        super().__init__()
        self.config = config
        c = config.encoder_config
        self.feature_projection = nn.Module()
        self.feature_projection.layer_norm = nn.LayerNorm(160, eps=c['layer_norm_eps'])
        self.feature_projection.projection = nn.Linear(160, 1024)
        self.layers = nn.ModuleList(_BertLayer(c) for _ in range(c['num_hidden_layers']))
        self.required_layers = sorted(set(config.global_layers) | {config.heads['stress']['layer'], config.ctc['layer']}
                                      | {h['layer'] for name in ('age', 'gender') for h in config.heads[name]['frames']})
        self.requires_grad_(False)

    @torch.no_grad()
    def forward(self, waveform, sample_rate=16000):
        """Accept one mono waveform, or a batch of different-length waveforms.

        Single input returns a states dictionary; list/2D input returns a list
        of dictionaries in input order. A sample rate list is also accepted.
        The neural encoder runs once for the padded, masked batch. Frontend
        and prosody statistics are computed separately for each recording.
        """
        items, rates, batched = _audio_items(waveform, sample_rate)
        if not items:
            return []
        prepared = []
        for item, rate in zip(items, rates):
            audio = _resample(item.detach().cpu().float(), rate)
            features, valid, spectral, raw = w2vbert_fbank(audio, spectral_statistics=True, raw_features=True)
            if features.shape[1] > 2000:
                raise ValueError('Use audio clips of at most 40 seconds')
            prepared.append((audio, features, valid, spectral, raw))
        length = max(item[1].shape[1] for item in prepared)
        features = torch.cat([F.pad(item[1], (0, 0, 0, length-item[1].shape[1])) for item in prepared])
        counts = torch.tensor([item[2] for item in prepared], device=self.device)
        valid = prepared[0][2] if len(items) == 1 else counts
        x = self.feature_projection.projection(self.feature_projection.layer_norm(features.to(self.device)))
        x = x.masked_fill(torch.arange(length, device=x.device)[None, :, None] >= counts[:, None, None], 0)
        taps = [{} for _ in items]
        for number, layer in enumerate(self.layers, 1):
            x = layer(x, valid)
            if number in self.required_layers:
                for index, item in enumerate(prepared):
                    taps[index][number] = x[index, :item[2]]
        states = []
        for hidden, (audio, _, count, spectral, raw) in zip(taps, prepared):
            statistics = recording_statistics(audio, count, self.device) if self.config.use_prosody_features else None
            recording = shared_recording_features(hidden, self.config.global_layers, statistics)
            states.append({'hidden_states': hidden, 'recording_features': recording,
                           'spectral_features': spectral.to(self.device), 'raw_features': raw.to(self.device),
                           'frame_count': count, 'frame_seconds': FRAME_SECONDS})
        return states if batched else states[0]


class _RecordingMLP(nn.Module):
    def __init__(self, dimension, widths, outputs):
        super().__init__()
        self.register_buffer('mean', torch.zeros(dimension))
        self.register_buffer('scale', torch.ones(dimension))
        layers, previous = [], dimension
        for width in widths:
            layers.extend((nn.Linear(previous, width), nn.ReLU(), nn.Dropout(.25)))
            previous = width
        layers.append(nn.Linear(previous, outputs))
        self.network = nn.Sequential(*layers)

    def forward(self, features):
        features = features.to(self.mean.device)
        return self.network(((features-self.mean)/self.scale).clamp(-10, 10))


def _recording_input(states, dimension, spectral=False):
    if _is_batch(states):
        return (torch.stack([_recording_input(item, dimension, spectral) for item in states])
                if states else torch.empty(0, dimension))
    values = states['recording_features']
    if spectral or values.numel() != dimension:
        values = torch.cat((values, states['spectral_features'].to(values.device)))
    if values.numel() != dimension:
        raise ValueError('Encoder states do not match this head configuration')
    return values


class EmotionHead(_Component):
    filename = 'emotion.safetensors'

    def __init__(self, config):
        super().__init__()
        spec = config.heads['emotion']
        self.labels = config.labels['emotion']
        self.recording = _RecordingMLP(spec['dimension'], spec['widths'], len(self.labels))

    def logits(self, states):
        return self.recording(_recording_input(states, self.recording.mean.numel()))

    @torch.inference_mode()
    def forward(self, states):
        indices = self.logits(states).argmax(-1)
        return [self.labels[index] for index in indices.tolist()] if _is_batch(states) else self.labels[indices.item()]


class _FrameAttributeHead(nn.Module):
    def __init__(self, spec, outputs):
        super().__init__()
        width, projection, dimension = spec['width'], spec['projection_dim'], spec['global_dim']
        self.layer, self.frame_stride = spec['layer'], spec.get('frame_stride', 2)
        self.register_buffer('ssl_mean', torch.zeros(1024))
        self.register_buffer('ssl_projection', torch.zeros(1024, projection))
        self.register_buffer('mean', torch.zeros(dimension))
        self.register_buffer('scale', torch.ones(dimension))
        self.ssl_norm, self.mel_norm = nn.LayerNorm(projection), nn.LayerNorm(160)
        self.frame = nn.Linear(projection+160, width)
        self.temporal = nn.Conv1d(width, width, 5, padding=2, groups=width)
        self.attention = nn.Sequential(nn.Linear(width, 64), nn.Tanh(), nn.Linear(64, 1))
        self.body = nn.Sequential(nn.Linear(dimension+width*2, 256), nn.ReLU(), nn.Dropout(.25))
        self.output = nn.Linear(256, outputs)

    def forward(self, states):
        batched = _is_batch(states)
        items = states if batched else [states]
        if not items:
            return self.mean.new_empty((0, self.output.out_features))
        sequences = []
        for item in items:
            ssl = item['hidden_states'][self.layer].to(self.mean.device).float()
            raw = item['raw_features'].to(self.mean.device).float()
            sequences.append(torch.cat(((ssl-self.ssl_mean) @ self.ssl_projection, raw), -1)[::self.frame_stride])
        frames = nn.utils.rnn.pad_sequence(sequences, batch_first=True)
        lengths = torch.tensor([item.shape[0] for item in sequences], device=self.mean.device)
        mask = torch.arange(frames.shape[1], device=self.mean.device)[None] >= lengths[:, None]
        dimension = self.ssl_projection.shape[1]
        x = torch.cat((self.ssl_norm(frames[:, :, :dimension]), self.mel_norm(frames[:, :, dimension:])), -1)
        x = F.gelu(self.frame(x)).masked_fill(mask[:, :, None], 0)
        x = F.gelu(x+self.temporal(x.transpose(1, 2)).transpose(1, 2))
        weights = self.attention(x).squeeze(-1).masked_fill(mask, float('-inf')).softmax(-1)
        mean = (x*weights[:, :, None]).sum(1)
        std = ((x-mean[:, None]).square()*weights[:, :, None]).sum(1).clamp_min(1e-5).sqrt()
        recording = _recording_input(items, self.mean.numel(), spectral=True).to(self.mean.device)
        normalized = ((recording-self.mean)/self.scale).clamp(-10, 10)
        output = self.output(self.body(torch.cat((normalized, mean, std), -1)))
        return output if batched else output[0]


class _AttributeHead(_Component):
    task = None

    def __init__(self, config):
        super().__init__()
        spec = config.heads[self.task]
        self.labels = config.labels.get(self.task)
        outputs = 1 if self.task == 'age' else len(self.labels)
        self.recording = _RecordingMLP(spec['dimension'], spec['widths'], outputs)
        self.frames = nn.ModuleList(_FrameAttributeHead(item, outputs) for item in spec['frames'])
        self.blend_weight = spec['blend_weight']

    def logits(self, states):
        value = self.recording(_recording_input(states, self.recording.mean.numel()))
        for frame_head in self.frames:
            frame_value = frame_head(states)
            weight = self.blend_weight
            if weight == 1:
                value = frame_value
            elif weight != 0:
                value = ((1-weight)*value+weight*frame_value if self.task == 'age' else
                         ((1-weight)*value.softmax(-1)+weight*frame_value.softmax(-1)).clamp_min(1e-9).log())
        return value


class AgeHead(_AttributeHead):
    filename, task = 'age.safetensors', 'age'

    @torch.inference_mode()
    def forward(self, states):
        values = self.logits(states)
        return ([int(round(value*100)) for value in values.flatten().tolist()] if _is_batch(states)
                else int(round(values.item()*100)))


class GenderHead(_AttributeHead):
    filename, task = 'gender.safetensors', 'gender'

    @torch.inference_mode()
    def forward(self, states):
        indices = self.logits(states).argmax(-1)
        return [self.labels[index] for index in indices.tolist()] if _is_batch(states) else self.labels[indices.item()]


class CTCHead(_Component):
    """Native 20 ms character emissions; used internally by StressHead."""
    filename = 'ctc.safetensors'

    def __init__(self, config):
        super().__init__()
        self.layer, self.blank = config.ctc['layer'], config.ctc['blank']
        self.vocab = config.vocab
        self.norm = nn.LayerNorm(1024)
        self.temporal = nn.Conv1d(1024, 1024, 5, padding=2, groups=1024)
        self.network = nn.Sequential(nn.Linear(1024, config.ctc['hidden_size']), nn.GELU(),
                                     nn.Dropout(.15), nn.Linear(config.ctc['hidden_size'], config.ctc['output_size']))

    def forward(self, states):
        batched = _is_batch(states)
        items = states if batched else [states]
        if not items:
            return []
        sequences = [self.norm(item['hidden_states'][self.layer].to(self.device).float()) for item in items]
        values = nn.utils.rnn.pad_sequence(sequences, batch_first=True)
        values = values+self.temporal(values.transpose(1, 2)).transpose(1, 2)
        logits = self.network(values)
        return [item[:state['frame_count']] for item, state in zip(logits, items)] if batched else logits[0]

    @torch.inference_mode()
    def align(self, states, text):
        texts = _matching_texts(states, text)
        if _is_batch(states):
            return [align(transcript, logits, self.vocab, self.blank, state['frame_count'])
                    for state, transcript, logits in zip(states, texts, self(states))]
        return align(texts[0], self(states), self.vocab, self.blank, states['frame_count'])


class StressHead(_Component):
    """States + transcript -> CTC timings -> word-local stressed text."""

    def __init__(self, config):
        super().__init__()
        spec = config.heads['stress']
        self.layer = spec['layer']
        self.ranker = _RecordingMLP(spec['dimension'], spec['widths'], 1)
        self.ctc_head = CTCHead(config)

    def _load_weights(self, source, options):
        _load_state(self.ranker, source, 'stress.safetensors', options)
        self.ctc_head._load_weights(source, options)

    @torch.inference_mode()
    def forward(self, states, text):
        _matching_texts(states, text)
        batched = _is_batch(states)
        items = states if batched else [states]
        if not items:
            return []
        aligned = self.ctc_head.align(states, text)
        alignments = aligned if batched else [aligned]
        pooled = [wav2vec_vowel_features(alignment, state['hidden_states'][self.layer], hidden_size=1024)
                  for state, alignment in zip(items, alignments)]
        scores = self.ranker(torch.cat([item[0] for item in pooled])).squeeze(-1)
        if not torch.isfinite(scores).all():
            raise ValueError('Non-finite stress scores')
        results, offset = [], 0
        for alignment, (features, words) in zip(alignments, pooled):
            rendered = list(alignment['normalized_text'])
            for word in words:
                chosen = scores[offset+word['row_start']:offset+word['row_end']].argmax().item()
                index = word['vowels'][chosen]
                rendered[index] = '+'+rendered[index]
            results.append(''.join(rendered))
            offset += features.shape[0]
        return results if batched else results[0]


class TeraAnaly(nn.Module):
    """Independently callable encoder, emotion, age, gender, stress and CTC."""
    config_class = TeraAnalyConfig

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.encoder = TeraEncoder(config)
        self.emotion_head = EmotionHead(config)
        self.age_head = AgeHead(config)
        self.gender_head = GenderHead(config)
        self.stress_head = StressHead(config)

    @property
    def ctc_head(self):
        return self.stress_head.ctc_head

    @property
    def device(self):
        return self.encoder.device

    @classmethod
    def register_for_auto_class(cls, auto_class='AutoModel'):
        cls._auto_class = auto_class if isinstance(auto_class, str) else auto_class.__name__

    @classmethod
    def from_pretrained(cls, source, *args, config=None, device='cpu', **options):
        """Load separate safetensors; also called by HF AutoModel."""
        if args:
            raise TypeError('Pass options as keyword arguments')
        config, options = _load_config(source, config, options)
        with torch.device('meta'):
            model = cls(config)
        for module in (model.encoder, model.emotion_head, model.age_head, model.gender_head, model.stress_head):
            module._load_weights(source, options)
        device_map = options.get('device_map')
        if device_map == 'auto':
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
        elif device_map is not None:
            if not isinstance(device_map, (str, torch.device)):
                raise ValueError('Use a single device or call model.to(device)')
            device = device_map
        model.to(device).eval()
        model.encoder.requires_grad_(False)
        return model

    @staticmethod
    def read_wav(path):
        return read_wav(path)

    @torch.inference_mode()
    def forward(self, waveform, text, sample_rate=16000):
        """Return one result dict, or a list for a waveform/text batch."""
        items, _, batched = _audio_items(waveform, sample_rate)
        _matching_texts(items if batched else waveform, text)
        states = self.encoder(waveform, sample_rate)
        outputs = {'age': self.age_head(states), 'gender': self.gender_head(states),
                   'emotion': self.emotion_head(states), 'stressed_text': self.stress_head(states, text)}
        return ([dict(zip(outputs, values)) for values in zip(*outputs.values())] if batched else outputs)