Spaces:

Tom-Brenner
/

PhonePaint

Sleeping

PhonePaint Space: MAPS-only, ZeroGPU, assets fetched at startup (part 3)
a22c748 verified about 2 months ago
51.6 kB
#!/usr/bin/env python3
"""mfa_align_from_json.py – MFA alignment from pre-provided text or CSV
Skips missing audio and missing .TextGrid outputs; processes intersection only.
CSV rule: row j maps to wav "{j:04d}.wav"; text taken from column 'transcript'.
"""

from __future__ import annotations
import argparse, csv, json, os, shutil, subprocess as sp, tempfile, re
from collections import defaultdict
from pathlib import Path
from typing import List, Tuple, Optional, Iterable
from textgrid import TextGrid, IntervalTier  # pip install textgrid

_MFA_PRETRAINED_BASE = Path.home() / "Documents" / "MFA" / "pretrained_models"
_MFA_DICT_DIR = _MFA_PRETRAINED_BASE / "dictionary"
_MFA_ACOUSTIC_DIR = _MFA_PRETRAINED_BASE / "acoustic"
_MFA_G2P_DIR = _MFA_PRETRAINED_BASE / "g2p"
_DEFAULT_MFA_CONDA_ENV = "mfa_env"
_PHONEPAINT_MFA_ENVS = frozenset({"PhonePaintMFA"})


def _active_conda_env_name() -> Optional[str]:
    name = (os.environ.get("CONDA_DEFAULT_ENV") or "").strip()
    if name:
        return name
    prefix = (os.environ.get("CONDA_PREFIX") or "").strip()
    if prefix:
        return Path(prefix).name
    return None


def _is_phonepaint_mfa_env() -> bool:
    return _active_conda_env_name() in _PHONEPAINT_MFA_ENVS


def resolve_mfa_bin() -> str:
    """Locate the MFA CLI binary; never import montreal_forced_aligner.
    Resolution order:
      1. MFA_BIN
      2. Active env: ``$CONDA_PREFIX/bin/mfa`` (same-env; preferred)
      3. ``mfa`` on PATH when it lives under ``$CONDA_PREFIX``
      4. Fallback conda env ``MFA_CONDA_ENV`` / ``mfa_env`` (legacy sidecar) —
         skipped while PhonePaintMFA is active (MFA+MAPS share that env)
      5. Any ``mfa`` on PATH (also skipped while PhonePaintMFA is active)
    """
    explicit = os.environ.get("MFA_BIN")
    if explicit:
        path = Path(explicit).expanduser()
        if path.is_file():
            return str(path)
        raise FileNotFoundError(f"MFA_BIN is set but not a file: {explicit}")

    conda_prefix = os.environ.get("CONDA_PREFIX")
    candidates: List[Path] = []
    if conda_prefix:
        candidates.append(Path(conda_prefix) / "bin" / "mfa")
        which = shutil.which("mfa")
        if which:
            which_path = Path(which).resolve()
            prefix_path = Path(conda_prefix).resolve()
            try:
                which_path.relative_to(prefix_path)
                candidates.append(which_path)
            except ValueError:
                pass

    allow_sidecar = not _is_phonepaint_mfa_env()
    if allow_sidecar:
        env_name = os.environ.get("MFA_CONDA_ENV", _DEFAULT_MFA_CONDA_ENV)
        if conda_prefix:
            candidates.append(Path(conda_prefix).parent / env_name / "bin" / "mfa")
        home = Path.home()
        for base in (home / "miniconda3", home / "anaconda3", home / "mambaforge", home / "miniforge3"):
            candidates.append(base / "envs" / env_name / "bin" / "mfa")
        which = shutil.which("mfa")
        if which:
            candidates.append(Path(which))

    seen: set[str] = set()
    for cand in candidates:
        key = str(cand)
        if key in seen:
            continue
        seen.add(key)
        if cand.is_file():
            return str(cand)
    if _is_phonepaint_mfa_env():
        env = _active_conda_env_name() or "PhonePaint"
        raise FileNotFoundError(
            f"{env} is active but MFA CLI was not found in this env "
            "($CONDA_PREFIX/bin/mfa). Install MFA in-env or set MFA_BIN; "
            f"the mfa_env sidecar is not used while {env} is active."
        )
    env_name = os.environ.get("MFA_CONDA_ENV", _DEFAULT_MFA_CONDA_ENV)
    raise FileNotFoundError(
        "MFA CLI not found. Install MFA in the active env (e.g. PhonePaint), "
        "create mfa_env (./create_mfa_envs.sh), or set MFA_BIN "
        f"(looked for active env and fallback {env_name!r})."
    )


def mfa_cli_env(**extra: str) -> dict:
    """Env for subprocesses that run the MFA CLI.
    Calling ``mfa`` by absolute path is not enough: MFA's third-party check and
    Kaldi helpers need ``sox`` / Kaldi binaries from the MFA conda env on PATH
    (and their shared libs).
    """
    env = os.environ.copy()
    mfa_bin = Path(resolve_mfa_bin())
    env_bin = mfa_bin.parent
    env_root = env_bin.parent
    env["PATH"] = f"{env_bin}{os.pathsep}{env.get('PATH', '')}"
    lib_dir = env_root / "lib"
    if lib_dir.is_dir():
        prev_ld = env.get("LD_LIBRARY_PATH", "")
        env["LD_LIBRARY_PATH"] = (
            f"{lib_dir}{os.pathsep}{prev_ld}" if prev_ld else str(lib_dir)
        )
    # MFA looks for conda-installed third-party tools relative to this prefix.
    env["CONDA_PREFIX"] = str(env_root)
    env["MFA_BIN"] = str(mfa_bin)
    if extra:
        env.update(extra)
    return env


def _resolve_model_file(kind: str, model: str) -> Path:
    p = Path(os.path.expanduser(os.path.expandvars(model)))
    if p.exists():
        return p
    if kind == "dictionary":
        cand = _MFA_DICT_DIR / f"{model}.dict"
    elif kind == "acoustic":
        cand = _MFA_ACOUSTIC_DIR / f"{model}.zip"
    elif kind == "g2p":
        cand = _MFA_G2P_DIR / f"{model}.zip"
    else:
        raise ValueError(kind)
    if not cand.exists():
        raise FileNotFoundError(f"Missing MFA {kind} model file: {cand}")
    return cand


def _load_dictionary_vocab(dict_path: Path) -> set[str]:
    vocab: set[str] = set()
    for line in dict_path.read_text(encoding="utf-8", errors="ignore").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        w = line.split()[0]
        w = re.sub(r"\(\d+\)$", "", w)
        vocab.add(w.lower())
    return vocab


# ARPABET to IPA conversion mapping for MFA compatibility
ARPABET_TO_IPA = {
    # Vowels
    'AA': 'ɑ', 'AA0': 'ɑ', 'AA1': 'ɑ', 'AA2': 'ɑ',
    'AE': 'æ', 'AE0': 'æ', 'AE1': 'æ', 'AE2': 'æ',
    'AH': 'ʌ', 'AH0': 'ə', 'AH1': 'ʌ', 'AH2': 'ʌ',  # AH0 -> schwa
    'AO': 'ɔ', 'AO0': 'ɔ', 'AO1': 'ɔ', 'AO2': 'ɔ',
    'AW': 'aʊ', 'AW0': 'aʊ', 'AW1': 'aʊ', 'AW2': 'aʊ',
    'AY': 'aɪ', 'AY0': 'aɪ', 'AY1': 'aɪ', 'AY2': 'aɪ',
    'EH': 'ɛ', 'EH0': 'ɛ', 'EH1': 'ɛ', 'EH2': 'ɛ',
    'ER': 'ɜr', 'ER0': 'ər', 'ER1': 'ɜr', 'ER2': 'ɜr',  # ER0 -> schwar
    'EY': 'eɪ', 'EY0': 'eɪ', 'EY1': 'eɪ', 'EY2': 'eɪ',
    'IH': 'ɪ', 'IH0': 'ɪ', 'IH1': 'ɪ', 'IH2': 'ɪ',
    'IY': 'i', 'IY0': 'i', 'IY1': 'i', 'IY2': 'i',
    'OW': 'oʊ', 'OW0': 'oʊ', 'OW1': 'oʊ', 'OW2': 'oʊ',
    'OY': 'ɔɪ', 'OY0': 'ɔɪ', 'OY1': 'ɔɪ', 'OY2': 'ɔɪ',
    'UH': 'ʊ', 'UH0': 'ʊ', 'UH1': 'ʊ', 'UH2': 'ʊ',
    'UW': 'u', 'UW0': 'u', 'UW1': 'u', 'UW2': 'u',
    
    # Consonants
    'B': 'b', 'CH': 'tʃ', 'D': 'd', 'DH': 'ð', 'F': 'f', 'G': 'ɡ',
    'HH': 'h', 'JH': 'dʒ', 'K': 'k', 'L': 'l', 'M': 'm', 'N': 'n',
    'NG': 'ŋ', 'P': 'p', 'R': 'r', 'S': 's', 'SH': 'ʃ', 'T': 't',
    'TH': 'θ', 'V': 'v', 'W': 'w', 'Y': 'j', 'Z': 'z', 'ZH': 'ʒ',
}


def convert_arpabet_to_ipa(phone: str) -> str:
    """Convert ARPABET phone to IPA phone used by english_mfa.
    
    Args:
        phone: ARPABET phone (e.g., 'AY1', 'B', 'SH')
        
    Returns:
        IPA phone (e.g., 'aɪ', 'b', 'ʃ')
    """
    phone = phone.strip().upper()
    
    # Direct lookup
    if phone in ARPABET_TO_IPA:
        return ARPABET_TO_IPA[phone]
    
    # Try without stress marker for vowels
    base_phone = re.sub(r'[012]$', '', phone)
    if base_phone in ARPABET_TO_IPA:
        return ARPABET_TO_IPA[base_phone]
    
    # If no conversion found, return original (might already be IPA)
    return phone.lower()


def _ffmpeg_convert(src: Path, dst: Path, gain_db: float | None = None) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    cmd = [
        "ffmpeg",
        "-y",
        "-loglevel",
        "error",
        "-i",
        str(src),
        "-map",
        "0:a:0",
        "-ac",
        "1",
        "-ar",
        "16000",
        "-c:a",
        "pcm_s16le",
    ]
    # Optional volume gain for the WAV that MFA aligns against.
    # Applied only when explicitly requested to preserve default behavior.
    if gain_db is not None:
        cmd += ["-af", f"volume={gain_db}dB"]
    cmd.append(str(dst))
    sp.run(cmd, check=True)


def run_mfa(
    corpus: Path,
    acoustic: str,
    dictionary: str,
    threads: int,
    *,
    out_dir: Optional[Path] = None,
    tag: Optional[str] = None,
) -> Path:
    """Run MFA alignment into a *dedicated* output directory.
    Important: the old implementation always wrote to corpus/out and thus overwrote
    previous passes. Incremental/fallback logic requires distinct output dirs.
    """
    if out_dir is None:
        out_dir = corpus / "out"
    out_dir = Path(out_dir)

    # Keep per-pass history/log to avoid clobbering across multiple runs.
    tag = tag or out_dir.name
    hist = corpus / f"history_{tag}"
    hist.mkdir(exist_ok=True)
    log = corpus / f"mfa_{tag}.log"

    mfa_cmd = resolve_mfa_bin()
    dict_arg = str(_resolve_model_file("dictionary", dictionary))
    acoustic_arg = str(_resolve_model_file("acoustic", acoustic))
    n_jobs = max(1, int(threads))
    cmd = [
        mfa_cmd,
        "align",
        str(corpus),
        dict_arg,
        acoustic_arg,
        str(out_dir),
        "--clean",
        "--single_speaker",
        "--num_jobs",
        str(n_jobs),
        "--output_format",
        "long_textgrid",
        "--debug",
        "--verbose",
        "--overwrite",
    ]
    # --use_mp with num_jobs=1 still pays multiprocessing overhead for no gain.
    if n_jobs > 1:
        cmd.append("--use_mp")
    env = mfa_cli_env(MFA_HISTORY_DIR=str(hist))

    try:
        result = sp.run(cmd, capture_output=True, text=True, check=False, env=env)
        log.write_text((result.stdout or "") + "\n" + (result.stderr or ""), encoding="utf-8", errors="ignore")
        if result.returncode != 0:
            print(f"MFA command failed with return code {result.returncode}")
            print(f"Command: {' '.join(cmd)}")
            print(f"Log file: {log}")
            if log.exists():
                print("Last 10 lines of log:")
                with log.open("r") as rf:
                    _lines = rf.readlines()
                    for line in _lines[-10:]:
                        print(f"  {line.rstrip()}")
            raise RuntimeError(f"MFA failed (see {log})")
    except FileNotFoundError as exc:
        raise RuntimeError(
            f"MFA not found at {mfa_cmd}. Install MFA in the active env, set "
            "MFA_BIN, or create mfa_env (./create_mfa_envs.sh)."
        ) from exc
    return out_dir


def count_spn_instances(output_dir: Path) -> int:
    """Count total 'spn' phone instances in all TextGrid files."""
    spn_count = 0
    for tg_file in output_dir.glob("*.TextGrid"):
        try:
            words, phones = parse_textgrid(tg_file)
            spn_count += sum(1 for _, _, phone in phones if phone.strip().lower() == 'spn')
        except Exception as e:
            print(f"Warning: Could not parse {tg_file}: {e}")
    return spn_count


def extract_spn_segments(output_dir: Path) -> dict:
    """Extract information about spn segments for targeted re-alignment."""
    spn_segments = {}
    for tg_file in output_dir.glob("*.TextGrid"):
        try:
            words, phones = parse_textgrid(tg_file)
            file_spns = []
            for start, end, phone in phones:
                if phone.strip().lower() == 'spn':
                    file_spns.append((start, end, phone))
            if file_spns:
                spn_segments[tg_file.stem] = file_spns
        except Exception as e:
            print(f"Warning: Could not parse {tg_file}: {e}")
    return spn_segments


def create_spn_targeted_corpus(original_corpus: Path, spn_segments: dict, text_data: dict) -> Path:
    """Create a new corpus containing only audio/text segments that had spn instances."""
    spn_corpus = original_corpus.parent / f"{original_corpus.name}_spn_targeted"
    spn_corpus.mkdir(exist_ok=True)
    
    # Copy files that had spn instances
    for filename in spn_segments.keys():
        wav_file = f"{filename}.wav"
        lab_file = f"{filename}.lab"
        
        src_wav = original_corpus / wav_file
        src_lab = original_corpus / lab_file
        
        if src_wav.exists() and src_lab.exists():
            shutil.copy2(src_wav, spn_corpus / wav_file)
            shutil.copy2(src_lab, spn_corpus / lab_file)
    
    return spn_corpus


def find_time_overlapping_segments(segments_list: list, target_start: float, target_end: float, tolerance: float = 0.01) -> list:
    """Find segments that overlap with a target time range."""
    overlapping = []
    for i, (start, end, text) in enumerate(segments_list):
        # Check for any overlap with tolerance
        if not (end <= target_start - tolerance or start >= target_end + tolerance):
            overlapping.append((i, start, end, text))
    return overlapping

def _build_aligned_dict(out_dir: Path, included: list[str]) -> dict[str, dict]:
    """Create the aligned JSON dict from TextGrids in out_dir for the given files."""
    words_all, phones_all = defaultdict(list), defaultdict(list)
    for fn in included:
        tg_path = out_dir / Path(fn).with_suffix(".TextGrid")
        if not tg_path.exists():
            continue
        w, p = parse_textgrid(tg_path)
        words_all[fn], phones_all[fn] = w, p

    aligned: dict[str, dict] = {}
    for fn in included:
        if fn not in words_all and fn not in phones_all:
            continue
        aligned[fn] = {
            "words": {
                str(i): {"xmin": round(x0, 3), "xmax": round(x1, 3), "text": t}
                for i, (x0, x1, t) in enumerate(words_all.get(fn, []))
            },
            "phones": {
                str(i): {"xmin": round(x0, 3), "xmax": round(x1, 3), "text": t}
                for i, (x0, x1, t) in enumerate(phones_all.get(fn, []))
            },
        }
    return aligned


def _alignment_needs_fallback(out_dir: Path, included: list[str]) -> bool:
    """True when alignment still has OOV/spn/sil gaps or a missing TextGrid."""
    targets = {"<unk>", "unk", "</unk>", "spn", "sil"}
    for fn in included:
        tg_path = out_dir / Path(fn).with_suffix(".TextGrid")
        if not tg_path.exists():
            return True
        try:
            words, phones = parse_textgrid(tg_path)
        except Exception:
            return True
        for _, _, w in words:
            if (w or "").strip().lower() in targets:
                return True
        for _, _, p in phones:
            if (p or "").strip().lower() in targets:
                return True
    return False


def _needs_g2p(out_dir: Path, included: list[str]) -> bool:
    """Return True if any included file still has unk/spn/sil tokens in words or phones."""
    return _alignment_needs_fallback(out_dir, included)

def run_incremental_alignment_with_model_pairs(
    corpus: Path,
    model_pairs: list[tuple[str, str]],
    threads: int,
    text_data: dict[str, dict],
    use_g2p: bool = False,
    g2p_model: str = "english_us_mfa",
    *,
    pre_g2p_json_path: Optional[Path] = None,
    included_fns: Optional[list[str]] = None,
) -> tuple[Path, dict]:
    """
    IMPROVED Pipeline (fixes 'homie' -> '<unk>' issue):
      1) Align with primary pair → primary out_dir.
      2) For each subsequent pair (e.g., ARPA): align only when primary still has
         <unk>/spn/sil gaps; surgically replace spn in primary using that pass.
      3) Optional G2P pass LAST: build enhanced dict from input texts; align; surgically replace remaining spn.
      4) Return the modified primary out_dir and stats.
    Replacement semantics:
      - Replace ONLY primary 'spn' phones using fallback phones clipped to the primary word span.
      - Words tier: replace '<unk>' with fallback word(s) if available; keep recognized primary words as-is.
      - No borrowing from neighbors; original primary word boundaries preserved.
      - G2P is used LAST, only for remaining OOV issues after dictionary fallbacks.
    """
    included = included_fns or list(text_data.keys())

    # 1) Primary alignment
    primary_acoustic, primary_dict = model_pairs[0]
    primary_out = run_mfa(corpus, primary_acoustic, primary_dict, threads, out_dir=(corpus / "out_primary"), tag="primary")

    # Stats init
    stats = {
        "models_tried": [f"{primary_acoustic}+{primary_dict}"],
        "spn_counts": [],
        "improvements": [],  # delta vs previous
    }
    spn0 = count_spn_instances(primary_out)
    stats["spn_counts"].append(spn0)

    # 2) Subsequent model pairs as additional fallbacks (e.g., english_us_arpa)
    for k, (acoustic, dictionary) in enumerate(model_pairs[1:], start=1):
        if not _alignment_needs_fallback(primary_out, included):
            print(
                f"  Skipping fallback pair {acoustic}+{dictionary}: "
                "primary alignment complete (no <unk>/spn/sil)."
            )
            break
        print(f"\n  Running fallback pair {acoustic}+{dictionary} …")
        fb_out = run_mfa(
            corpus,
            acoustic,
            dictionary,
            threads,
            out_dir = (corpus / f"out_fallback_{k:02d}"),
            tag = f"fallback_{k:02d}")
        is_arpa_model = 'arpa' in acoustic.lower() or 'arpa' in dictionary.lower()
        surgical_spn_replacement(primary_out, fb_out, convert_arpabet=is_arpa_model)
        spn_now = count_spn_instances(primary_out)
        stats["models_tried"].append(f"{acoustic}+{dictionary}")
        stats["improvements"].append(stats["spn_counts"][-1] - spn_now)
        stats["spn_counts"].append(spn_now)
        if not _alignment_needs_fallback(primary_out, included):
            print("  Alignment complete after fallback; skipping remaining dictionary pairs.")
            break

    # 3) Optional G2P pass LAST (only after dictionary fallbacks have been tried)
    # Pre-G2P JSON snapshot, if requested
    if use_g2p and pre_g2p_json_path is not None:
        if not included:
            raise RuntimeError("included_fns must be provided when pre_g2p_json_path is set.")
        pre_aligned = _build_aligned_dict(primary_out, included)
        pre_g2p_json_path.parent.mkdir(parents=True, exist_ok=True)
        pre_g2p_json_path.write_text(json.dumps(pre_aligned, indent=2, ensure_ascii=False))
        print(f"[pre-G2P] JSON snapshot written to: {pre_g2p_json_path}")

    if use_g2p:
        # Only run G2P if targeted tokens remain
        if not _needs_g2p(primary_out, included):
            print("  Skipping G2P fallback: no <unk>/unk/</unk>/spn/sil tokens detected.")
            return primary_out, stats

        print(f"\n  Applying G2P fallback (last resort for remaining OOV words)...")
        tmp_g2p_dir = corpus / "_g2p"
        tmp_g2p_dir.mkdir(exist_ok=True)
        all_words = extract_words_from_texts(text_data)  # naive tokenization
        if generate_g2p_pronunciations(all_words, g2p_model, tmp_g2p_dir):
            enhanced_dict = create_enhanced_dictionary(primary_dict, tmp_g2p_dir / "g2p_pronunciations.dict", tmp_g2p_dir)
            if enhanced_dict and enhanced_dict.exists():
                g2p_out = run_mfa(corpus, primary_acoustic, str(enhanced_dict), threads, out_dir=(corpus / "out_g2p"),tag="g2p")
                surgical_spn_replacement(primary_out, g2p_out, convert_arpabet=False)
                spn_g2p = count_spn_instances(primary_out)
                stats["models_tried"].append(f"{primary_acoustic}+G2P({g2p_model})")
                stats["improvements"].append(stats["spn_counts"][-1] - spn_g2p)
                stats["spn_counts"].append(spn_g2p)
            else:
                print("  G2P fallback skipped: could not create enhanced dictionary")
        else:
            print("  G2P fallback skipped: could not generate pronunciations")

    # 4) Return the modified primary out_dir and stats
    return primary_out, stats



def _stretch_segments_to_primary_boundaries(primary_start: float, primary_end: float, fallback_segments: List[Tuple[float, float, str]]) -> List[Tuple[float, float, str]]:
    """Stretch fallback segments (phones or words) to fit primary boundaries using proportional scaling.
    
    NEVER use fallback timestamps directly. Always scale relative to primary boundaries.
    
    Example: 'homie' primary (0.85-1.12), fallback phones [(0.85, 0.94, 'h'), (0.94, 0.97, 'oʊ'), (0.97, 1.04, 'm'), (1.04, 1.10, 'i')]
    
    Algorithm:
    1. Convert fallback to relative times (0-based)
    2. Calculate scaling factor to fit primary duration 
    3. Scale each segment proportionally
    4. Add primary_start to get final timestamps
    """
    if not fallback_segments:
        return []
    
    # Get fallback boundaries
    fallback_start = min(s for s, e, lab in fallback_segments)
    fallback_end = max(e for s, e, lab in fallback_segments)
    fallback_duration = fallback_end - fallback_start
    
    if fallback_duration <= 1e-9:
        # Degenerate case: all segments at same time, distribute equally
        primary_duration = primary_end - primary_start
        segment_duration = primary_duration / len(fallback_segments)
        result = []
        for i, (_, _, lab) in enumerate(fallback_segments):
            start = primary_start + i * segment_duration
            end = primary_start + (i + 1) * segment_duration
            result.append((start, end, lab))
        return result
    
    # Calculate scaling factor
    primary_duration = primary_end - primary_start
    scale_factor = primary_duration / fallback_duration
    
    # Scale each segment
    result = []
    for fb_start, fb_end, lab in fallback_segments:
        # Convert to relative times (0-based from fallback_start)
        rel_start = fb_start - fallback_start
        rel_end = fb_end - fallback_start
        
        # Scale to primary duration
        scaled_start = rel_start * scale_factor
        scaled_end = rel_end * scale_factor
        
        # Add primary_start to get final absolute times
        final_start = primary_start + scaled_start
        final_end = primary_start + scaled_end
        
        result.append((final_start, final_end, lab))
    
    # Ensure last phone ends exactly at primary_end (handle floating point precision)
    if result:
        last_start, _, last_lab = result[-1]
        result[-1] = (last_start, primary_end, last_lab)
    
    return result


def word_bounds_for(primary_words: Optional[IntervalTier], span_start: float, span_end: float) -> Tuple[float, float]:
    """Find the primary word boundaries that contain the given span.
    
    CRITICAL: Always return primary model boundaries, never fallback boundaries.
    This ensures timestamps like 'homie' (0.85-1.12 primary vs 0.85-1.10 fallback) 
    use the primary boundaries (0.85-1.12).
    """
    if primary_words is None:
        return (span_start, span_end)
    
    # Find the primary word that overlaps with this span
    for w in primary_words:
        # Check if primary word overlaps with the span (not strict containment)
        if (w.maxTime > span_start + 1e-9 and w.minTime < span_end - 1e-9):
            if os.environ.get("MFA_DEBUG_BOUNDARIES"):
                word_label = (w.mark or "").strip()
                print(f"    word_bounds_for({span_start:.3f}-{span_end:.3f}) -> primary word '{word_label}' ({w.minTime:.3f}-{w.maxTime:.3f})")
            return (w.minTime, w.maxTime)
    
    # If no overlapping word found, return the span as-is
    if os.environ.get("MFA_DEBUG_BOUNDARIES"):
        print(f"    word_bounds_for({span_start:.3f}-{span_end:.3f}) -> no overlapping word found, using span")
    return (span_start, span_end)


def _iter_textgrid_paths(root: Path) -> Iterable[Path]:
    patterns = ("*.TextGrid", "*.Textgrid", "*.textgrid")
    for pat in patterns:
        yield from root.rglob(pat)



def surgical_spn_replacement(primary_path: Path, fallback_path: Path, convert_arpabet: bool = True) -> None:
    """
    Recalculate per pseudocode. Copy primary word boundaries; only replace primary 'spn' with fallback phones
    remapped to primary span by equal per-phone adjustment; never copy fallback timestamps.
    """
        # Handle directory vs file inputs up front
    if Path(primary_path).is_dir() and Path(fallback_path).is_dir():
        fb_index = {q.stem: q for q in Path(fallback_path).rglob("*.TextGrid")}
        for prim in Path(primary_path).rglob("*.TextGrid"):
            mate = fb_index.get(prim.stem)
            if mate is None:
                continue
            surgical_spn_replacement(prim, mate, convert_arpabet)
        return
    if Path(primary_path).is_file() and Path(fallback_path).is_file():
        pass
    else:
        raise ValueError("surgical_spn_replacement expects both paths to be files or both to be directories")
    tol = 1e-3

    def pick_tiers(tg: TextGrid):
        phones_tier = None
        words_tier = None
        for t in tg.tiers:
            if not isinstance(t, IntervalTier):
                continue
            lname = (t.name or "").lower()
            if ("phone" in lname) or ("segment" in lname):
                phones_tier = t
            if "word" in lname:
                words_tier = t
        return phones_tier, words_tier

    def find_span_indices(tier: IntervalTier, start: float, end: float):
        s = None; e = None
        for i, it in enumerate(tier):
            if it.maxTime < start - tol:
                continue
            if it.minTime > end + tol:
                break
            if s is None and abs(it.minTime - start) <= 5e-3:
                s = i
            if abs(it.maxTime - end) <= 5e-3:
                e = i
        if s is None:
            for i, it in enumerate(tier):
                if it.minTime >= start - tol:
                    s = i; break
        if e is None:
            for i in range(len(tier)-1, -1, -1):
                it = tier[i]
                if it.maxTime <= end + tol:
                    e = i; break
        return s, e

    def get_word_by_index(tier: IntervalTier, k: int):
        return tier[k] if (tier is not None and 0 <= k < len(tier)) else None

    primary_tg = TextGrid(); primary_tg.read(primary_path)
    fallback_tg = TextGrid(); fallback_tg.read(fallback_path)

    p_phones, p_words = pick_tiers(primary_tg)
    f_phones, f_words = pick_tiers(fallback_tg)

    if p_words is None or p_phones is None:
        raise ValueError("Primary TextGrid missing phones/words")

    new_tiers = []
    rewritten_phones = IntervalTier(name=p_phones.name or "phones", minTime=p_phones.minTime, maxTime=p_phones.maxTime)
    rewritten_words = IntervalTier(name=p_words.name or "words", minTime=p_words.minTime, maxTime=p_words.maxTime)

    for wi, pw in enumerate(p_words):
        p_xmin, p_xmax = pw.minTime, pw.maxTime
        p_text = (pw.mark or "").strip()

        pi_s, pi_e = find_span_indices(p_phones, p_xmin, p_xmax)

        fw = get_word_by_index(f_words, wi)
        f_xmin = fw.minTime if fw is not None else p_xmin
        f_xmax = fw.maxTime if fw is not None else p_xmax
        f_text = (fw.mark or "").strip() if fw is not None else ""

        fi_s = fi_e = None
        if f_phones is not None and fw is not None:
            fi_s, fi_e = find_span_indices(f_phones, f_xmin, f_xmax)

        if pi_s is None or pi_e is None or pi_s > pi_e:
            rewritten_words.add(p_xmin, p_xmax, p_text)
            continue

        p_first = (p_phones[pi_s].mark or "").strip().lower()
        f_first = ""
        if fi_s is not None and fi_e is not None and fi_s <= fi_e and f_phones is not None:
            f_first = (f_phones[fi_s].mark or "").strip().lower()

        use_fallback = (p_first == "spn" and f_first != "spn" and (fi_s is not None and fi_e is not None))

        if not use_fallback:
            for idx in range(pi_s, pi_e + 1):
                ph = p_phones[idx]
                rewritten_phones.add(ph.minTime, ph.maxTime, ph.mark)
            rewritten_words.add(p_xmin, p_xmax, p_text)
        else:
            n = fi_e - fi_s + 1
            if n <= 0:
                for idx in range(pi_s, pi_e + 1):
                    ph = p_phones[idx]
                    rewritten_phones.add(ph.minTime, ph.maxTime, ph.mark)
                rewritten_words.add(p_xmin, p_xmax, p_text)
            else:
                fb = []
                for idx in range(fi_s, fi_e + 1):
                    it = f_phones[idx]
                    lab = (it.mark or "").strip()
                    if convert_arpabet:
                        lab = convert_arpabet_to_ipa(lab)
                    fb.append((it.minTime, it.maxTime, lab))
                f_dur = f_xmax - f_xmin
                p_dur = p_xmax - p_xmin
                delta = f_dur - p_dur
                cur = p_xmin
                for j, (s, e, lab) in enumerate(fb):
                    d = max(0.0, e - s)
                    d_adj = d - (delta / n)
                    if j == n - 1:
                        start = cur
                        end = p_xmax
                    else:
                        d_adj = max(1e-4, d_adj)
                        start = cur
                        end = cur + d_adj
                    if start < p_xmin: start = p_xmin
                    if end > p_xmax: end = p_xmax
                    rewritten_phones.add(start, end, lab)
                    cur = end
                final_word_text = f_text if f_text else p_text
                rewritten_words.add(p_xmin, p_xmax, final_word_text)

    for t in primary_tg.tiers:
        lname = (t.name or "").lower()
        if "phone" in lname or "segment" in lname:
            continue
        if "word" in lname:
            continue
        new_tiers.append(t)

    merged = TextGrid(minTime=primary_tg.minTime, maxTime=primary_tg.maxTime)
    merged.append(rewritten_phones)
    merged.append(rewritten_words)
    for t in new_tiers:
        merged.append(t)
    merged.write(primary_path)
    return

def parse_textgrid(tg_path: Path):
    tg = TextGrid()
    tg.read(tg_path)
    words, phones = [], []
    for tier in tg:
        name = (tier.name or "").lower()
        tgt = words if "word" in name else phones if ("phone" in name or "segment" in name) else None
        if tgt is None:
            continue
        tgt += [(i.minTime, i.maxTime, i.mark) for i in tier if (i.mark or "").strip()]
    return words, phones

def detect_fail_words_from_textgrid(tg_path: Path, eps_s: float = 0.010) -> list[int]:
    """
    fail_word definition (your condition 1):
      Word k is a fail_word iff within the word interval [ws,we]:
        - there exists exactly one phone interval P with label in {'spn','sil'}
          such that |P.start-ws|<=eps and |P.end-we|<=eps
        - and there are no other phones overlapping the word by more than eps
    Returns word indices (0-based) into the words tier order.
    """
    tg = TextGrid(); tg.read(tg_path)
    phones_tier, words_tier = None, None
    for t in tg.tiers:
        if not isinstance(t, IntervalTier):
            continue
        lname = (t.name or "").lower()
        if ("phone" in lname) or ("segment" in lname):
            phones_tier = t
        if "word" in lname:
            words_tier = t
    if phones_tier is None or words_tier is None:
        return []

    def overlap(a0, a1, b0, b1) -> float:
        return max(0.0, min(a1, b1) - max(a0, b0))

    out: list[int] = []
    for k, w in enumerate(words_tier):
        ws, we = float(w.minTime), float(w.maxTime)
        # Candidate "single phone equals word interval"
        candidates = []
        for p in phones_tier:
            ps, pe = float(p.minTime), float(p.maxTime)
            lab = (p.mark or "").strip().lower()
            if lab in {"spn", "sil"} and abs(ps - ws) <= eps_s and abs(pe - we) <= eps_s:
                candidates.append(p)
        if len(candidates) != 1:
            continue
        # Ensure no other phone overlaps the word substantially
        ok = True
        for p in phones_tier:
            ps, pe = float(p.minTime), float(p.maxTime)
            if overlap(ws, we, ps, pe) > eps_s:
                if p is not candidates[0]:
                    ok = False
                    break
        if ok:
            out.append(k)
    return out

def patch_fail_words_from_fallback(
    primary_out_dir: Path,
    fallback_out_dir: Path,
    fail_words_by_wav: dict[str, list[int]],
    *,
    eps_s: float = 0.010,
    convert_arpabet: bool = False,
) -> None:
    """
    Implements your condition 3:
      - keep primary words tier (timestamps + labels)
      - replace only fail_words (by word index) where primary has a single spn/sil phone == word interval
      - use only fallback phones inside fallback word interval, discard fallback spn/sil/empty
      - retimestamp by preserving fallback phone duration ratios and projecting onto primary word interval
    """
    def pick_tiers(tg: TextGrid):
        phones_tier, words_tier = None, None
        for t in tg.tiers:
            if not isinstance(t, IntervalTier):
                continue
            lname = (t.name or "").lower()
            if ("phone" in lname) or ("segment" in lname):
                phones_tier = t
            if "word" in lname:
                words_tier = t
        return phones_tier, words_tier

    def merge_adjacent(segs: list[tuple[float,float,str]]) -> list[tuple[float,float,str]]:
        out: list[tuple[float,float,str]] = []
        for s,e,lab in segs:
            if e <= s + 1e-9:
                continue
            if out and out[-1][2] == lab and abs(out[-1][1] - s) <= 1e-6:
                out[-1] = (out[-1][0], e, lab)
            else:
                out.append((s,e,lab))
        return out

    def clamp(x, lo, hi): return lo if x < lo else hi if x > hi else x

    fb_index = {q.stem: q for q in Path(fallback_out_dir).rglob("*.TextGrid")}
    for prim in Path(primary_out_dir).rglob("*.TextGrid"):
        mate = fb_index.get(prim.stem)
        if mate is None:
            continue
        wav_name = prim.with_suffix(".wav").name  # matches included_wavs naming convention
        fail_idxs = fail_words_by_wav.get(wav_name, [])
        if not fail_idxs:
            continue

        ptg = TextGrid(); ptg.read(prim)
        ftg = TextGrid(); ftg.read(mate)
        p_phones, p_words = pick_tiers(ptg)
        f_phones, f_words = pick_tiers(ftg)
        if p_phones is None or p_words is None or f_phones is None or f_words is None:
            continue

        # Build phone replacement plan: list of (word_start, word_end, replacement_segments)
        plan: list[tuple[float,float,list[tuple[float,float,str]]]] = []
        for k in fail_idxs:
            if not (0 <= k < len(p_words)) or not (0 <= k < len(f_words)):
                continue

            pw = p_words[k]; fw = f_words[k]
            pws, pwe = float(pw.minTime), float(pw.maxTime)
            fws, fwe = float(fw.minTime), float(fw.maxTime)

            # Verify primary really has the single spn/sil phone == word interval (within eps)
            carrier = None
            for p in p_phones:
                lab = (p.mark or "").strip().lower()
                if lab in {"spn","sil"} and abs(float(p.minTime)-pws) <= eps_s and abs(float(p.maxTime)-pwe) <= eps_s:
                    carrier = p
                    break
            if carrier is None:
                continue

            # Collect fallback phones within fallback word interval; discard spn/sil/empty
            fb = []
            for p in f_phones:
                ps, pe = float(p.minTime), float(p.maxTime)
                if pe <= fws or ps >= fwe:
                    continue
                lab = (p.mark or "").strip()
                if convert_arpabet:
                    lab = convert_arpabet_to_ipa(lab)
                lab_l = lab.strip().lower()
                if not lab_l or lab_l in {"spn","sil"}:
                    continue
                cs = clamp(ps, fws, fwe)
                ce = clamp(pe, fws, fwe)
                if ce > cs + 1e-9:
                    fb.append((cs, ce, lab))
            if not fb:
                continue

            durs = [ce - cs for cs, ce, _ in fb]
            total = sum(durs)
            if total <= 1e-9:
                continue

            # Project ratios onto primary word interval
            span = pwe - pws
            cur = pws
            repl: list[tuple[float,float,str]] = []
            for i, (_, _, lab) in enumerate(fb):
                if i == len(fb) - 1:
                    ns, ne = cur, pwe
                else:
                    frac = durs[i] / total
                    ne = cur + frac * span
                    ns = cur
                repl.append((ns, ne, lab))
                cur = ne
            repl = merge_adjacent(repl)
            plan.append((pws, pwe, repl))

        if not plan:
            continue
        plan.sort(key=lambda x: x[0])

        # Rebuild phones tier with replacements; keep all other tiers untouched
        new_phones = IntervalTier(name=p_phones.name or "phones", minTime=p_phones.minTime, maxTime=p_phones.maxTime)
        i_plan = 0
        for p in p_phones:
            ps, pe = float(p.minTime), float(p.maxTime)
            lab = (p.mark or "")
            if i_plan < len(plan):
                ws, we, repl = plan[i_plan]
                if abs(ps - ws) <= eps_s and abs(pe - we) <= eps_s and (lab.strip().lower() in {"spn","sil"}):
                    for ns, ne, nl in repl:
                        new_phones.add(ns, ne, nl)
                    i_plan += 1
                    continue
            new_phones.add(ps, pe, lab)

        # Write merged TextGrid: replace phones tier only
        merged = TextGrid(minTime=ptg.minTime, maxTime=ptg.maxTime)
        for t in ptg.tiers:
            lname = (t.name or "").lower()
            if ("phone" in lname) or ("segment" in lname):
                merged.append(new_phones)
            else:
                merged.append(t)
        merged.write(prim)

def extract_words_from_texts(text_data: dict) -> set:
    """Extract unique words from all text entries."""
    all_words = set()
    for entry in text_data.values():
        text = entry.get('text', '')
        if isinstance(text, str):
            # Basic word extraction - split on whitespace and punctuation
            words = re.findall(r"\b[a-zA-Z]+\b", text.lower())
            all_words.update(words)
    return all_words


def detect_oov_words(words: set, dictionary: str) -> set:
    dict_path = _resolve_model_file("dictionary", dictionary)
    vocab = _load_dictionary_vocab(dict_path)
    return {w for w in words if w.lower() not in vocab}


def generate_g2p_pronunciations(oov_words: set, g2p_model: str, output_path: Path) -> bool:
    """Generate pronunciations for OOV words using G2P model."""
    if not oov_words:
        print("  [g2p] OOV count: 0 → skipping G2P generation")
        return False
        
    try:
        print(f"  [g2p] OOV count: {len(oov_words)}")
        # Create temporary file with OOV words
        oov_file = output_path / "oov_words.txt"
        with open(oov_file, 'w') as f:
            for word in sorted(oov_words):
                f.write(f"{word}\n")
        
        # Generate pronunciations using MFA G2P (sidecar CLI)
        g2p_output = output_path / "g2p_pronunciations.dict"
        mfa_cmd = resolve_mfa_bin()
        g2p_model_path = str(_resolve_model_file("g2p", g2p_model))
        result = sp.run([
            mfa_cmd, "g2p", str(oov_file), g2p_model_path, str(g2p_output)
        ], capture_output=True, text=True, check=True, env=mfa_cli_env())
        if g2p_output.exists():
            lines = g2p_output.read_text(encoding="utf-8", errors="ignore").strip().splitlines()
            print(f"  [g2p] output lines: {len(lines)}")
            return len(lines) > 0
        else:
            print("  [g2p] output file missing after g2p")
            return False
        
    except sp.CalledProcessError as e:
        print(f"Warning: G2P generation failed: {e}")
        return False

def create_enhanced_dictionary(base_dict: str, g2p_dict_path: Path, output_path: Path) -> Optional[Path]:
    """Create an enhanced dictionary: base dictionary + G2P pronunciations.
    This is *not* MFA's built-in "enhanced dictionary" logic; it's a simple merge to
    ensure only OOV words gain pronunciations while preserving the original dictionary.
    """
    output_path.mkdir(parents=True, exist_ok=True)
    enhanced_dict = output_path / f"{base_dict.replace('/', '_')}_enhanced.dict"

    # Read G2P pronunciations (may be empty if G2P failed)
    g2p_content = ""
    if g2p_dict_path.exists():
        g2p_content = g2p_dict_path.read_text(encoding="utf-8", errors="ignore").strip()

    base_path = _resolve_model_file("dictionary", base_dict)
    base_dict_content = base_path.read_text(encoding="utf-8", errors="ignore").strip()

    # If neither base nor G2P content exists, skip
    if not (base_dict_content or g2p_content):
        print("Warning: Could not create enhanced dictionary, skipping G2P pass")
        return None

    # Write and then verify *after closing* (avoids zero-size stat on some filesystems)
    with enhanced_dict.open("w", encoding="utf-8") as out_f:
        if base_dict_content:
            out_f.write(base_dict_content)
            out_f.write("\n")
        if g2p_content:
            out_f.write(g2p_content)
            out_f.write("\n")
            print("Added G2P pronunciations to enhanced dictionary")

    if enhanced_dict.exists() and enhanced_dict.stat().st_size > 0:
        return enhanced_dict

    print("Warning: Could not create enhanced dictionary, skipping G2P pass")
    return None


def _load_texts_from_json(p: Path) -> dict[str, dict]:
    data = json.loads(p.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise SystemExit("--input_json must map filename → {'text': str}")
    return data


def _load_texts_from_csv(p: Path,prefix:str) -> dict[str, dict]:
    # Row position j (0-based) → filename "{j:04d}.wav"; text from 'transcript'
    out: dict[str, dict] = {}
    with p.open(newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        if "transcript" not in (reader.fieldnames or []):
            raise SystemExit("CSV must contain a 'transcript' column")
        for j, row in enumerate(reader):
            txt = (row.get("transcript") or "").strip()
            if not txt:
                continue  # skip missing/empty rows
            fn = f"{prefix}_{j:04d}.wav"
            out[fn] = {"text": txt}
    return out


def parse_model_pairs(args) -> list[tuple[str, str]]:
    """Parse model pairs from arguments, with smart defaults."""
    model_pairs = []
    
    # If explicit model pairs provided, use those
    if args.mfa_model_pairs and any(args.mfa_model_pairs):
        for pair_list in args.mfa_model_pairs:
            for pair_str in pair_list:
                if ':' in pair_str:
                    acoustic, dictionary = pair_str.split(':', 1)
                    model_pairs.append((acoustic.strip(), dictionary.strip()))
                else:
                    # If no colon, assume it's a model name to use for both
                    model_pairs.append((pair_str.strip(), pair_str.strip()))
    else:
        # Use smart defaults based on backward compatibility
        # Create matched pairs from the fallback lists
        if hasattr(args, 'mfa_acoustic_fallbacks') and hasattr(args, 'mfa_dict_fallbacks'):
            # Primary pair: first of each list
            primary_acoustic = args.mfa_acoustic_fallbacks[0]
            primary_dict = args.mfa_dict_fallbacks[0]
            model_pairs.append((primary_acoustic, primary_dict))
            
            # ARPA pair only if explicitly listed in both fallback lists
            if 'english_us_arpa' in args.mfa_acoustic_fallbacks and 'english_us_arpa' in args.mfa_dict_fallbacks:
                model_pairs.append(('english_us_arpa', 'english_us_arpa'))
            
        # Fallback to primary english_mfa only (G2P handles OOVs when requested)
        if not model_pairs:
            model_pairs = [
                ('english_mfa', 'english_mfa'),
            ]
    
    return model_pairs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--audio_dir", type=Path, required=True)
    mx = ap.add_mutually_exclusive_group(required=True)
    mx.add_argument("--input_json", type=Path, help="filename → {'text': str}")
    mx.add_argument("--input_csv", type=Path, help="CSV with a 'transcript' column")
    ap.add_argument("--output_json", type=Path, required=True)
    ap.add_argument("--prefix", type=str, help="audio filename prefix when csv is used")
    ap.add_argument(
        "--gain",
        type=float,
        default=None,
        help="Optional volume gain (dB) applied when converting audio to 16k mono WAV for MFA alignment.",
    )
    ap.add_argument("-j", "--threads", type=int, default=os.cpu_count() or 4)
    ap.add_argument("--mfa_acoustic", default="english_mfa", 
                    help="Primary acoustic model (deprecated, use --mfa_model_pairs)")
    ap.add_argument("--mfa_dict", default="english_mfa",
                    help="Primary dictionary (deprecated, use --mfa_model_pairs)")
    ap.add_argument("--mfa_model_pairs", nargs='*', action='append',
                    help="Model pairs in format 'acoustic:dictionary' (e.g. english_mfa:english_mfa english_us_arpa:english_us_arpa)")
    ap.add_argument("--mfa_dict_fallbacks", nargs='+', 
                    default=["english_mfa"], 
                    help="Dictionary models to try in order (deprecated, use --mfa_model_pairs)")
    ap.add_argument("--mfa_acoustic_fallbacks", nargs='+',
                    default=["english_mfa"], 
                    help="Acoustic models to try in order (deprecated, use --mfa_model_pairs)")  
    ap.add_argument("--use_g2p_fallback", action='store_true',
                    help="After dictionary pass(es), run G2P only for remaining OOV/fail words")
    ap.add_argument("--g2p_model", default="english_us_mfa",
                    help="G2P model for OOV pronunciation generation (default: english_us_mfa)")
    ap.add_argument("--debug_keep_tmp", action="store_true")
    ap.add_argument("--fallback_tolerance_ms", type=float, default=50.0,
                    help = "Tolerance window around primary spn intervals for capturing fallback segments (ms)")
    ap.add_argument("--pre_g2p_output_json", type=Path,
                    help = "When --use_g2p_fallback is set, write an additional JSON just before G2P fallback to this path")
    args = ap.parse_args()
    
    # Parse model pairs
    model_pairs = parse_model_pairs(args)
    print(f"Using model pairs: {' → '.join([f'{a}+{d}' for a, d in model_pairs])}")
    if args.use_g2p_fallback and not args.pre_g2p_output_json:
        raise SystemExit("--pre_g2p_output_json is required when --use_g2p_fallback is used")


    if args.input_json:
        data = _load_texts_from_json(args.input_json)
    else:
        data = _load_texts_from_csv(args.input_csv,args.prefix)  # type: ignore[arg-type]

    if args.debug_keep_tmp:
        tmp = Path(tempfile.mkdtemp(prefix="mfa_corpus_"))
        cleanup = False
    else:
        td = tempfile.TemporaryDirectory()
        tmp = Path(td.name)
        cleanup = True

    included: list[str] = []
    for fn, val in data.items():
        src = args.audio_dir / fn
        if not src.exists():
            continue  # skip missing audio
        text = (val or {}).get("text", "")
        if not isinstance(text, str) or not text.strip():
            continue  # skip malformed/empty entry
        dst_wav = tmp / fn
        _ffmpeg_convert(src, dst_wav, gain_db=args.gain)
        (tmp / Path(fn).with_suffix(".lab")).write_text(text.strip(), encoding="utf-8")
        included.append(fn)

    # Use incremental fallback system with model pairs for alignment
    if included:
        try:
            out_dir, alignment_stats = run_incremental_alignment_with_model_pairs(
                tmp, 
                model_pairs,
                args.threads,
                text_data=data,
                use_g2p=args.use_g2p_fallback,
                g2p_model=args.g2p_model,
                pre_g2p_json_path=args.pre_g2p_output_json,
                included_fns = included
            )
            
            # Store alignment stats for summary
            globals()['alignment_stats'] = alignment_stats
            
        except RuntimeError as e:
            print(f"Error: {e}")
            print("Alignment failed completely.")
            out_dir = tmp / "out"  # Create empty output dir
            out_dir.mkdir(exist_ok=True)
            globals()['alignment_stats'] = {'models_tried': [], 'spn_counts': [], 'improvements': []}
    else:
        out_dir = tmp / "out"
        globals()['alignment_stats'] = {'models_tried': [], 'spn_counts': [], 'improvements': []}

    # Expose tolerance to surgical replacer
    os.environ["MFA_FALLBACK_TOL"] = f"{max(0.0, (args.fallback_tolerance_ms / 1000.0)):.3f}"
    aligned = _build_aligned_dict(out_dir, included)

    args.output_json.write_text(json.dumps(aligned, indent=2, ensure_ascii=False))
    
    # Print detailed summary with alignment statistics
    # Print detailed summary with alignment statistics
    print(f"\n=== Alignment Completed Successfully! ===")
    print(f"Output written to: {args.output_json}")
    print(f"Total files processed: {len(aligned)}")
    
    # Count final spn instances in output
    total_spn_instances = 0
    for file_data in aligned.values():
        phones = file_data.get('phones', {})
        total_spn_instances += sum(1 for phone_data in phones.values() 
                                 if phone_data.get('text', '').strip().lower() == 'spn')
    
    print(f"Final 'spn' instances in output: {total_spn_instances}")
    
    # Print alignment statistics if available
    if 'alignment_stats' in globals() and globals()['alignment_stats']['models_tried']:
        stats = globals()['alignment_stats']
        print(f"\n=== Fallback Performance ===")
        for i, model in enumerate(stats['models_tried']):
            spn_count = stats['spn_counts'][i]
            if i > 0 and i-1 < len(stats['improvements']):
                improvement = stats['improvements'][i-1]
                print(f"{model}: {spn_count} spn instances (improved by {improvement})")
            else:
                print(f"{model}: {spn_count} spn instances")
                
        if len(stats['spn_counts']) > 1:
            total_improvement = stats['spn_counts'][0] - stats['spn_counts'][-1]
            reduction_pct = (total_improvement / stats['spn_counts'][0] * 100) if stats['spn_counts'][0] > 0 else 0
            print(f"Total improvement: -{total_improvement} spn instances ({reduction_pct:.1f}% reduction)")
    
    print(f"\nConfiguration used:")
    print(f"  Model pairs: {' → '.join([f'{a}+{d}' for a, d in model_pairs])}")
    if args.use_g2p_fallback:
        print(f"  G2P model: {args.g2p_model}")
    else:
        print(f"  G2P fallback: Disabled")

    if cleanup:
        shutil.rmtree(tmp, ignore_errors=True)
    elif not args.debug_keep_tmp:
        print(f"[debug] corpus kept at {tmp}")


if __name__ == "__main__":
    main()