ml/model/engine.py

"""
ml/model/engine.py - Complete Multi-Modal Speech Diagnostic Engine
===================================================================
Orchestrates:
1. Signal conditioning & raw-signal Multi-Factor VAD
2. Acoustic & Voice Phonation Analysis (Praat PointProcess)
3. Neural ASR & DP Phonetic Alignment (wav2vec2-CTC)
4. Neural Disfluency Detection (wav2vec2 + LoRA)
5. Multi-Modal Decision Fusion, Concern Bands & Confidence Margins
"""
from __future__ import annotations
import gc
import json
import os
import time
from pathlib import Path
from typing import Any, Dict, Generator, List, Optional

import numpy as np
import torch
from peft import PeftModel
from transformers import (
    Wav2Vec2FeatureExtractor,
    Wav2Vec2ForCTC,
    Wav2Vec2ForSequenceClassification,
    Wav2Vec2Processor,
)

from ml.model import fusion, pron_eval

SR = 16000
MAX_SECONDS = 6.0
CKPT_PATH = "ml/models/stutter/stutter_lora"
CLASS_MAP_PATH = "ml/models/stutter/class_map.json"
MODEL_BASE = "facebook/wav2vec2-base"
CTC_MODEL_NAME = "facebook/wav2vec2-base-960h"


class SpeechDiagnosticEngine:
    _instance: Optional[SpeechDiagnosticEngine] = None

    def __init__(self, ckpt_dir: str = CKPT_PATH, device: Optional[str] = None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        print(f"[SpeechDiagnosticEngine] Initializing on device: {self.device}")

        # 1. Load LoRA Stutter Classifier
        self.stutter_model = None
        self.stutter_feat = None
        self.id2label = {0: "fluent", 1: "stutter"}
        
        ckpt = Path(ckpt_dir)
        if ckpt.exists():
            try:
                cm_file = ckpt.parent / "class_map.json"
                if cm_file.exists():
                    with open(cm_file, "r", encoding="utf-8") as f:
                        data = json.load(f)
                        self.id2label = {int(k): v for k, v in data.get("id2label", {}).items()}
                
                base = Wav2Vec2ForSequenceClassification.from_pretrained(
                    MODEL_BASE, num_labels=len(self.id2label), ignore_mismatched_sizes=True
                )
                self.stutter_model = PeftModel.from_pretrained(base, str(ckpt))
                self.stutter_model.to(self.device)
                self.stutter_model.eval()
                self.stutter_feat = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_BASE)
                print(f"[SpeechDiagnosticEngine] Loaded LoRA stutter classifier from {ckpt}")
            except Exception as e:
                print(f"[SpeechDiagnosticEngine] Warning: Could not load LoRA classifier: {e}")

        # 2. Load Neural ASR / CTC Pronunciation Model
        print(f"[SpeechDiagnosticEngine] Loading ASR model: {CTC_MODEL_NAME}...")
        self.asr_processor = Wav2Vec2Processor.from_pretrained(CTC_MODEL_NAME)
        self.asr_model = Wav2Vec2ForCTC.from_pretrained(CTC_MODEL_NAME)
        self.asr_model.to(self.device)
        self.asr_model.eval()
        
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        print("[SpeechDiagnosticEngine] Engine initialized and ready.")

    @classmethod
    def get_instance(cls, ckpt_dir: str = CKPT_PATH) -> SpeechDiagnosticEngine:
        """Singleton accessor."""
        if cls._instance is None:
            cls._instance = SpeechDiagnosticEngine(ckpt_dir)
        return cls._instance

    @torch.inference_mode()
    def transcribe_and_align(self, audio_input: Any, reference: str) -> dict:
        """Perform neural ASR decoding and DP phonetic word alignment."""
        raw_arr = pron_eval._load_wave(audio_input)
        if pron_eval.is_silent_or_empty(raw_arr, SR):
            ref_norm = pron_eval._norm(reference)
            return {
                "asr_hypothesis": "",
                "reference_normalized": ref_norm,
                "word_error": len(ref_norm.split()),
                "n_reference_words": len(ref_norm.split()),
                "wer": 1.0,
                "goodness": 0.0,
                "pron_score": 0.0,
                "alignment": [{"expected": w, "spoken": "—", "status": "omission"} for w in ref_norm.split()],
                "is_silent": True,
                "length_warning": None,
            }

        # Normalize audio strictly for neural model forward pass
        norm_arr = pron_eval.normalize_for_neural_inference(raw_arr)
        inp = self.asr_processor(norm_arr, sampling_rate=SR, return_tensors="pt")
        inp = {k: v.to(self.device) for k, v in inp.items()}
        logits = self.asr_model(**inp).logits
        pred_ids = torch.argmax(logits, dim=-1)
        hypothesis = pron_eval._norm(self.asr_processor.batch_decode(pred_ids)[0])

        ref = pron_eval._norm(reference)
        alignment = pron_eval.align_words(ref, hypothesis)
        wer, pron_score, counts = pron_eval.compute_standard_wer_and_pron_score(ref, hypothesis, alignment)

        length_warning = None
        if len(hypothesis.split()) == 1 and len(ref.split()) >= 4:
            length_warning = f"Only 1 word ('{hypothesis}') was recognized out of {len(ref.split())} target words."

        return {
            "asr_hypothesis": hypothesis,
            "reference_normalized": ref,
            "word_error": counts["substitutions"] + counts["deletions"] + counts["insertions"],
            "n_reference_words": counts["n_ref"],
            "wer": wer,
            "goodness": pron_score,
            "pron_score": pron_score,
            "alignment": alignment,
            "is_silent": False,
            "length_warning": length_warning,
        }

    @torch.inference_mode()
    def predict_stutter_probs(self, audio_input: Any) -> Optional[List[float]]:
        """Predict neural stutter probabilities [P(fluent), P(stutter)]."""
        if self.stutter_model is None:
            return None

        raw_arr = pron_eval._load_wave(audio_input)
        if pron_eval.is_silent_or_empty(raw_arr, SR):
            return None

        norm_arr = pron_eval.normalize_for_neural_inference(raw_arr)
        arr_clipped = norm_arr[: int(SR * MAX_SECONDS)]
        inp = self.stutter_feat(arr_clipped, sampling_rate=SR, return_tensors="pt", padding=True)
        inp = {k: v.to(self.device) for k, v in inp.items()}
        logits = self.stutter_model(**inp).logits
        return torch.softmax(logits, dim=-1)[0].tolist()

    def diagnose_audio(
        self,
        audio_input: Any,
        target_phrase: str,
        normal_calibration_audio: Optional[Any] = None,
    ) -> Dict[str, Any]:
        """Complete, unified diagnostic pipeline with latency timing."""
        t0 = time.perf_counter()

        # 1. Load raw audio and check VAD
        raw_arr = pron_eval._load_wave(audio_input)
        is_silent = pron_eval.is_silent_or_empty(raw_arr, SR)

        # 2. Voice Acoustics & Phonation (Praat on raw waveform)
        artic = pron_eval.praat_metrics_arr(raw_arr, SR)

        # 3. Neural ASR & DP Alignment
        pron = self.transcribe_and_align(raw_arr, target_phrase)

        # 4. Neural Disfluency Detection
        stut_probs = self.predict_stutter_probs(raw_arr) if not is_silent else None
        p_stut = float(stut_probs[1]) if (stut_probs and len(stut_probs) > 1) else (
            float(np.sum(stut_probs[1:])) if stut_probs else 0.0
        )

        # 5. Acoustic-Phonetic Flaw Rules
        flaws = pron_eval.analyze_speech_flaws(
            reference=target_phrase,
            hypothesis=pron.get("asr_hypothesis", ""),
            alignment=pron.get("alignment", []),
            praat_dict=artic,
            stutter_prob=p_stut,
            is_silent=is_silent,
        )

        # 6. Baseline Normalization ("My Normal")
        cal = None
        if normal_calibration_audio is not None and self.stutter_model is not None:
            norm_probs = self.predict_stutter_probs(normal_calibration_audio)
            if norm_probs is not None:
                cal = fusion.calibrate_from_normal(norm_probs)

        # 7. Multi-Modal Decision Fusion
        decision = fusion.diag_statistics(stut_probs, pron, artic, cal)
        latency_ms = round((time.perf_counter() - t0) * 1000, 1)

        return {
            "is_silent": is_silent,
            "decision": decision,
            "pronunciation": pron,
            "flaws": flaws,
            "articulation": artic,
            "stutter_probs": stut_probs,
            "latency_ms": latency_ms,
            "duration_s": round(len(raw_arr) / SR, 2),
            "device": self.device,
        }

    def diagnose_audio_stream(
        self,
        audio_input: Any,
        target_phrase: str,
        normal_calibration_audio: Optional[Any] = None,
    ) -> Generator[Dict[str, Any], None, Dict[str, Any]]:
        """Streaming generator yielding step-by-step progress and telemetry."""
        t0 = time.perf_counter()

        # Step 1: Conditioning & Raw-Signal VAD
        t_step = time.perf_counter()
        raw_arr = pron_eval._load_wave(audio_input)
        is_silent = pron_eval.is_silent_or_empty(raw_arr, SR)
        t_s1 = round((time.perf_counter() - t_step) * 1000, 1)
        yield {
            "step": 1,
            "total": 5,
            "label": "Acoustic Signal Preconditioning",
            "detail": f"16kHz PCM Resampling, 60Hz High-Pass, Multi-Factor VAD ({t_s1} ms)",
            "progress": 0.20,
            "elapsed_ms": round((time.perf_counter() - t0) * 1000, 1),
        }

        # Step 2: Phonation & Voice Analysis (Praat)
        t_step = time.perf_counter()
        artic = pron_eval.praat_metrics_arr(raw_arr, SR)
        t_s2 = round((time.perf_counter() - t_step) * 1000, 1)
        f0_val = f"{artic.get('f0_median_hz', 0):.1f}Hz" if artic.get('f0_median_hz') is not None else "N/A"
        hnr_val = f"{artic.get('hnr_db', 0):.1f}dB" if artic.get('hnr_db') is not None else "N/A"
        yield {
            "step": 2,
            "total": 5,
            "label": "Acoustic & Voice Phonation Analysis",
            "detail": f"Praat PointProcess Pitch F0={f0_val}, HNR={hnr_val} ({t_s2} ms)",
            "progress": 0.40,
            "elapsed_ms": round((time.perf_counter() - t0) * 1000, 1),
        }

        # Step 3: Neural ASR & DP Phonetic Alignment
        t_step = time.perf_counter()
        pron = self.transcribe_and_align(raw_arr, target_phrase)
        t_s3 = round((time.perf_counter() - t_step) * 1000, 1)
        yield {
            "step": 3,
            "total": 5,
            "label": "Neural ASR & Phonetic Alignment",
            "detail": f"wav2vec2-CTC: \"{pron.get('asr_hypothesis','')}\" | WER: {pron.get('wer',0)*100:.1f}% ({t_s3} ms)",
            "progress": 0.60,
            "elapsed_ms": round((time.perf_counter() - t0) * 1000, 1),
        }

        # Step 4: Neural Disfluency Detection
        t_step = time.perf_counter()
        stut_probs = self.predict_stutter_probs(raw_arr) if not is_silent else None
        p_stut = float(stut_probs[1]) if (stut_probs and len(stut_probs) > 1) else (
            float(np.sum(stut_probs[1:])) if stut_probs else 0.0
        )
        t_s4 = round((time.perf_counter() - t_step) * 1000, 1)
        yield {
            "step": 4,
            "total": 5,
            "label": "Neural Disfluency Classification (LoRA)",
            "detail": f"Wav2Vec2 LoRA Stutter Probability: {p_stut*100:.1f}% ({t_s4} ms)",
            "progress": 0.80,
            "elapsed_ms": round((time.perf_counter() - t0) * 1000, 1),
        }

        # Step 5: Sound Flaws & Decision Fusion
        t_step = time.perf_counter()
        flaws = pron_eval.analyze_speech_flaws(
            reference=target_phrase,
            hypothesis=pron.get("asr_hypothesis", ""),
            alignment=pron.get("alignment", []),
            praat_dict=artic,
            stutter_prob=p_stut,
            is_silent=is_silent,
        )

        cal = None
        if normal_calibration_audio is not None and self.stutter_model is not None:
            norm_probs = self.predict_stutter_probs(normal_calibration_audio)
            if norm_probs is not None:
                cal = fusion.calibrate_from_normal(norm_probs)

        decision = fusion.diag_statistics(stut_probs, pron, artic, cal)
        t_s5 = round((time.perf_counter() - t_step) * 1000, 1)
        latency_ms = round((time.perf_counter() - t0) * 1000, 1)

        final_res = {
            "is_silent": is_silent,
            "decision": decision,
            "pronunciation": pron,
            "flaws": flaws,
            "articulation": artic,
            "stutter_probs": stut_probs,
            "latency_ms": latency_ms,
            "duration_s": round(len(raw_arr) / SR, 2),
            "device": self.device,
            "step_timings_ms": {
                "preconditioning": t_s1,
                "phonation_praat": t_s2,
                "neural_asr": t_s3,
                "neural_disfluency": t_s4,
                "fusion_and_flaws": t_s5,
            }
        }

        yield {
            "step": 5,
            "total": 5,
            "label": "Multi-Modal Decision Fusion & Clinical Report",
            "detail": f"Screening Index: {decision.get('fluency_100', 0)}/100 | Concern Band: {decision['buckets']['overall'].upper()} ({t_s5} ms)",
            "progress": 1.0,
            "elapsed_ms": latency_ms,
            "final_result": final_res,
        }
        return final_res