#!/usr/bin/env python
"""Machine-fault-detector: audio file in -> diagnosis (or transcript) -> spoken reply WAV.

  pip install -r requirements.txt
  python inference.py clip.wav                       # -> clip_reply.wav + printed text
  python inference.py clip.wav --out reply.wav --no-voice

Loads openai/whisper-large-v3 + this LoRA adapter, decodes like training (16 kHz mono,
<= 30 s chunks, peak -1 dBFS, greedy, suppress_tokens=[] because ':' and ';' are template
tokens). Per chunk:
  * a complete machine template -> "Sound check. This is a fan under high load. ..."
  * otherwise speech            -> "I heard: <transcript>"
  * noise / known hallucination -> silent
The reply is spoken with facebook/mms-tts-eng (CC-BY-NC-4.0). Research prototype: trained
on synthetic machine sounds only; do not use for maintenance or safety decisions.
"""

import argparse
import copy
import re
import sys
from math import gcd
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
from scipy.signal import resample_poly

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
from diagnosis import describe, parse_diagnosis  # noqa: E402
from tts_text import speakable, split_sentences  # noqa: E402

SR = 16000
REPO = "muruganmath/Machine-fault-detector"
# outputs that are noise, not speech (measured on the training project's clips)
NOT_SPEECH = re.compile(r"^\s*(?:(?:(?:machine|engine)(?:\s+sounds?)?)[\s.,;:]*)+$", re.I)
HALLUCINATIONS = {"you", "so", "the", "okay", "ok", "bye", "uh", "um", "hmm", "oh", "thank you",
                  "thanks", "thank you very much", "thanks for watching", "thank you for watching"}


def load_audio(path: str) -> np.ndarray:
    x, sr = sf.read(path, dtype="float32", always_2d=True)
    x = x.mean(axis=1)
    if sr != SR:
        g = gcd(sr, SR)
        x = resample_poly(x, SR // g, sr // g).astype(np.float32)
    return x


def chunks(x: np.ndarray, max_s: float = 30.0):
    n = int(max_s * SR)
    for i in range(0, len(x), n):
        c = x[i:i + n]
        if c.size >= SR and 20 * np.log10(max(float(np.sqrt(np.mean(c ** 2))), 1e-10)) >= -60.0:
            yield i / SR, c                       # same rejection as training: < 1 s or silent


class Detector:
    def __init__(self, repo: str = REPO, base: str = "openai/whisper-large-v3", device: str | None = None):
        from peft import PeftModel
        from transformers import WhisperForConditionalGeneration, WhisperProcessor

        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
        self.proc = WhisperProcessor.from_pretrained(base)
        model = WhisperForConditionalGeneration.from_pretrained(base, dtype=self.dtype, attn_implementation="sdpa")
        self.model = PeftModel.from_pretrained(model, repo).merge_and_unload().to(self.device).eval()
        gc = copy.deepcopy(self.model.generation_config)
        gc.suppress_tokens, gc.language, gc.task, gc.forced_decoder_ids = [], "en", "transcribe", None
        self.gen_config = gc

    @torch.inference_mode()
    def decode(self, chunk: np.ndarray) -> str:
        x = chunk * (10 ** (-1.0 / 20) / max(float(np.abs(chunk).max()), 1e-9))        # peak -1 dBFS
        feats = self.proc.feature_extractor([x], sampling_rate=SR, return_tensors="pt").input_features
        out = self.model.generate(input_features=feats.to(self.device, self.dtype), generation_config=self.gen_config,
                                  max_new_tokens=128, num_beams=1)
        return self.proc.tokenizer.batch_decode(out, skip_special_tokens=True)[0].strip()


def reply_for(text: str) -> tuple[str, str | None]:
    """(kind, reply text or None)."""
    d = parse_diagnosis(text)
    if d is not None:
        return "machine", describe(d)
    key = re.sub(r"\s+", " ", re.sub(r"[^a-z0-9' ]", " ", text.lower())).strip()
    if not re.search(r"[A-Za-z0-9]", text) or NOT_SPEECH.match(text) or key in HALLUCINATIONS:
        return "noise", None
    return "speech", f"I heard: {text}"


class Voice:
    def __init__(self, repo: str = "facebook/mms-tts-eng", device: str | None = None):
        from transformers import AutoTokenizer, VitsModel
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.model = VitsModel.from_pretrained(repo).to(self.device).eval()
        self.tok = AutoTokenizer.from_pretrained(repo)
        self.sr = int(self.model.config.sampling_rate)

    @torch.inference_mode()
    def speak(self, text: str) -> np.ndarray:
        parts, gap = [], np.zeros(int(0.12 * self.sr), np.float32)
        for s in split_sentences(speakable(text)):
            ids = self.tok(s, return_tensors="pt").to(self.device)
            if ids["input_ids"].shape[-1] == 0:
                continue
            torch.manual_seed(0)                   # VITS is stochastic; keep replies reproducible
            parts += [self.model(**ids).waveform[0].float().cpu().numpy(), gap]
        if not parts:
            return np.zeros(0, np.float32)
        y = np.concatenate(parts[:-1])
        return (y * min(1.0, 0.9 / max(float(np.abs(y).max()), 1e-9))).astype(np.float32)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("audio")
    ap.add_argument("--out", help="reply WAV (default: <audio>_reply.wav next to the input)")
    ap.add_argument("--repo", default=REPO, help="adapter repo id or local folder")
    ap.add_argument("--base", default="openai/whisper-large-v3", help="base model repo id or local folder")
    ap.add_argument("--tts", default="facebook/mms-tts-eng", help="TTS repo id or local folder")
    ap.add_argument("--no-voice", action="store_true", help="print the text only")
    args = ap.parse_args()

    det = Detector(args.repo, args.base)
    replies = []
    for t0, c in chunks(load_audio(args.audio)):
        text = det.decode(c)
        kind, reply = reply_for(text)
        print(f"[{t0:6.1f}s] {kind:7s} model: {text}")
        print(f"{'':16}reply: {reply or '(silent)'}")
        if reply:
            replies.append(reply)
    if args.no_voice or not replies:
        return
    voice = Voice(args.tts)
    gap = np.zeros(int(0.4 * voice.sr), np.float32)
    wav = np.concatenate([np.concatenate([voice.speak(r), gap]) for r in replies])[: -gap.size]
    out = args.out or str(Path(args.audio).with_name(Path(args.audio).stem + "_reply.wav"))
    sf.write(out, wav, voice.sr, subtype="PCM_16")
    print(f"voice reply: {out} ({wav.size / voice.sr:.1f} s)")


if __name__ == "__main__":
    main()