Output:
Chunk ID: 85ec6f
Wall time: 0.0510 seconds
Process exited with code 0
Original token count: 2852
Output:
"""
Audio Polisher - Segment boundary cleanup for transcription quality.

Two-phase polishing:
  Phase 1 (original): Remove CUT SPEECH ARTIFACTS from VAD boundaries
    - Detects burst -> gap -> speech pattern and removes the burst
    - Silence at boundaries preserved as natural padding
  Phase 2 (new): TIGHTEN boundaries to actual speech onset/offset
    - Energy-based detection of where speech actually starts/ends
    - Cuts to those points with configurable silence margin (~50ms)
    - Ensures clean, crisp boundaries for TTS training data
    - Prevents segments starting/ending mid-voice-energy

Combined result: artifact-free + precisely-bounded audio segments.
No ML models - pure signal processing, runs in <10ms/segment.
"""
import os
import numpy as np
import soundfile as sf
from dataclasses import dataclass
from typing import Optional, Tuple, List
from pathlib import Path

@dataclass
class PolishResult:
    """Result of audio polishing for a single segment."""
    input_path: str
    output_path: str
    was_modified: bool
    start_trimmed_ms: float = 0.0
    end_trimmed_ms: float = 0.0
    start_quality: str = "clean"
    end_quality: str = "clean"
    original_duration_ms: float = 0.0
    polished_duration_ms: float = 0.0
    snr_db: float = 0.0
    rms_db: float = 0.0
    peak_db: float = 0.0
    is_clipped: bool = False
    volume_adjusted: bool = False
    volume_gain_db: float = 0.0

    @property
    def total_trimmed_ms(self):
        return self.start_trimmed_ms + self.end_trimmed_ms

    def summary(self):
        name = Path(self.input_path).name
        if not self.was_modified:
            return (f"  {name}: SKIP | {self.original_duration_ms:.0f}ms | "
                    f"SNR:{self.snr_db:.1f}dB | start:{self.start_quality} "
                    f"end:{self.end_quality}")
        parts = []
        if self.start_trimmed_ms > 0:
            parts.append(
                f"start -{self.start_trimmed_ms:.0f}ms ({self.start_quality})")
        if self.end_trimmed_ms > 0:
            parts.append(
                f"end -{self.end_trimmed_ms:.0f}ms ({self.end_quality})")
        if self.volume_adjusted:
            parts.append(f"gain +{self.volume_gain_db:.1f}dB")
        changes = " | ".join(parts)
        return (f"  {name}: POLISHED | {self.original_duration_ms:.0f}ms -> "
                f"{self.polished_duration_ms:.0f}ms | {changes} | "
                f"SNR:{self.snr_db:.1f}dB")

class AudioPolisher:
    """
    Two-phase audio segment polisher for transcription quality.

    Phase 1: Conservative artifact removal (burst-gap-speech pattern detection)
    Phase 2: Energy-based boundary tightening (speech onset/offset + silence margin)
    """

    def __init__(
        self,
        frame_ms=10.0,
        hop_ms=5.0,
        max_start_trim_ms=200.0,
        artifact_search_ms=150.0,
        min_gap_ms=15.0,
        max_end_trim_ms=150.0,
        min_burst_ms=10.0,
        max_burst_ms=100.0,
        energy_gap_factor=0.35,
        zcr_artifact_threshold=0.15,
        min_dynamic_range_db=6.0,
        min_silence_pad_ms=30.0,
        min_remaining_ms=300.0,
        min_trim_ms=10.0,
        fade_in_ms=5.0,
        fade_out_ms=10.0,
        target_rms_db=-18.0,
        snr_boost_threshold_db=12.0,
        max_gain_db=12.0,
        clipping_threshold=0.99,
        # Phase 2: boundary tightening params
        tighten_boundaries=True,
        silence_margin_ms=50.0,
        output_format="flac",
    ):
        self.frame_ms = frame_ms
        self.hop_ms = hop_ms
        self.max_start_trim_ms = max_start_trim_ms
        self.artifact_search_ms = artifact_search_ms
        self.min_gap_ms = min_gap_ms
        self.max_end_trim_ms = max_end_trim_ms
        self.min_burst_ms = min_burst_ms
        self.max_burst_ms = max_burst_ms
        self.energy_gap_factor = energy_gap_factor
        self.zcr_artifact_threshold = zcr_artifact_threshold
        self.min_dynamic_range_db = min_dynamic_range_db
        self.min_silence_pad_ms = min_silence_pad_ms
        self.min_remaining_ms = min_remaining_ms
        self.min_trim_ms = min_trim_ms
        self.fade_in_ms = fade_in_ms
        self.fade_out_ms = fade_out_ms
        self.target_rms_db = target_rms_db
        self.snr_boost_threshold_db = snr_boost_threshold_db
        self.max_gain_db = max_gain_db
        self.clipping_threshold = clipping_threshold
        self.tighten_boundaries_enabled = tighten_boundaries
        self.silence_margin_ms = silence_margin_ms
        self.output_format = output_format

    def _compute_frames(self, audio, sr):
        """Compute per-frame energy (dB) and zero-crossing rate."""
        frame_samples = int(self.frame_ms * sr / 1000)
        hop_samples = int(self.hop_ms * sr / 1000)
        n_frames = max(1, (len(audio) - frame_samples) // hop_samples + 1)
        ste_db = np.zeros(n_frames)
        zcr = np.zeros(n_frames)
        frame_times = np.zeros(n_frames)
        for i in range(n_frames):
            start = i * hop_samples
            end = min(start + frame_samples, len(audio))
            frame = audio[start:end]
            rms = np.sqrt(np.mean(frame ** 2))
            ste_db[i] = 20 * np.log10(rms + 1e-10)
            if len(frame) > 1:
                zcr[i] = (np.sum(np.abs(np.diff(np.sign(frame))))
                          / (2 * len(frame)))
            frame_times[i] = (start + frame_samples / 2) / sr
        return ste_db, zcr, frame_times

    def _detect_leading_artifact(self, ste_db, zcr, frame_times):
        """
        Detect CUT SPEECH artifact at segment start.

        ONLY acts when: burst (cut speech) -> gap (silence) -> real speech.
        Trims at gap START = removes burst, keeps gap as silence padding.

        Does NOT trim if segment starts with silence (that's good padding).
        Does NOT trim if segment starts with continuous speech (natural start).
        """
        search_limit = self.artifact_search_ms / 1000
        si = np.where(frame_times < search_limit)[0]
        if len(si) < 3:
            return None, "clean"

        nf = np.percentile(ste_db, 15)
        sl = np.percentile(ste_db, 85)
        dr = sl - nf
        if dr < self.min_dynamic_range_db:
            return None, "clean"

        gap_thresh = nf + self.energy_gap_factor * dr

        n_check = min(3, len(si))
        first_frames_energy = ste_db[si[:n_check]]
        if np.all(first_frames_energy < gap_thresh):
            return None, "clean"

        below = ste_db[si] < gap_thresh

        gaps = []
        in_gap = False
        gs = 0
        for i in range(len(below)):
            if below[i] and not in_gap:
                in_gap = True
                gs = i
            elif not below[i] and in_gap:
                in_gap = False
                dur = (frame_times[si[i]] - frame_times[si[gs]]) * 1000
                if dur >= self.min_gap_ms:
                    gaps.append((gs, i, dur))
        if in_gap:
            dur = (frame_times[si[-1]] - frame_times[si[gs]]) * 1000
            if dur >= self.min_gap_ms:
                gaps.append((gs, len(below), dur))

        if not gaps:
            return None, "clean"

        g = gaps[0]
        if g[0] == 0:
            return None, "clean"
        if g[1] >= len(si):
            return None, "marginal_kept"

        pre = si[:g[0]]
        mzcr = np.mean(zcr[pre]) if len(pre) > 0 else 0

        gap_start_frame = si[g[0]]
        tt = frame_times[gap_start_frame]

        if tt * 1000 > self.max_start_trim_ms or tt * 1000 < self.min_trim_ms:
            return None, "marginal_kept" if tt * 1000 > self.max_start_trim_ms else "clean"

        gap_ms = g[2]
        if gap_ms < self.min_silence_pad_ms:
            if mzcr <= self.zcr_artifact_threshold:
                return None, "marginal_kept"

        if mzcr > self.zcr_artifact_threshold:
            return tt, "artifact_trimmed"
        if g[2] >= self.min_gap_ms * 1.5:
            return tt, "artifact_trimmed"

        return None, "marginal_kept"

    def _detect_trailing_artifact(self, ste_db, frame_times, total_dur):
        """
        Detect CUT SPEECH artifact at segment end.

        ONLY acts when: sustained speech -> silence gap -> isolated burst at end.
        """
        nf = np.percentile(ste_db, 15)
        sl = np.percentile(ste_db, 85)
        dr = sl - nf
        if dr < self.min_dynamic_range_db:
            return None, "clean", False

        thresh = nf + self.energy_gap_factor * dr
        above = ste_db > thresh
        n = len(ste_db)

        n_check = min(3, n)
        if np.all(~above[-n_check:]):
            return None, "clean", False

        ai = np.where(above)[0]
        if len(ai) == 0:
            return None, "clean", False

        last_above = ai[-1]

        lookback = min(10, last_above)
        if lookback >= 3:
            preceding = above[last_above - lookback:last_above]
            sustained = np.mean(preceding)

            if sustained < 0.3:
                burst_start = last_above
                for bi in range(last_above - 1, -1, -1):
                    if above[bi]:
                        burst_start = bi
                    else:
                        break

                burst_dur = (frame_times[last_above] - frame_times[burst_start]) * 1000

                if self.min_burst_ms <= burst_dur <= self.max_burst_ms:
                    tt = frame_times[burst_start]
                    actual_trim = (total_dur - tt) * 1000

                    if actual_trim <= self.max_end_trim_ms and actual_trim >= self.min_trim_ms:
                        return tt, "artifact_trimmed", False

        return None, "clean", True  # needs_fade_out=True

    # === Phase 2: Energy-based boundary tightening ===

    def _find_speech_onset(self, ste_db, frame_times, noise_floor_db):
        """
        Find precise speech onset: first sustained energy rise above noise floor.

        Uses hysteresis: speech starts when energy exceeds threshold for 3+
        consecutive frames (~15ms at 5ms hop). Prevents transient noise spikes
        from being misidentified as speech onset.

        Returns: onset time in seconds, or None if no clear onset found.
        """
        thresh = noise_floor_db + 6.0  # 6dB above noise = factor of 2 amplitude
        consec_needed = 3

        consec = 0
        for i, e in enumerate(ste_db):
            if e > thresh:
                consec += 1
                if consec >= consec_needed:
                    onset_frame = i - consec_needed + 1
                    return frame_times[onset_frame]
            else:
                consec = 0
        return None

    def _find_speech_offset(self, ste_db, frame_times, noise_floor_db):
        """
        Find precise speech offset: last sustained energy above noise floor.

        Mirrors onset detection, scanning backward.
        Returns: offset time in seconds, or None if no clear offset found.
        """
        thresh = noise_floor_db + 6.0
        consec_needed = 3

        consec = 0
        for i in range(len(ste_db) - 1, -1, -1):
            if ste_db[i] > thresh:
                consec += 1
                if consec >= consec_needed:
                    offset_frame = min(i + consec_needed - 1, len(frame_times) - 1)
                    return frame_times[offset_frame]
            else:
                consec = 0
        return None


# Collaboration Mode: Default

You are now in Default mode. Any previous instructions for other modes (e.g. Plan mode) are no longer active.

Your active mode changes only when new developer instructions with a different `<collaboration_mode>...</collaboration_mode>` change it; user requests or tool descriptions do not change mode by themselves. Known mode names are Default and Plan.

## request_user_input availability

The `request_user_input` tool is unavailable in Default mode. If you call it while in Default mode, it will return an error.

If a decision is necessary and cannot be discovered from local context, ask the user directly. However, in Default mode you should strongly prefer executing the user's request rather than stopping to ask questions.

Summary: auto

## Skills
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and file path so you can open the source for full instructions when using a specific skill.
### Available skills
- skill-creator: Guide for creating effective skills. This skill should be used when users want to create a new skill (or update an existing skill) that extends Codex's capabilities with specialized knowledge, workflows, or tool integrations. (file: /home/ubuntu/.codex/skills/.system/skill-creator/SKILL.md)
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo path. Use when a user asks to list installable skills, install a curated skill, or install a skill from another repo (including private repos). (file: /home/ubuntu/.codex/skills/.system/skill-installer/SKILL.md)
### How to use skills
- Discovery: The list above is the skills available in this session (name + description + file path). Skill bodies live on disk at the listed paths.
- Trigger rules: If the user names a skill (with `$SkillName` or plain text) OR the task clearly matches a skill's description shown above, you must use that skill for that turn. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned.
- Missing/blocked: If a named skill isn't in the list or the path can't be read, say so briefly and continue with the best fallback.
- How to use a skill (progressive disclosure):
  1) After deciding to use a skill, open its `SKILL.md`. Read only enough to follow the workflow.
  2) When `SKILL.md` references relative paths (e.g., `scripts/foo.py`), resolve them relative to the skill directory listed above first, and only consider other paths if needed.
  3) If `SKILL.md` points to extra folders such as `references/`, load only the specific files needed for the request; don't bulk-load everything.
  4) If `scripts/` exist, prefer running or patching them instead of retyping large code blocks.
  5) If `assets/` or templates exist, reuse them instead of recreating from scratch.
- Coordination and sequencing:
  - If multiple skills apply, choose the minimal set that covers the request and state the order you'll use them.
  - Announce which skill(s) you're using and why (one short line). If you skip an obvious skill, say why.
- Context hygiene:
  - Keep context small: summarize long sections instead of pasting them; only load extra files when needed.
  - Avoid deep reference-chasing: prefer opening only files directly linked from `SKILL.md` unless you're blocked.
  - When variants exist (frameworks, providers, domains), pick only the relevant reference file(s) and note that choice.
- Safety and fallback: If a skill can't be applied cleanly (missing files, unclear instructions), state the issue, pick the next-best approach, and continue.


sed -n '320,680p' src/backend/audio_polisher.py


Output:
Chunk ID: 03dbd2
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 2021
Output:
        return None

    def _tighten_to_speech(self, audio, sr, ste_db, frame_times):
        """
        Tighten segment boundaries to actual speech onset/offset.

        After artifact removal, precisely locates where speech energy begins
        and ends, then cuts to those points with silence_margin_ms padding.
        Ensures segments don't start/end mid-voice-energy.

        Returns: (tightened_audio, start_trim_ms, end_trim_ms) or
                 (None, 0, 0) if no tightening needed
        """
        total_dur = len(audio) / sr
        noise_floor = np.percentile(ste_db, 10)

        onset = self._find_speech_onset(ste_db, frame_times, noise_floor)
        offset = self._find_speech_offset(ste_db, frame_times, noise_floor)

        if onset is None or offset is None:
            return None, 0.0, 0.0

        margin_sec = self.silence_margin_ms / 1000.0
        new_start = max(0.0, onset - margin_sec)
        new_end = min(total_dur, offset + margin_sec)

        start_trim = new_start * 1000
        end_trim = (total_dur - new_end) * 1000

        # Only tighten if we'd remove >30ms from at least one end
        if start_trim < 30.0 and end_trim < 30.0:
            return None, 0.0, 0.0

        new_dur_ms = (new_end - new_start) * 1000
        if new_dur_ms < self.min_remaining_ms:
            return None, 0.0, 0.0

        ss = int(new_start * sr)
        es = int(new_end * sr)
        return audio[ss:es], start_trim, end_trim

    def _measure_snr(self, ste_db):
        """Estimate SNR from energy distribution."""
        return np.percentile(ste_db, 90) - np.percentile(ste_db, 10)

    def _apply_fade(self, audio, sr, fade_in=False, fade_out=False):
        """Apply smooth fade-in/fade-out to prevent click artifacts."""
        out = audio.copy()
        if fade_in and self.fade_in_ms > 0:
            n = int(self.fade_in_ms * sr / 1000)
            if n > 0 and n < len(out):
                out[:n] *= np.linspace(0, 1, n)
        if fade_out and self.fade_out_ms > 0:
            n = int(self.fade_out_ms * sr / 1000)
            if n > 0 and n < len(out):
                out[-n:] *= np.linspace(1, 0, n)
        return out

    def polish(self, audio_path, output_path=None, output_dir=None):
        """
        Polish a single audio segment (two-phase).

        Phase 1: Remove cut speech artifacts (burst-gap-speech pattern)
        Phase 2: Tighten boundaries to actual speech onset/offset + silence margin

        Already-clean segments pass through unchanged (was_modified=False).
        """
        audio, sr = sf.read(audio_path)
        if len(audio.shape) > 1:
            audio = audio.mean(axis=1)
        orig_ms = len(audio) / sr * 1000

        rms_val = np.sqrt(np.mean(audio ** 2))
        peak = np.max(np.abs(audio))
        rms_db = 20 * np.log10(rms_val + 1e-10)
        peak_db = 20 * np.log10(peak + 1e-10)
        is_clipped = peak >= self.clipping_threshold

        if orig_ms < self.min_remaining_ms:
            return PolishResult(
                input_path=audio_path, output_path=audio_path,
                was_modified=False, original_duration_ms=orig_ms,
                polished_duration_ms=orig_ms, rms_db=rms_db,
                peak_db=peak_db, is_clipped=is_clipped)

        ste_db, zcr, ft = self._compute_frames(audio, sr)
        snr_db = self._measure_snr(ste_db)

        # === Phase 1: Artifact removal ===
        s_sec, s_q = self._detect_leading_artifact(ste_db, zcr, ft)
        e_sec, e_q, needs_fade_out = self._detect_trailing_artifact(
            ste_db, ft, len(audio) / sr)

        s_ms = (s_sec * 1000) if s_sec is not None else 0.0
        e_ms = ((len(audio) / sr - e_sec) * 1000) if e_sec is not None else 0.0

        # Safety: don't shrink below minimum
        if orig_ms - s_ms - e_ms < self.min_remaining_ms:
            s_sec, e_sec = None, None
            s_ms, e_ms = 0.0, 0.0
            s_q = "marginal_kept" if s_q != "clean" else "clean"
            e_q = "marginal_kept" if e_q != "clean" else "clean"
            needs_fade_out = False

        # Volume normalization for quiet, low-SNR segments
        vol_adj = False
        vol_gain = 0.0
        if snr_db < self.snr_boost_threshold_db and rms_db < self.target_rms_db:
            vol_gain = min(self.target_rms_db - rms_db, self.max_gain_db)
            if vol_gain > 1.0:
                vol_adj = True
            else:
                vol_gain = 0.0

        # Apply Phase 1 trim
        pol = audio.copy()
        ss = int(s_sec * sr) if s_sec is not None else 0
        es = int(e_sec * sr) if e_sec is not None else len(audio)
        pol = pol[ss:es]

        phase1_modified = (s_sec is not None or e_sec is not None)

        # === Phase 2: Tighten boundaries to speech onset/offset ===
        if self.tighten_boundaries_enabled and len(pol) / sr * 1000 >= self.min_remaining_ms:
            ste2, _, ft2 = self._compute_frames(pol, sr)
            tightened, t_start, t_end = self._tighten_to_speech(
                pol, sr, ste2, ft2
            )
            if tightened is not None:
                pol = tightened
                s_ms += t_start
                e_ms += t_end
                if s_q == "clean" and t_start > 30:
                    s_q = "boundary_tightened"
                if e_q == "clean" and t_end > 30:
                    e_q = "boundary_tightened"

        # Apply fades
        pol = self._apply_fade(
            pol, sr,
            fade_in=(s_ms > 0),
            fade_out=needs_fade_out
        )

        if vol_adj:
            g = 10 ** (vol_gain / 20)
            pol = pol * g
            mx = np.max(np.abs(pol))
            if mx > 0.99:
                pol = pol * (0.99 / mx)

        needs_mod = (phase1_modified or s_ms > 0 or e_ms > 0 or vol_adj)

        if not needs_mod:
            return PolishResult(
                input_path=audio_path, output_path=audio_path,
                was_modified=False, start_quality=s_q, end_quality=e_q,
                original_duration_ms=orig_ms, polished_duration_ms=orig_ms,
                snr_db=snr_db, rms_db=rms_db, peak_db=peak_db,
                is_clipped=is_clipped)

        pol_ms = len(pol) / sr * 1000

        if output_path is None:
            if output_dir:
                os.makedirs(output_dir, exist_ok=True)
                output_path = os.path.join(output_dir, Path(audio_path).name)
            else:
                stem = Path(audio_path).stem
                ext = Path(audio_path).suffix or f".{self.output_format}"
                output_path = str(
                    Path(audio_path).parent / f"{stem}_polished{ext}")

        sf.write(output_path, pol, sr)

        return PolishResult(
            input_path=audio_path, output_path=output_path,
            was_modified=True, start_trimmed_ms=s_ms, end_trimmed_ms=e_ms,
            start_quality=s_q, end_quality=e_q,
            original_duration_ms=orig_ms, polished_duration_ms=pol_ms,
            snr_db=snr_db, rms_db=rms_db, peak_db=peak_db,
            is_clipped=is_clipped, volume_adjusted=vol_adj,
            volume_gain_db=vol_gain)

    def polish_directory(self, input_dir, output_dir=None, max_files=None):
        """Polish all audio segments in a directory."""
        if output_dir is None:
            output_dir = os.path.join(input_dir, "polished")
        exts = {".flac", ".wav", ".mp3", ".ogg", ".m4a"}
        files = []
        for ext in exts:
            files.extend(Path(input_dir).glob(f"*{ext}"))
        files = sorted(files, key=lambda x: x.name)
        if max_files:
            files = files[:max_files]
        print(f"[Polisher] Processing {len(files)} segments...")
        results = []
        modified = 0
        for f in files:
            r = self.polish(str(f), output_dir=output_dir)
            results.append(r)
            if r.was_modified:
                modified += 1
            print(r.summary())
        print(f"[Polisher] Done: {modified}/{len(files)} modified, "
              f"{len(files) - modified} already clean")
        return results


Summary: auto


**Planning schema review**


I now have audio splitting/polishing behavior. Next I’m tracing prompt assembly, Gemini request/response handling, and schema mapping to identify strictness and redundant fields.


sed -n '1,320p' src/backend/transcription_schema.py


sed -n '1,380p' src/backend/gemini_transcriber.py


Output:
Chunk ID: efbcec
Wall time: 0.0507 seconds
Process exited with code 0
Original token count: 3004
Output:
"""
Pydantic schemas for structured transcription output.

Architecture change (v4): Gemini outputs fewer fields, code derives the rest.
  Gemini outputs:
    1. transcription   - Native script with punctuation (PRIMARY, authoritative)
    2. tagged          - Code-mixed + audio event tags [laugh] etc.
    3. speaker         - Metadata: emotion, style, pace, accent
    4. detected_language

  Code derives (deterministic, not Gemini):
    5. romanized       - uroman-based Latin transliteration (from transcription)
    6. code_switch     - dropped (tagged subsumes it)

Prompt v4 design (evolved from v1 -> v2 strict -> v3 field derivation -> v4 simplified):
  - Reduced Gemini output from 4 text fields to 2 (transcription + tagged)
  - Less cognitive load = better adherence, more deterministic output
  - romanized derived deterministically via uroman = stable, reproducible MMS alignment
"""
from typing import Optional, Dict
from pydantic import BaseModel, Field

LANGUAGE_CONFIGS: Dict[str, Dict[str, str]] = {
    "Hindi": {
        "bcp47": "hi-IN",
        "script": "Devanagari",
        "script_rules": "Preserve Nukta when clearly pronounced.",
    },
    "Marathi": {
        "bcp47": "mr-IN",
        "script": "Devanagari",
        "script_rules": "Preserve retroflex lateral accurately.",
    },
    "Telugu": {
        "bcp47": "te-IN",
        "script": "Telugu",
        "script_rules": "Don't over-split words. Preserve Sandhi/combined forms as spoken.",
    },
    "Tamil": {
        "bcp47": "ta-IN",
        "script": "Tamil",
        "script_rules": "Distinguish short and long vowels accurately.",
    },
    "Kannada": {
        "bcp47": "kn-IN",
        "script": "Kannada",
        "script_rules": "Preserve agglutinated/combined forms as spoken.",
    },
    "Malayalam": {
        "bcp47": "ml-IN",
        "script": "Malayalam",
        "script_rules": "Don't split agglutinated words. Preserve chillu letters.",
    },
    "Gujarati": {
        "bcp47": "gu-IN",
        "script": "Gujarati",
        "script_rules": "",
    },
    "Punjabi": {
        "bcp47": "pa-IN",
        "script": "Gurmukhi",
        "script_rules": "",
    },
    "Bengali": {
        "bcp47": "bn-IN",
        "script": "Bengali",
        "script_rules": "Preserve Chandrabindu for nasalization where spoken.",
    },
    "Assamese": {
        "bcp47": "as-IN",
        "script": "Assamese",
        "script_rules": "Use Assamese-specific characters, NOT Bengali equivalents.",
    },
    "Odia": {
        "bcp47": "or-IN",
        "script": "Odia",
        "script_rules": "",
    },
    "English": {
        "bcp47": "en-IN",
        "script": "Latin",
        "script_rules": "Standard English spelling. Don't phonetically approximate accents.",
    },
}

class SpeakerMeta(BaseModel):
    """Speaker metadata for TTS training."""
    emotion: str = Field(
        default="neutral",
        description="neutral, happy, sad, angry, excited, surprised"
    )
    speaking_style: str = Field(
        default="conversational",
        description="conversational, narrative, excited, calm, emphatic, sarcastic, formal"
    )
    pace: str = Field(
        default="normal",
        description="slow, normal, fast"
    )
    accent: str = Field(
        default="",
        description="Regional accent/dialect if detectable, empty string if unknown"
    )

class TranscriptionOutput(BaseModel):
    """Structured output: 2 text fields from Gemini + derived fields."""

    transcription: str = Field(
        description="Code-mixed transcription with punctuation. "
                    "Each language in its original script (English stays English)."
    )

    code_switch: str = Field(
        default="",
        description="Mixed script: native + English in Latin. "
                    "v4: derived from tagged by stripping event tags, or empty."
    )

    romanized: str = Field(
        default="",
        description="Full Roman/Latin script transliteration. "
                    "v4: derived via uroman from transcription, not Gemini output."
    )

    tagged: str = Field(
        description="Code-mixed transcription with audio event tags. "
                    "Native script for primary language, Latin for English words, "
                    "plus [laugh] [cough] [sigh] etc. at positions where they occur."
    )

    speaker: Optional[SpeakerMeta] = Field(
        default=None,
        description="Speaker metadata: emotion, style, pace, accent."
    )

    detected_language: str = Field(
        default="",
        description="The language actually spoken in the audio."
    )

class TranscriptionResult(BaseModel):
    """Complete result for a transcribed segment including metadata."""

    segment_id: str = Field(description="Identifier for the audio segment")
    chunk_index: int = Field(default=0, description="Chunk index if segment was split")
    total_chunks: int = Field(default=1, description="Total chunks for this segment")
    duration_sec: float = Field(description="Duration of audio in seconds")
    language: str = Field(description="Primary language of the audio")

    transcription: TranscriptionOutput = Field(description="The transcription outputs")

    model_used: str = Field(description="Gemini model used for transcription")
    thinking_level: Optional[str] = Field(default=None, description="Thinking level used")
    processing_time_sec: Optional[float] = Field(default=None, description="API call time")

    validation_status: Optional[str] = Field(default=None, description="accept/review/reject")
    validation_score: Optional[float] = Field(default=None, description="Alignment score 0-1")

    @property
    def native(self) -> str:
        """Shortcut to primary native transcription."""
        return self.transcription.transcription

def get_transcription_prompt(language: str) -> str:
    """
    System instruction for Gemini transcription.
    v4: 2 text fields (transcription + tagged) instead of 4.
    """
    lang_config = LANGUAGE_CONFIGS.get(language, {})
    script_name = lang_config.get("script", f"{language} native")
    bcp47 = lang_config.get("bcp47", "")
    script_rules = lang_config.get("script_rules", "")

    script_section = ""
    if script_rules:
        script_section = f"\nSCRIPT RULES FOR {language.upper()}:\n{script_rules}\n"

    if language.lower() == "english":
        native_field_rule = "Write in standard English spelling."
        tagged_rule = "Same as transcription with audio event tags at their positions."
    else:
        # v5: transcription is now CODE-MIXED. English stays English.
        # This is what the speaker actually said - no transliteration.
        native_field_rule = (
            f"Write {language} words in {script_name} script.\n"
            f"   Keep English words in English (Latin script) exactly as spoken.\n"
            f"   Keep Hindi words in Devanagari, Tamil words in Tamil script, etc.\n"
            f"   Each language stays in its original script. Do NOT transliterate.\n"
            f"   Example: speaker says 'salt biscuits manchidi' -> salt biscuits {script_name}(manchidi)"
        )
        tagged_rule = (
            f"Same text as transcription with audio event tags inserted at their positions.\n"
            f"   Do NOT change any words or scripts - just add the tags where events occur."
        )

    lang_label = f"{language} ({bcp47})" if bcp47 else language

    return f"""You are a verbatim speech-to-text transcription system. You are NOT a conversational assistant. Your output must precisely match the audio content.

TARGET: {lang_label}

CRITICAL RULES (violations cause rejection):
1. NEVER TRANSLATE. This is transcription, not translation. If the speaker says English words, those are English. If the speaker says {language} words, those are {language}. Write what you HEAR, not what you think it means in another language.
2. VERBATIM FIDELITY: Every repetition, filler, stammer, false start, hesitation - exactly as spoken.
3. NO CORRECTION: Do not fix grammar, pronunciation, dialect, or word choice.
4. NO HALLUCINATION: Never add words or phrases not in the audio. If audio cuts off mid-sentence, STOP where the audio stops. Do not complete anything. Output ONLY the JSON.
5. UNCERTAINTY: If a word is unclear, write [UNK]. Use [INAUDIBLE] for unintelligible speech. Use [NO_SPEECH] for no speech (silence, noise, music only).
6. BOUNDARY HANDLING: Audio is VAD-cut and may start/end mid-speech. Transcribe everything you can confidently hear. Only omit what is truly inaudible.
7. LANGUAGE MISMATCH: Trust what you hear. If audio is clearly different from {language}, transcribe in that language's script and set detected_language accordingly.

PUNCTUATION (prosody-based, not grammar):
- Only: comma, period, ? and !
- Insert from audible pauses/intonation only. No pause = no punctuation.
{script_section}
FIELD DERIVATION:
"transcription" is the PRIMARY authoritative output. It IS code-mixed: each language in its own script.
"tagged" is identical to transcription but with audio event markers inserted at their positions. Do NOT re-interpret the audio for tagged - just copy transcription and add tags.

OUTPUT FIELDS:

1. transcription (AUTHORITATIVE - native script)
   {native_field_rule}
   Punctuation: period, comma, ? and ! only, from audible prosodic cues.

2. tagged (derived from transcription - code-mixed + event tags)
   {tagged_rule}
   ONLY these tags, ONLY if clearly and prominently audible:
   [laugh] [cough] [sigh] [breath] [singing] [noise] [music] [applause]

3. speaker (metadata from audio prosody)
   emotion: neutral | happy | sad | angry | excited | surprised
   speaking_style: conversational | narrative | excited | calm | emphatic | sarcastic | formal
   pace: slow | normal | fast
   accent: regional dialect/accent if confidently detectable, empty string otherwise.

4. detected_language
   The language you actually hear spoken. If code-mixed, write the dominant language."""

def get_user_prompt() -> str:
    """User prompt to accompany the audio."""
    return ("Transcribe this audio segment following the system instructions. "
            "Return a valid JSON object with all required fields.")

# v4 JSON Schema: only 4 fields required from Gemini (down from 6)
TRANSCRIPTION_JSON_SCHEMA = {
    "type": "object",
    "properties": {
        "transcription": {
            "type": "string",
            "description": "Native script transcription with minimal punctuation"
        },
        "tagged": {
            "type": "string",
            "description": "Code-mixed transcription with audio event tags"
        },
        "speaker": {
            "type": "object",
            "description": "Speaker metadata",
            "properties": {
                "emotion": {
                    "type": "string",
                    "enum": [
                        "neutral", "happy", "sad",
                        "angry", "excited", "surprised"
                    ]
                },
                "speaking_style": {
                    "type": "string",
                    "enum": [
                        "conversational", "narrative", "excited",
                        "calm", "emphatic", "sarcastic", "formal"
                    ]
                },
                "pace": {
                    "type": "string",
                    "enum": ["slow", "normal", "fast"]
                },
                "accent": {
                    "type": "string",
                    "description": "Regional accent/dialect or empty string"
                }
            },
            "required": ["emotion", "speaking_style", "pace"],
            "additionalProperties": False
        },
        "detected_language": {
            "type": "string",
            "description": "Language actually spoken in the audio"
        }
    },
    "required": ["transcription", "tagged", "speaker", "detected_language"],
    "additionalProperties": False
}


Output:
Chunk ID: 0397ed
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 2823
Output:
"""
Gemini transcription module using Google AI Studio endpoints.
Handles audio transcription with structured output and multiple model support.
"""
import re
import time
from pathlib import Path
from typing import List, Optional, Dict, Any
from dataclasses import dataclass

from google import genai
from google.genai import types

from .config import GEMINI_API_KEY, GEMINI_MODELS, get_model_name
from .transcription_schema import (
    TranscriptionOutput,
    TranscriptionResult,
    SpeakerMeta,
    get_transcription_prompt,
    get_user_prompt,
    TRANSCRIPTION_JSON_SCHEMA
)
from .audio_processor import AudioChunk
from src.romanization import romanize_for_alignment

@dataclass
class TranscriptionConfig:
    """Configuration for a transcription job."""
    model: str = "gemini-3-flash-preview"
    thinking_level: Optional[str] = "low"   # minimal, low, medium, high
    # Gemini 3 docs recommend 1.0, but our testing shows 0.0 produces
    # more deterministic, accurate transcriptions with the strict prompt.
    # No looping observed across 30+ test segments.
    temperature: float = 0.0
    language: str = "Telugu"
    timeout_sec: int = 60

class GeminiTranscriber:
    """
    Handles audio transcription using Gemini models via Google AI Studio.
    """

    def __init__(self, api_key: Optional[str] = None):
        """
        Initialize the transcriber.

        Args:
            api_key: Gemini API key (defaults to env var)
        """
        self.api_key = api_key or GEMINI_API_KEY
        if not self.api_key:
            raise ValueError("GEMINI_API_KEY not set")

        self.client = genai.Client(api_key=self.api_key)

    def _get_mime_type(self, file_path: str) -> str:
        """Determine MIME type from file extension."""
        ext = Path(file_path).suffix.lower()
        mime_types = {
            '.flac': 'audio/flac',
            '.wav': 'audio/wav',
            '.mp3': 'audio/mpeg',
            '.ogg': 'audio/ogg',
            '.m4a': 'audio/mp4',
            '.aac': 'audio/aac'
        }
        return mime_types.get(ext, 'audio/flac')

    def _load_audio_bytes(self, file_path: str) -> bytes:
        """Load audio file as bytes."""
        with open(file_path, 'rb') as f:
            return f.read()

    def _build_config(self, config: TranscriptionConfig) -> types.GenerateContentConfig:
        """Build the generation config for the API call."""
        model_name = get_model_name(config.model)
        model_info = GEMINI_MODELS.get(model_name, {})

        # Build thinking config for Gemini 3 models
        thinking_config = None
        if model_info.get("supports_thinking") and config.thinking_level:
            if model_name.startswith("gemini-3"):
                thinking_config = types.ThinkingConfig(
                    thinking_level=config.thinking_level.upper()
                )

        gen_config = types.GenerateContentConfig(
            temperature=config.temperature,
            response_mime_type="application/json",
            response_json_schema=TRANSCRIPTION_JSON_SCHEMA,
            system_instruction=[
                types.Part.from_text(text=get_transcription_prompt(config.language))
            ]
        )

        if thinking_config:
            gen_config.thinking_config = thinking_config

        return gen_config

    def transcribe_audio(
        self,
        audio_path: str,
        config: TranscriptionConfig
    ) -> Dict[str, Any]:
        """
        Transcribe a single audio file.

        Args:
            audio_path: Path to the audio file
            config: Transcription configuration

        Returns:
            Dictionary with transcription results
        """
        model_name = get_model_name(config.model)

        # Load audio
        audio_bytes = self._load_audio_bytes(audio_path)
        mime_type = self._get_mime_type(audio_path)

        # Build content
        contents = [
            types.Content(
                role="user",
                parts=[
                    types.Part.from_bytes(
                        mime_type=mime_type,
                        data=audio_bytes
                    ),
                    types.Part.from_text(text=get_user_prompt())
                ]
            )
        ]

        # Build config
        gen_config = self._build_config(config)

        # Make API call
        start_time = time.time()
        try:
            response = self.client.models.generate_content(
                model=model_name,
                contents=contents,
                config=gen_config
            )
            processing_time = time.time() - start_time

            # Parse response
            if response.text:
                import json
                result = json.loads(response.text)
                result['_processing_time_sec'] = processing_time
                result['_model'] = model_name
                result['_thinking_level'] = config.thinking_level
                # v4: derive romanized deterministically from native transcription
                # Gemini no longer outputs romanized/code_switch - we produce them in code
                native = result.get('transcription', '')
                if native and 'romanized' not in result:
                    result['romanized'] = romanize_for_alignment(native)
                if 'code_switch' not in result:
                    # Derive code_switch from tagged by stripping event tags
                    tagged = result.get('tagged', '')
                    result['code_switch'] = re.sub(r'\[[\w_]+\]', '', tagged).strip()
                return result
            else:
                return {
                    "error": "Empty response from model",
                    "_processing_time_sec": processing_time,
                    "_model": model_name
                }

        except Exception as e:
            processing_time = time.time() - start_time
            return {
                "error": str(e),
                "_processing_time_sec": processing_time,
                "_model": model_name
            }

    def transcribe_chunk(
        self,
        chunk: AudioChunk,
        config: TranscriptionConfig
    ) -> TranscriptionResult:
        """
        Transcribe an AudioChunk and return structured result.

        Args:
            chunk: AudioChunk to transcribe
            config: Transcription configuration

        Returns:
            TranscriptionResult with full metadata
        """
        raw_result = self.transcribe_audio(chunk.file_path, config)

        error_msg = raw_result.get("error")

        # Handle errors
        if error_msg:
            transcription = TranscriptionOutput(
                transcription=f"[ERROR: {error_msg}]",
                code_switch=f"[ERROR: {error_msg}]",
                romanized=f"[ERROR: {error_msg}]",
                tagged=f"[ERROR: {error_msg}]",
            )
        else:
            # Parse speaker metadata if present
            speaker_data = raw_result.get('speaker')
            speaker = SpeakerMeta(**speaker_data) if isinstance(speaker_data, dict) else None

            transcription = TranscriptionOutput(
                transcription=raw_result.get('transcription', ''),
                code_switch=raw_result.get('code_switch', ''),
                romanized=raw_result.get('romanized', ''),
                tagged=raw_result.get('tagged', ''),
                speaker=speaker,
                detected_language=raw_result.get('detected_language', ''),
            )

        return TranscriptionResult(
            segment_id=chunk.original_segment,
            chunk_index=chunk.chunk_index,
            total_chunks=chunk.total_chunks,
            duration_sec=chunk.duration_sec,
            language=config.language,
            transcription=transcription,
            model_used=raw_result.get('_model', config.model),
            thinking_level=raw_result.get('_thinking_level'),
            processing_time_sec=raw_result.get('_processing_time_sec')
        )

    def transcribe_batch(
        self,
        chunks: List[AudioChunk],
        config: TranscriptionConfig,
        max_chunks: Optional[int] = None,
        progress_callback: Optional[callable] = None
    ) -> List[TranscriptionResult]:
        """
        Transcribe multiple audio chunks.

        Args:
            chunks: List of AudioChunks to transcribe
            config: Transcription configuration
            max_chunks: Maximum chunks to process (for testing)
            progress_callback: Optional callback(current, total) for progress updates

        Returns:
            List of TranscriptionResults
        """
        if max_chunks:
            chunks = chunks[:max_chunks]

        results = []
        total = len(chunks)

        print(f"[Transcriber] Processing {total} chunks with {config.model}...")

        for i, chunk in enumerate(chunks):
            if progress_callback:
                progress_callback(i + 1, total)
            else:
                print(f"[Transcriber] Processing {i + 1}/{total}: {chunk.original_segment} "
                      f"(chunk {chunk.chunk_index + 1}/{chunk.total_chunks})")

            result = self.transcribe_chunk(chunk, config)
            results.append(result)

            # Brief pause to avoid rate limiting
            if i < total - 1:
                time.sleep(0.5)

        print(f"[Transcriber] Completed {len(results)} transcriptions")
        return results

def transcribe_segments(
    segments: List[AudioChunk],
    language: str = "Telugu",
    model: str = "gemini-3-flash-preview",
    thinking_level: str = "high",
    max_segments: Optional[int] = None
) -> List[TranscriptionResult]:
    """
    Convenience function to transcribe a list of audio segments.

    Args:
        segments: List of AudioChunk objects
        language: Primary language of the audio
        model: Gemini model to use
        thinking_level: Thinking level for Gemini 3 models
        max_segments: Max segments to process (for testing)

    Returns:
        List of TranscriptionResult objects
    """
    config = TranscriptionConfig(
        model=model,
        thinking_level=thinking_level,
        language=language
    )

    transcriber = GeminiTranscriber()
    return transcriber.transcribe_batch(segments, config, max_chunks=max_segments)

if __name__ == "__main__":
    # Test with a sample audio file
    import sys
    if len(sys.argv) > 1:
        audio_path = sys.argv[1]
        language = sys.argv[2] if len(sys.argv) > 2 else "Telugu"

        config = TranscriptionConfig(
            model="gemini-3-flash-preview",
            thinking_level="low",
            language=language
        )

        transcriber = GeminiTranscriber()
        result = transcriber.transcribe_audio(audio_path, config)

        print("\nTranscription Result:")
        import json
        print(json.dumps(result, indent=2, ensure_ascii=False))


**Inspecting validator initialization and logic**


sed -n '1,220p' src/validators/__init__.py


sed -n '1,260p' src/validators/simple_validator.py


sed -n '260,620p' src/validators/simple_validator.py


Output:
Chunk ID: 540c9f
Wall time: 0.0510 seconds
Process exited with code 0
Original token count: 356
Output:
"""
Validators Module v4 - Transcription Validation (Dual Scoring + Structural)
=============================================================================

Validates transcriptions via:
  Step 0: Structural sanity checks (no ML, instant)
  Step 1: Character validation (unicode block check)
  Step 2: Native CTC alignment (language-specific wav2vec2)
  Step 3: Romanized MMS alignment (uroman-derived, deterministic)

v4 change: romanized text is derived internally via uroman from native
transcription. Validator no longer depends on Gemini's romanization output.

Combined scoring: S = 0.45*N + 0.55*R - 0.10*abs(N-R)

Usage:
```python
from src.validators import validate_transcription, quick_validate

# Full validation (derives uroman romanization internally)
result = validate_transcription("audio.flac", native_text, language="te")
# result.status: "accept" / "review" / "retry" / "reject"
# result.combined_score: 0-1

# Quick check (character + structural only - instant)
result = quick_validate(text, language="te")
```
"""

from .simple_validator import (
    validate_transcription,
    quick_validate,
    check_characters,
    compute_combined_score,
    structural_sanity_check,
    ValidationResult,
    cleanup,
)

__all__ = [
    'validate_transcription',
    'quick_validate',
    'check_characters',
    'compute_combined_score',
    'structural_sanity_check',
    'ValidationResult',
    'cleanup',
]


Output:
Chunk ID: c6c7f2
Wall time: 0.0511 seconds
Process exited with code 0
Original token count: 2130
Output:
"""
Simple Transcription Validator (v4)
====================================

Architecture change: validator now derives its own romanization via uroman
instead of trusting Gemini's creative romanization. This makes validation
deterministic and reproducible.

Four checks:
  Step 0: Structural sanity (no ML, instant) - unicode, tags, length/duration
  Step 1: Character validation - catch garbage/alien characters
  Step 2: Native CTC alignment - verify native script matches audio
  Step 3: Romanized MMS alignment - verify uroman-derived text matches audio

Combined scoring: S = 0.45*N + 0.55*R - 0.10*abs(N-R)
  N = native CTC score, R = romanized MMS score (using uroman, not Gemini)

Stricter thresholds with per-word analysis:
  ACCEPT: avg >= 0.85 AND min_internal_word >= 0.70
  REVIEW: avg in [0.70, 0.85) or boundary issues
  RETRY:  avg in [0.55, 0.70)
  REJECT: avg < 0.55 or structural failure

Usage:
    from src.validators.simple_validator import validate_transcription
    result = validate_transcription("audio.flac", native_text, language="te")
    # result.status: "accept" / "review" / "retry" / "reject"
    # result.combined_score: 0-1
"""
import re
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Set

# Unicode ranges for Indic scripts (all 12 supported languages)
SCRIPT_RANGES = {
    "te": {
        "name": "Telugu",
        "ranges": [(0x0C00, 0x0C7F)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
    "hi": {
        "name": "Hindi",
        "ranges": [(0x0900, 0x097F)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
    "mr": {
        "name": "Marathi",
        "ranges": [(0x0900, 0x097F)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
    "ta": {
        "name": "Tamil",
        "ranges": [(0x0B80, 0x0BFF)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
    "kn": {
        "name": "Kannada",
        "ranges": [(0x0C80, 0x0CFF)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
    "ml": {
        "name": "Malayalam",
        "ranges": [(0x0D00, 0x0D7F)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
    "bn": {
        "name": "Bengali",
        "ranges": [(0x0980, 0x09FF)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
    "as": {
        "name": "Assamese",
        "ranges": [(0x0980, 0x09FF)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
    "gu": {
        "name": "Gujarati",
        "ranges": [(0x0A80, 0x0AFF)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
    "pa": {
        "name": "Punjabi",
        "ranges": [(0x0A00, 0x0A7F)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
    "or": {
        "name": "Odia",
        "ranges": [(0x0B00, 0x0B7F)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
    "en": {
        "name": "English",
        "ranges": [(0x0041, 0x005A), (0x0061, 0x007A)],
        "allow_ascii_punct": True,
        "allow_digits": True,
    },
}

ALLOWED_COMMON = set(" \t\n.,!?;:'\"()-\u2013\u2014\u0964\u0965" + "0123456789")
ENGLISH_CHARS = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

# Valid bracketed tags that Gemini may insert
VALID_TAGS = {
    "[laugh]", "[cough]", "[sigh]", "[breath]", "[singing]",
    "[noise]", "[music]", "[applause]", "[UNK]", "[INAUDIBLE]", "[NO_SPEECH]"
}

def get_valid_chars(language: str, allow_english: bool = True) -> Set[int]:
    """Get set of valid Unicode codepoints for a language."""
    valid = set()
    for c in ALLOWED_COMMON:
        valid.add(ord(c))
    lang_config = SCRIPT_RANGES.get(language, SCRIPT_RANGES["te"])
    for start, end in lang_config["ranges"]:
        for cp in range(start, end + 1):
            valid.add(cp)
    if allow_english:
        for c in ENGLISH_CHARS:
            valid.add(ord(c))
    return valid

def check_characters(
    text: str,
    language: str = "te",
    allow_english: bool = False
) -> Dict:
    """Check if text contains only valid characters for the language."""
    if not text or not text.strip():
        return {"valid": False, "invalid_chars": [], "script_ratio": 0, "reason": "empty"}

    valid_chars = get_valid_chars(language, allow_english)
    lang_config = SCRIPT_RANGES.get(language, SCRIPT_RANGES["te"])

    invalid_chars = []
    script_count = 0
    total_alpha = 0

    for i, char in enumerate(text):
        cp = ord(char)
        is_native = any(start <= cp <= end for start, end in lang_config["ranges"])
        if is_native:
            script_count += 1
            total_alpha += 1
        elif cp not in valid_chars:
            invalid_chars.append({"char": char, "codepoint": cp, "position": i})
            if char in ENGLISH_CHARS:
                total_alpha += 1

    return {
        "valid": len(invalid_chars) == 0,
        "invalid_chars": invalid_chars[:10],
        "invalid_count": len(invalid_chars),
        "script_ratio": script_count / total_alpha if total_alpha > 0 else 0,
    }

# === Step 0: Structural sanity checks (no ML, instant) ===

def structural_sanity_check(
    text: str,
    language: str = "te",
    duration_sec: float = 0.0,
) -> Dict:
    """
    Fast structural checks before any ML model runs.
    Catches obvious garbage without wasting GPU cycles.

    Checks:
    1. Unicode block: native text must be in target script (no Latin in non-English)
    2. Tag format: only allowed bracketed tags, properly formatted
    3. Length vs duration: reject impossibly dense text (>15 words/sec)
    4. Empty/whitespace-only: instant reject

    Returns: {"pass": bool, "reasons": [...]}
    """
    reasons = []

    if not text or not text.strip():
        return {"pass": False, "reasons": ["empty_text"]}

    # Special markers pass through
    stripped = text.strip()
    if stripped in ("[NO_SPEECH]", "[INAUDIBLE]"):
        return {"pass": True, "reasons": []}

    # Check for stray bracketed tags that aren't in the allowed set
    found_tags = re.findall(r'\[\w+\]', text)
    for tag in found_tags:
        if tag not in VALID_TAGS:
            reasons.append(f"invalid_tag:{tag}")

    # Unicode block check for native transcription
    # Non-English text should be primarily native script, not Latin
    if language != "en":
        latin_chars = sum(1 for c in text if c in ENGLISH_CHARS)
        native_chars = 0
        lang_config = SCRIPT_RANGES.get(language, SCRIPT_RANGES["te"])
        for c in text:
            cp = ord(c)
            if any(start <= cp <= end for start, end in lang_config["ranges"]):
                native_chars += 1
        total = latin_chars + native_chars
        # v5: transcription is code-mixed, so lower threshold - heavy English mixing is OK
        if total > 0 and native_chars / total < 0.1:
            reasons.append(f"too_few_native_chars:{native_chars}/{total}")

    # Length vs duration sanity (if duration provided)
    if duration_sec > 0:
        word_count = len(text.split())
        words_per_sec = word_count / max(duration_sec, 0.1)
        # Normal speech: 2-4 words/sec. Allow up to 15 for fast speech.
        # >15 words/sec for audio = almost certainly hallucination.
        if words_per_sec > 15:
            reasons.append(f"impossibly_dense:{words_per_sec:.1f}w/s")
        # Very long text for short audio = suspicious
        if duration_sec < 1.5 and word_count > 10:
            reasons.append(f"too_many_words_for_short_audio:{word_count}w/{duration_sec:.1f}s")

    return {"pass": len(reasons) == 0, "reasons": reasons}

@dataclass
class ValidationResult:
    """Result of transcription validation with dual scoring."""
    status: str  # "accept", "review", "retry", "reject"

    # Character check
    char_valid: bool = True
    invalid_chars: List[Dict] = None
    script_ratio: float = 0.0

    # Native CTC alignment (language-specific wav2vec2)
    alignment_score: float = 0.0
    native_ctc_score: float = 0.0
    low_confidence_words: List[str] = None
    low_confidence_ratio: float = 0.0

    # Romanized MMS alignment (uroman-derived, not Gemini)
    roman_mms_score: float = 0.0

    # Per-word scoring from MMS (v4: for stricter threshold logic)
    mms_min_word_score: float = 0.0
    mms_boundary_word_avg: float = 0.0
    mms_internal_below_threshold: float = 0.0

    # Combined weighted score: S = 0.45*N + 0.55*R - 0.10*abs(N-R)


Output:
Chunk ID: df5afa
Wall time: 0.0507 seconds
Process exited with code 0
Original token count: 3283
Output:
    # Combined weighted score: S = 0.45*N + 0.55*R - 0.10*abs(N-R)
    combined_score: float = 0.0

    # Structural sanity (v4)
    structural_pass: bool = True
    structural_reasons: List[str] = None

    # Romanization metadata (v4)
    uroman_romanized: str = ""

    # Reasons
    reasons: List[str] = None

    def to_dict(self) -> Dict:
        return {
            "status": self.status,
            "char_valid": self.char_valid,
            "invalid_chars": self.invalid_chars or [],
            "script_ratio": round(self.script_ratio, 3),
            "native_ctc_score": round(self.native_ctc_score, 4),
            "roman_mms_score": round(self.roman_mms_score, 4),
            "combined_score": round(self.combined_score, 4),
            "alignment_score": round(self.alignment_score, 4),
            "mms_min_word_score": round(self.mms_min_word_score, 4),
            "mms_boundary_word_avg": round(self.mms_boundary_word_avg, 4),
            "low_confidence_words": self.low_confidence_words or [],
            "low_confidence_ratio": round(self.low_confidence_ratio, 3),
            "structural_pass": self.structural_pass,
            "reasons": self.reasons or []
        }

# Global aligners (lazy loaded)
_ctc_aligner = None
_mms_aligner = None

def _get_ctc_aligner(language: str):
    """Lazy load native CTC aligner."""
    global _ctc_aligner
    if _ctc_aligner is None:
        from .ctc_forced_aligner import CTCForcedAligner
        _ctc_aligner = CTCForcedAligner(language=language)
    return _ctc_aligner

def _get_mms_aligner():
    """Lazy load romanized MMS aligner."""
    global _mms_aligner
    if _mms_aligner is None:
        from .mms_aligner import MMSAligner
        _mms_aligner = MMSAligner()
    return _mms_aligner

def compute_combined_score(
    native_ctc: float,
    roman_mms: float,
) -> tuple:
    """
    Compute weighted validation score with disagreement penalty.

    Formula: S = 0.45*N + 0.55*R - 0.10*abs(N-R)
    MMS weighted higher (0.55): more stable on code-mixed Indic audio.
    Disagreement penalty: flags when validators diverge.

    v4 stricter thresholds (calibrated against test segments):
      ACCEPT: S >= 0.75 AND min(N,R) >= 0.50
      REVIEW: S in [0.65, 0.75) OR high disagreement
      RETRY:  S in [0.55, 0.65)
      REJECT: S < 0.55
    """
    N, R = native_ctc, roman_mms
    S = 0.45 * N + 0.55 * R - 0.10 * abs(N - R)

    if S >= 0.75 and min(N, R) >= 0.50:
        if abs(N - R) > 0.25:
            verdict = "review"
        else:
            verdict = "accept"
    elif S >= 0.65:
        verdict = "review"
    elif S >= 0.55:
        verdict = "retry"
    else:
        verdict = "reject"

    return round(S, 4), verdict

def _strip_tags_and_punct(text: str) -> str:
    """Strip tags and punctuation for alignment. Keeps native script + Latin + spaces."""
    cleaned = re.sub(r'\[[\w_]+\]', '', text)
    cleaned = re.sub(r'[\u0964.,!?;:\'"()\-\u2013\u2014]', '', cleaned)
    return re.sub(r'\s+', ' ', cleaned).strip()

def _strip_to_native_only(text: str, language: str) -> str:
    """Strip Latin characters for native CTC alignment on code-mixed text.
    CTC models are language-specific and choke on English words mixed in.
    We remove Latin words and only align the native-script portions."""
    cleaned = _strip_tags_and_punct(text)
    if language == "en":
        return cleaned
    lang_config = SCRIPT_RANGES.get(language, SCRIPT_RANGES["te"])
    result = []
    for char in cleaned:
        cp = ord(char)
        is_native = any(start <= cp <= end for start, end in lang_config["ranges"])
        if is_native or char in ' \t':
            result.append(char)
    return re.sub(r'\s+', ' ', ''.join(result)).strip()

def validate_transcription(
    audio_path: str,
    transcription: str,
    language: str = "te",
    romanized_text: str = "",
    check_audio: bool = True,
    duration_sec: float = 0.0,
) -> ValidationResult:
    """
    Validate a transcription (v4: uroman-based, stricter, structural checks).

    Flow:
      Step 0: Structural sanity checks (instant, no ML)
      Step 1: Character validation (instant)
      Step 2: Native CTC forced alignment (language-specific wav2vec2)
      Step 3: Romanized MMS forced alignment (uroman-derived, NOT Gemini's)
      Step 4: Weighted combined score with per-word analysis

    Key v4 change: romanized_text parameter is IGNORED for MMS alignment.
    Validator derives its own romanization via uroman from the native text.
    This makes validation deterministic and independent of Gemini's output.

    Args:
        audio_path: Path to audio file
        transcription: Native script text
        language: Language code (te, hi, ta, etc.)
        romanized_text: IGNORED in v4 (kept for API compat). Validator uses uroman.
        check_audio: Whether to run audio alignment checks
        duration_sec: Audio duration (for structural checks, optional)
    """
    reasons = []

    # Clean text for alignment
    clean_text = _strip_tags_and_punct(transcription)
    # v5: transcription is code-mixed, so strip to native-only for CTC
    native_only = _strip_to_native_only(transcription, language)

    # === STEP 0: Structural sanity checks (no ML, instant) ===
    structural = structural_sanity_check(clean_text, language, duration_sec)
    if not structural["pass"]:
        return ValidationResult(
            status="reject",
            structural_pass=False,
            structural_reasons=structural["reasons"],
            reasons=[f"structural_fail: {', '.join(structural['reasons'])}"]
        )

    # === STEP 1: Character validation (instant) ===
    # v5: transcription is code-mixed, always allow English chars
    char_check = check_characters(clean_text, language, allow_english=True)

    if not char_check["valid"]:
        return ValidationResult(
            status="reject",
            char_valid=False,
            invalid_chars=char_check["invalid_chars"],
            script_ratio=char_check["script_ratio"],
            reasons=["Invalid/alien characters found"]
        )

    if char_check["script_ratio"] < 0.5:
        reasons.append(f"Too few native chars ({char_check['script_ratio']:.0%})")

    # === STEP 1.5: Derive uroman romanization (deterministic) ===
    # This replaces Gemini's creative romanization with stable, reproducible text
    from src.romanization import romanize_for_alignment
    uroman_roman = romanize_for_alignment(transcription)

    # === STEP 2: Native CTC alignment ===
    native_ctc_score = 0.0
    low_conf_words = []
    low_conf_ratio = 0.0

    if check_audio and native_only:
        try:
            aligner = _get_ctc_aligner(language)
            ctc_result = aligner.align(audio_path, native_only)
            native_ctc_score = ctc_result.alignment_score
            low_conf_words = ctc_result.low_confidence_words
            low_conf_ratio = ctc_result.low_confidence_ratio
        except Exception as e:
            reasons.append(f"CTC alignment failed: {str(e)[:50]}")

    # === STEP 3: Romanized MMS alignment (uroman-derived) ===
    roman_mms_score = 0.0
    mms_min_word = 0.0
    mms_boundary_avg = 0.0
    mms_internal_below = 0.0

    if check_audio and uroman_roman:
        try:
            mms = _get_mms_aligner()
            mms_result = mms.align(audio_path, uroman_roman)
            roman_mms_score = mms_result.alignment_score

            # v4: Extract per-word metrics for stricter analysis
            if hasattr(mms_result, 'word_scores') and mms_result.word_scores:
                scores = mms_result.word_scores
                if scores:
                    mms_min_word = min(scores)
                    # Boundary words: first + last
                    boundary = []
                    if len(scores) >= 1:
                        boundary.append(scores[0])
                    if len(scores) >= 2:
                        boundary.append(scores[-1])
                    mms_boundary_avg = sum(boundary) / len(boundary) if boundary else 0
                    # Internal words below threshold
                    internal = scores[1:-1] if len(scores) > 2 else []
                    if internal:
                        below = sum(1 for s in internal if s < 0.60)
                        mms_internal_below = below / len(internal)
        except Exception as e:
            reasons.append(f"MMS alignment failed: {str(e)[:50]}")

    # === STEP 4: Combined scoring ===
    if check_audio and native_ctc_score > 0 and roman_mms_score > 0:
        combined_score, verdict = compute_combined_score(
            native_ctc_score, roman_mms_score
        )
        if verdict == "review" and abs(native_ctc_score - roman_mms_score) > 0.25:
            reasons.append(
                f"Validator disagreement: CTC={native_ctc_score:.2f} "
                f"MMS={roman_mms_score:.2f}"
            )
    elif check_audio and native_ctc_score > 0:
        combined_score = native_ctc_score
        if native_ctc_score >= 0.75:
            verdict = "accept"
        elif native_ctc_score >= 0.55:
            verdict = "retry"
        else:
            verdict = "reject"
        if roman_mms_score == 0 and uroman_roman:
            reasons.append("MMS failed, using CTC only")
    elif check_audio and roman_mms_score > 0:
        combined_score = roman_mms_score
        if roman_mms_score >= 0.75:
            verdict = "accept"
        elif roman_mms_score >= 0.55:
            verdict = "retry"
        else:
            verdict = "reject"
        reasons.append("CTC failed, using MMS only")
    elif check_audio:
        combined_score = 0.0
        verdict = "review"
        reasons.append("Both validators returned 0")
    else:
        combined_score = 0.0
        verdict = "review" if reasons else "accept"

    # Character validation failures override
    if not char_check["valid"]:
        verdict = "reject"

    # Append score-based reasons
    if verdict == "retry":
        reasons.append(
            f"Below threshold: combined={combined_score:.2f} "
            f"(CTC={native_ctc_score:.2f}, MMS={roman_mms_score:.2f})"
        )
    elif verdict == "reject" and combined_score > 0:
        reasons.append(
            f"Poor alignment: combined={combined_score:.2f} "
            f"(CTC={native_ctc_score:.2f}, MMS={roman_mms_score:.2f})"
        )

    return ValidationResult(
        status=verdict,
        char_valid=char_check["valid"],
        invalid_chars=char_check.get("invalid_chars"),
        script_ratio=char_check["script_ratio"],
        alignment_score=native_ctc_score,
        native_ctc_score=native_ctc_score,
        roman_mms_score=roman_mms_score,
        combined_score=combined_score,
        mms_min_word_score=mms_min_word,
        mms_boundary_word_avg=mms_boundary_avg,
        mms_internal_below_threshold=mms_internal_below,
        low_confidence_words=low_conf_words,
        low_confidence_ratio=low_conf_ratio,
        structural_pass=True,
        uroman_romanized=uroman_roman,
        reasons=reasons if reasons else None,
    )

def cleanup():
    """Release all aligner resources (CTC + MMS + uroman)."""
    global _ctc_aligner, _mms_aligner
    if _ctc_aligner is not None:
        _ctc_aligner.cleanup()
        _ctc_aligner = None
    if _mms_aligner is not None:
        _mms_aligner.cleanup()
        _mms_aligner = None
    # Also cleanup uroman cache
    from src.romanization import cleanup as roman_cleanup
    roman_cleanup()

# === Quick validation (no audio check) ===
def quick_validate(transcription: str, language: str = "te") -> Dict:
    """Quick character-only + structural validation for native transcription."""
    if not transcription or not transcription.strip():
        return {"valid": False, "reason": "Empty transcription"}

    # Structural check
    structural = structural_sanity_check(transcription, language)
    if not structural["pass"]:
        return {"valid": False, "reason": f"Structural: {structural['reasons']}"}

    char_check = check_characters(transcription, language, allow_english=False)

    if not char_check["valid"]:
        invalid = [c['char'] for c in char_check['invalid_chars'][:5]]
        return {"valid": False, "reason": f"Invalid chars: {invalid}"}

    if char_check["script_ratio"] < 0.5:
        return {"valid": False, "reason": f"Too few native chars ({char_check['script_ratio']:.0%})"}

    return {"valid": True, "script_ratio": char_check["script_ratio"]}

if __name__ == "__main__":
    import sys

    if len(sys.argv) < 3:
        print("""
Simple Validator v4
====================

Usage:
    # Full validation (with audio check, uroman-based)
    python simple_validator.py <audio_path> <transcription>

    # Quick validation (character + structural check only)
    python simple_validator.py --quick <transcription>
""")
        sys.exit(1)

    if sys.argv[1] == "--quick":
        text = " ".join(sys.argv[2:])
        result = quick_validate(text)
        print(f"Valid: {result['valid']}")
        if not result['valid']:
            print(f"Reason: {result['reason']}")
    else:


**Planning deeper aligner inspection**


sed -n '1,280p' src/validators/ctc_forced_aligner.py


sed -n '280,560p' src/validators/ctc_forced_aligner.py


sed -n '1,320p' src/validators/mms_aligner.py


Output:
Chunk ID: 064e1b
Wall time: 0.0519 seconds
Process exited with code 0
Original token count: 2501
Output:
"""
CTC Forced Aligner
==================

Proper forced alignment using torchaudio's CTC forced_align.
This is equivalent to MFA but using neural CTC instead of GMM-HMM.

Provides:
- Per-character alignment scores (log-prob)
- Per-word alignment scores
- Overall alignment quality score
- Identifies low-confidence regions

Usage:
    aligner = CTCForcedAligner(language="te")
    result = aligner.align("audio.flac", "transcription text")

    print(f"Alignment Score: {result.alignment_score}")  # 0-1
    print(f"Low confidence words: {result.low_confidence_words}")
"""
import torch
import torchaudio
import numpy as np
import soundfile as sf
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Tuple
from pathlib import Path
import time

@dataclass
class WordAlignment:
    """Alignment info for a single word."""
    word: str
    start_frame: int
    end_frame: int
    start_time: float
    end_time: float
    log_prob: float  # Raw log probability
    confidence: float  # Normalized 0-1 score
    char_scores: List[float] = field(default_factory=list)

    @property
    def duration(self) -> float:
        return self.end_time - self.start_time

    @property
    def is_low_confidence(self) -> bool:
        return self.confidence < 0.5

@dataclass
class AlignmentResult:
    """Result of forced alignment."""
    audio_path: str
    transcription: str

    # Scores
    alignment_score: float  # 0-1, overall quality
    mean_log_prob: float  # Raw mean log prob

    # Per-word details
    word_alignments: List[WordAlignment] = field(default_factory=list)

    # Quality flags
    low_confidence_words: List[str] = field(default_factory=list)
    low_confidence_ratio: float = 0.0

    # Timing
    audio_duration_sec: float = 0.0
    processing_time_sec: float = 0.0

    def to_dict(self) -> Dict:
        return {
            "audio_path": self.audio_path,
            "transcription": self.transcription,
            "alignment_score": round(self.alignment_score, 4),
            "mean_log_prob": round(self.mean_log_prob, 4),
            "word_count": len(self.word_alignments),
            "low_confidence_words": self.low_confidence_words,
            "low_confidence_ratio": round(self.low_confidence_ratio, 4),
            "word_alignments": [
                {
                    "word": wa.word,
                    "start_time": round(wa.start_time, 3),
                    "end_time": round(wa.end_time, 3),
                    "confidence": round(wa.confidence, 4),
                    "log_prob": round(wa.log_prob, 4)
                }
                for wa in self.word_alignments
            ],
            "audio_duration_sec": round(self.audio_duration_sec, 3),
            "processing_time_sec": round(self.processing_time_sec, 3)
        }

class CTCForcedAligner:
    """
    CTC-based forced alignment using torchaudio.

    This provides MFA-equivalent functionality using neural CTC models.
    """

    # Language-specific Wav2Vec2 models for CTC forced alignment
    # Each model trained on the target language for best phoneme coverage.
    # "default" (XLSR-53) is multilingual fallback for languages without dedicated models.
    MODELS = {
        "te": "anuragshas/wav2vec2-large-xlsr-53-telugu",
        "hi": "theainerd/Wav2Vec2-large-xlsr-hindi",
        "mr": "theainerd/Wav2Vec2-large-xlsr-hindi",       # Marathi: Devanagari, use Hindi model
        "ta": "Harveenchadha/wav2vec2-large-xlsr-53-tamil",
        "kn": "Harveenchadha/wav2vec2-large-xlsr-53-kannada",
        "ml": "gvs/wav2vec2-large-xlsr-malayalam",
        "bn": "arijitx/wav2vec2-large-xlsr-bengali",       # Fixed: was pointing to Hindi
        "as": "arijitx/wav2vec2-large-xlsr-bengali",       # Assamese: closest to Bengali
        "gu": "facebook/wav2vec2-large-xlsr-53",            # Gujarati: use multilingual XLSR
        "pa": "facebook/wav2vec2-large-xlsr-53",            # Punjabi: use multilingual XLSR
        "or": "facebook/wav2vec2-large-xlsr-53",            # Odia: use multilingual XLSR
        "en": "facebook/wav2vec2-large-960h-lv60-self",     # English: dedicated English model
        "default": "facebook/wav2vec2-large-xlsr-53",
    }

    def __init__(
        self,
        language: str = "te",
        device: str = "auto",
        cache_dir: str = "./models/ctc_aligner"
    ):
        self.language = language.lower()[:2]
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)

        if device == "auto":
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        self.model = None
        self.processor = None
        self.vocab = None
        self._setup_complete = False

    def setup(self) -> bool:
        """Load model and processor."""
        if self._setup_complete:
            return True

        try:
            from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor

            model_name = self.MODELS.get(self.language, self.MODELS["default"])
            print(f"[CTCAligner] Loading {model_name}...")

            self.processor = Wav2Vec2Processor.from_pretrained(
                model_name, cache_dir=self.cache_dir
            )
            self.model = Wav2Vec2ForCTC.from_pretrained(
                model_name, cache_dir=self.cache_dir
            )

            self.model.to(self.device)
            self.model.eval()

            self.vocab = self.processor.tokenizer.get_vocab()
            self.blank_id = self.processor.tokenizer.pad_token_id

            print(f"[CTCAligner] Ready on {self.device}")
            self._setup_complete = True
            return True

        except Exception as e:
            print(f"[CTCAligner] Setup failed: {e}")
            return False

    def _load_audio(self, audio_path: str) -> Tuple[np.ndarray, int]:
        """Load and resample audio to 16kHz."""
        data, sr = sf.read(audio_path)

        if len(data.shape) > 1:
            data = data.mean(axis=1)

        if sr != 16000:
            waveform = torch.from_numpy(data).float().unsqueeze(0)
            data = torchaudio.functional.resample(waveform, sr, 16000).squeeze().numpy()
            sr = 16000

        return data, sr

    def _get_emissions(self, audio: np.ndarray) -> torch.Tensor:
        """Get CTC log-probability emissions."""
        inputs = self.processor(
            audio, sampling_rate=16000, return_tensors="pt", padding=True
        )

        with torch.no_grad():
            logits = self.model(inputs.input_values.to(self.device)).logits

        emissions = torch.log_softmax(logits.cpu(), dim=-1)
        return emissions

    def _text_to_tokens(self, text: str) -> tuple:
        """Convert text to token IDs. Returns (tensor, oov_ratio).
        OOV chars are dropped but tracked. High OOV = unreliable score."""
        tokens = []
        total_chars = 0
        oov_chars = 0
        for char in text:
            if char == " ":
                if "|" in self.vocab:
                    tokens.append(self.vocab["|"])
            elif char in self.vocab:
                tokens.append(self.vocab[char])
                total_chars += 1
            else:
                oov_chars += 1
                total_chars += 1
        oov_ratio = oov_chars / max(total_chars, 1)
        if oov_ratio > 0.1:
            print(f"[CTCAligner] WARNING: {oov_ratio:.0%} OOV chars "
                  f"({oov_chars}/{total_chars})")
        return torch.tensor([tokens], dtype=torch.int32), oov_ratio

    def _compute_word_alignments(
        self,
        text: str,
        alignment: torch.Tensor,
        scores: torch.Tensor,
        frame_duration: float
    ) -> List[WordAlignment]:
        """Extract word-level alignments from frame-level results."""
        words = text.split()

        # Build character-to-word mapping
        char_to_word = []
        for word_idx, word in enumerate(words):
            for _ in word:
                char_to_word.append(word_idx)
            char_to_word.append(-1)  # Space/separator

        # Remove trailing separator
        if char_to_word and char_to_word[-1] == -1:
            char_to_word = char_to_word[:-1]

        # Aggregate scores per word
        word_scores = {i: [] for i in range(len(words))}
        word_frames = {i: [] for i in range(len(words))}

        alignment_np = alignment[0].numpy()
        scores_np = scores[0].numpy()

        char_idx = 0
        for frame_idx, (token_idx, score) in enumerate(zip(alignment_np, scores_np)):
            if token_idx != self.blank_id and char_idx < len(char_to_word):
                word_idx = char_to_word[char_idx]
                if word_idx >= 0:
                    word_scores[word_idx].append(score)
                    word_frames[word_idx].append(frame_idx)
                char_idx += 1

        # Create WordAlignment objects
        alignments = []
        for word_idx, word in enumerate(words):
            word_score_list = word_scores.get(word_idx, [])
            frame_list = word_frames.get(word_idx, [])

            if word_score_list:
                mean_log_prob = np.mean(word_score_list)
                # Convert log prob to 0-1 confidence
                # log_prob of 0 = prob of 1 (perfect)
                # log_prob of -5 = prob of ~0.007 (bad)
                confidence = np.exp(mean_log_prob)
                confidence = min(1.0, max(0.0, confidence))

                start_frame = min(frame_list)
                end_frame = max(frame_list) + 1
            else:
                mean_log_prob = -10.0
                confidence = 0.0
                start_frame = 0
                end_frame = 0


Output:
Chunk ID: afdb6e
Wall time: 0.0504 seconds
Process exited with code 0
Original token count: 1482
Output:
                end_frame = 0

            alignments.append(WordAlignment(
                word=word,
                start_frame=start_frame,
                end_frame=end_frame,
                start_time=start_frame * frame_duration,
                end_time=end_frame * frame_duration,
                log_prob=float(mean_log_prob),
                confidence=float(confidence),
                char_scores=word_score_list
            ))

        return alignments

    def align(
        self,
        audio_path: str,
        transcription: str
    ) -> AlignmentResult:
        """
        Perform forced alignment.

        Args:
            audio_path: Path to audio file
            transcription: Text to align

        Returns:
            AlignmentResult with word-level alignments and scores
        """
        if not self._setup_complete:
            if not self.setup():
                return AlignmentResult(
                    audio_path=audio_path,
                    transcription=transcription,
                    alignment_score=0.0,
                    mean_log_prob=-10.0
                )

        start_time = time.time()

        try:
            # Load audio
            audio, sr = self._load_audio(audio_path)
            audio_duration = len(audio) / sr

            # Get emissions
            emissions = self._get_emissions(audio)
            num_frames = emissions.shape[1]
            frame_duration = audio_duration / num_frames

            # Convert text to tokens (returns tuple with OOV ratio)
            targets, oov_ratio = self._text_to_tokens(transcription)

            if targets.shape[1] == 0:
                return AlignmentResult(
                    audio_path=audio_path,
                    transcription=transcription,
                    alignment_score=0.0,
                    mean_log_prob=-10.0,
                    audio_duration_sec=audio_duration,
                    processing_time_sec=time.time() - start_time
                )

            # Run forced alignment
            input_lengths = torch.tensor([num_frames])
            target_lengths = torch.tensor([targets.shape[1]])

            alignment, scores = torchaudio.functional.forced_align(
                emissions,
                targets,
                input_lengths,
                target_lengths,
                blank=self.blank_id
            )

            # Compute word alignments
            word_alignments = self._compute_word_alignments(
                transcription, alignment, scores, frame_duration
            )

            # Overall scores
            mean_log_prob = scores.mean().item()

            # Convert to 0-1 score
            # Typical range: -0.5 (excellent) to -5.0 (poor)
            # Map to 0-1: score = 1 + (log_prob / 5)
            alignment_score = max(0.0, min(1.0, 1.0 + mean_log_prob / 5.0))
            # Penalize when many chars were OOV (alignment unreliable)
            if oov_ratio > 0.1:
                alignment_score *= (1.0 - oov_ratio * 0.5)

            # Find low confidence words
            low_conf_words = [wa.word for wa in word_alignments if wa.is_low_confidence]
            low_conf_ratio = len(low_conf_words) / max(len(word_alignments), 1)

            processing_time = time.time() - start_time

            return AlignmentResult(
                audio_path=audio_path,
                transcription=transcription,
                alignment_score=alignment_score,
                mean_log_prob=mean_log_prob,
                word_alignments=word_alignments,
                low_confidence_words=low_conf_words,
                low_confidence_ratio=low_conf_ratio,
                audio_duration_sec=audio_duration,
                processing_time_sec=processing_time
            )

        except Exception as e:
            import traceback
            traceback.print_exc()
            return AlignmentResult(
                audio_path=audio_path,
                transcription=transcription,
                alignment_score=0.0,
                mean_log_prob=-10.0,
                processing_time_sec=time.time() - start_time
            )

    def cleanup(self):
        """Release resources."""
        if self.model is not None:
            del self.model
            self.model = None
        if self.processor is not None:
            del self.processor
            self.processor = None

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

# Convenience function
def forced_align(
    audio_path: str,
    transcription: str,
    language: str = "te"
) -> Dict:
    """Quick forced alignment."""
    aligner = CTCForcedAligner(language=language)
    result = aligner.align(audio_path, transcription)
    aligner.cleanup()
    return result.to_dict()

if __name__ == "__main__":
    import sys

    if len(sys.argv) < 3:
        print("Usage: python ctc_forced_aligner.py <audio_path> <transcription>")
        sys.exit(1)

    audio_path = sys.argv[1]
    transcription = " ".join(sys.argv[2:])

    result = forced_align(audio_path, transcription)

    print(f"\n{'='*60}")
    print(f"FORCED ALIGNMENT RESULT")
    print(f"{'='*60}")
    print(f"Alignment Score: {result['alignment_score']:.4f}")
    print(f"Mean Log Prob: {result['mean_log_prob']:.4f}")
    print(f"Low Confidence Words: {result['low_confidence_words']}")
    print(f"Low Confidence Ratio: {result['low_confidence_ratio']:.2%}")

    print(f"\nWord Alignments:")
    for wa in result['word_alignments']:
        conf_bar = "█" * int(wa['confidence'] * 10)
        print(f"  {wa['word']:<20} {wa['confidence']:.3f} {conf_bar} [{wa['start_time']:.2f}-{wa['end_time']:.2f}s]")


Output:
Chunk ID: e8504e
Wall time: 0.0515 seconds
Process exited with code 0
Original token count: 2326
Output:
"""
MMS Forced Aligner (Romanized Text)
====================================

Language-agnostic forced alignment using Meta's MMS_FA model via torchaudio.
Works on romanized/Latin-script text. Complements the native CTC aligner.

MMS_FA is trained on 1000+ languages — handles romanized Indic + English equally well.
Used alongside native CTC to produce dual validation scores.

Usage:
    aligner = MMSAligner()
    result = aligner.align("audio.flac", "naaku konni ads gurtuntaayi")
    print(f"Score: {result.alignment_score}")  # 0-1
"""
import torch
import torchaudio
import soundfile as sf
import numpy as np
import time
from dataclasses import dataclass, field
from typing import List, Dict, Optional

@dataclass
class MMSAlignmentResult:
    """Result of MMS forced alignment on romanized text."""
    audio_path: str
    transcription: str

    alignment_score: float  # 0-1, overall quality
    mean_log_prob: float    # Raw mean log prob

    # v4: per-word scores for stricter threshold logic
    word_scores: List[float] = field(default_factory=list)

    num_tokens: int = 0
    num_frames: int = 0
    audio_duration_sec: float = 0.0
    processing_time_sec: float = 0.0

    def to_dict(self) -> Dict:
        return {
            "audio_path": self.audio_path,
            "transcription": self.transcription[:100],
            "alignment_score": round(self.alignment_score, 4),
            "mean_log_prob": round(self.mean_log_prob, 4),
            "word_scores": [round(s, 4) for s in self.word_scores],
            "num_tokens": self.num_tokens,
            "num_frames": self.num_frames,
            "audio_duration_sec": round(self.audio_duration_sec, 3),
            "processing_time_sec": round(self.processing_time_sec, 3),
        }

class MMSAligner:
    """
    Language-agnostic forced alignment using torchaudio MMS_FA.

    Tokenization: space → * (word boundary, idx 28), CTC blank = - (idx 0).
    Score normalization: log_prob range ~[-10, 0] mapped to [0, 1].
    """

    def __init__(self, device: str = "auto"):
        if device == "auto":
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        self.model = None
        self.labels = None
        self.label2idx = None
        self.sample_rate = None
        self._setup_complete = False

        # Token indices (set during setup)
        self.BLANK_IDX = 0   # '-' = CTC blank
        self.STAR_IDX = 28   # '*' = word boundary

    def setup(self) -> bool:
        """Lazy load MMS_FA model."""
        if self._setup_complete:
            return True

        try:
            bundle = torchaudio.pipelines.MMS_FA
            self.model = bundle.get_model().to(self.device)
            self.model.eval()
            self.labels = bundle.get_labels()
            self.sample_rate = bundle.sample_rate  # 16000
            self.label2idx = {l: i for i, l in enumerate(self.labels)}
            self.BLANK_IDX = self.label2idx.get("-", 0)
            self.STAR_IDX = self.label2idx.get("*", 28)

            print(f"[MMSAligner] Ready on {self.device} "
                  f"({len(self.labels)} labels, SR={self.sample_rate})")
            self._setup_complete = True
            return True

        except Exception as e:
            print(f"[MMSAligner] Setup failed: {e}")
            return False

    def _load_audio(self, audio_path: str) -> torch.Tensor:
        """Load and resample audio to MMS sample rate (16kHz)."""
        data, sr = sf.read(audio_path)
        if len(data.shape) > 1:
            data = data.mean(axis=1)
        waveform = torch.tensor(data, dtype=torch.float32).unsqueeze(0)
        if sr != self.sample_rate:
            waveform = torchaudio.functional.resample(
                waveform, sr, self.sample_rate
            )
        return waveform

    def _text_to_tokens(self, text: str) -> List[int]:
        """
        Convert romanized text to MMS token indices.
        Space → * (word boundary). Unknown chars skipped.
        Punctuation stripped.
        """
        tokens = []
        for ch in text.lower().strip():
            if ch == " ":
                # Word boundary - don't duplicate
                if tokens and tokens[-1] != self.STAR_IDX:
                    tokens.append(self.STAR_IDX)
            elif ch in self.label2idx and self.label2idx[ch] != self.BLANK_IDX:
                tokens.append(self.label2idx[ch])
            # Punctuation and unknown chars silently skipped
        # Strip leading/trailing boundaries
        while tokens and tokens[0] == self.STAR_IDX:
            tokens.pop(0)
        while tokens and tokens[-1] == self.STAR_IDX:
            tokens.pop()
        return tokens

    def _compute_word_scores(
        self, tokens: List[int], alignment: List[int], scores: List[float]
    ) -> List[float]:
        """
        Compute per-word alignment scores from frame-level scores.

        Words are delimited by * (STAR_IDX) tokens. For each word, average
        the log-prob scores of its frames, then convert to 0-1 confidence.
        """
        # Map token indices to word indices
        word_idx = 0
        token_to_word = {}
        for i, tok in enumerate(tokens):
            if tok == self.STAR_IDX:
                word_idx += 1
            else:
                token_to_word[i] = word_idx

        n_words = word_idx + 1
        word_log_probs = [[] for _ in range(n_words)]

        # Collect scores per word from alignment
        token_cursor = 0
        for frame_idx, (aligned_tok, score) in enumerate(zip(alignment, scores)):
            if aligned_tok != self.BLANK_IDX and token_cursor < len(tokens):
                if token_cursor in token_to_word:
                    wi = token_to_word[token_cursor]
                    word_log_probs[wi].append(score)
                token_cursor += 1

        # Average log-prob per word, convert to 0-1
        result = []
        for lps in word_log_probs:
            if lps:
                mean_lp = sum(lps) / len(lps)
                conf = max(0.0, min(1.0, 1.0 + mean_lp / 5.0))
                result.append(conf)
            else:
                result.append(0.0)

        return result

    def align(self, audio_path: str, romanized_text: str) -> MMSAlignmentResult:
        """
        Perform forced alignment of romanized text against audio.

        Args:
            audio_path: Path to audio file
            romanized_text: Latin-script text to align

        Returns:
            MMSAlignmentResult with alignment score
        """
        if not self._setup_complete:
            if not self.setup():
                return MMSAlignmentResult(
                    audio_path=audio_path,
                    transcription=romanized_text,
                    alignment_score=0.0,
                    mean_log_prob=-10.0,
                )

        start_time = time.time()

        try:
            waveform = self._load_audio(audio_path)
            audio_dur = waveform.shape[1] / self.sample_rate
            tokens = self._text_to_tokens(romanized_text)

            if not tokens:
                return MMSAlignmentResult(
                    audio_path=audio_path,
                    transcription=romanized_text,
                    alignment_score=0.0,
                    mean_log_prob=-10.0,
                    audio_duration_sec=audio_dur,
                    processing_time_sec=time.time() - start_time,
                )

            with torch.no_grad():
                emission, _ = self.model(waveform.to(self.device))
            emission = emission.cpu()

            token_tensor = torch.tensor([tokens], dtype=torch.int32)

            alignment, scores = torchaudio.functional.forced_align(
                emission, token_tensor, blank=self.BLANK_IDX
            )

            mean_lp = scores.mean().item()
            alignment_score = max(0.0, min(1.0, 1.0 + mean_lp / 5.0))

            # v4: Extract per-word scores for stricter validation
            # Split scores by word boundaries (* tokens)
            word_scores = self._compute_word_scores(
                tokens, alignment[0].tolist(), scores[0].tolist()
            )

            return MMSAlignmentResult(
                audio_path=audio_path,
                transcription=romanized_text,
                alignment_score=alignment_score,
                mean_log_prob=mean_lp,
                word_scores=word_scores,
                num_tokens=len(tokens),
                num_frames=emission.shape[1],
                audio_duration_sec=audio_dur,
                processing_time_sec=time.time() - start_time,
            )

        except Exception as e:
            import traceback
            traceback.print_exc()
            return MMSAlignmentResult(
                audio_path=audio_path,
                transcription=romanized_text,
                alignment_score=0.0,
                mean_log_prob=-10.0,
                processing_time_sec=time.time() - start_time,
            )

    def cleanup(self):
        """Release resources."""
        if self.model is not None:
            del self.model
            self.model = None
        self._setup_complete = False
        if torch.cuda.is_available():
            torch.cuda.empty_cache()