#!/usr/bin/env python3
"""
infer.py — susurro raw (PyTorch) inference: text + voicepack -> 24 kHz wav.

Self-contained: bundles the StyleTTS2 model code + utility-net assets under
./styletts2/. Pick a voice+register by loading its voicepack, type English text,
get audio. misaki[en] handles G2P; synthesis is the Kokoro-faithful voicepack path
(predict duration/F0/energy from the prosodic style, decode with the acoustic style —
no diffusion sampler).

    pip install -r requirements-raw.txt
    python infer.py --voicepack voicepacks/voice_a__tender.pt \
        --text "Hey, I wasn't expecting you tonight." --out hello.wav

For a dependency-light path (onnxruntime, no PyTorch/transformers) see infer_onnx.py.
"""
from __future__ import annotations

import argparse
import os
import os.path as osp
import sys
import wave
from pathlib import Path

HERE = Path(__file__).resolve().parent
REPO = HERE / "styletts2"
sys.path.insert(0, str(REPO))            # models / text_utils / Modules / Utils

VOICES = ["voice_a", "voice_b", "voice_c"]
REGISTERS = ["neutral", "breathless", "playful", "urgent", "tender", "whisper"]


def _patch_torch_load():
    # StyleTTS2 utility checkpoints need weights_only=False (PyTorch 2.6+ default flipped).
    import torch
    _orig = torch.load
    torch.load = lambda *a, **k: _orig(*a, **{**k, "weights_only": False})


def recursive_munch(d):
    from munch import Munch
    if isinstance(d, dict):
        return Munch((k, recursive_munch(v)) for k, v in d.items())
    if isinstance(d, list):
        return [recursive_munch(v) for v in d]
    return d


def load_model(config_path, ckpt_path, device):
    import yaml
    import torch
    from models import build_model, load_ASR_models, load_F0_models
    from text_utils import TextCleaner
    from Utils.PLBERT.util import load_plbert

    cfg = yaml.safe_load(open(config_path))
    mp = recursive_munch(cfg["model_params"])

    # asset paths in config.yml are relative to ./styletts2/
    asr = load_ASR_models(str(REPO / cfg["ASR_path"]), str(REPO / cfg["ASR_config"]))
    f0 = load_F0_models(str(REPO / cfg["F0_path"]))
    plbert = load_plbert(str(REPO / cfg["PLBERT_dir"]))

    model = build_model(mp, asr, f0, plbert)
    for k in model:
        model[k].eval().to(device)

    state = torch.load(ckpt_path, map_location="cpu")
    params = state
    for cand in ("net", "model", "state_dict"):
        if isinstance(state, dict) and isinstance(state.get(cand), dict):
            params = state[cand]
            break

    def strip_prefix(sd, mkeys):
        # single-GPU checkpoints save keys with a 'module.' prefix the model lacks;
        # strict=False would silently match nothing -> random weights -> noise. Strip it.
        if set(sd) & mkeys:
            return sd
        for p in ("module.", "_orig_mod."):
            if any(k.startswith(p) for k in sd):
                s = {(k[len(p):] if k.startswith(p) else k): v for k, v in sd.items()}
                if set(s) & mkeys:
                    return s
        return sd

    covs = []
    for key in model:
        sd = params.get(key) if isinstance(params, dict) else None
        if not isinstance(sd, dict):
            covs.append(0.0)
            continue
        mkeys = set(model[key].state_dict().keys())
        res = model[key].load_state_dict(strip_prefix(sd, mkeys), strict=False)
        covs.append((len(mkeys) - len(getattr(res, "missing_keys", []))) / max(1, len(mkeys)))
    mean_cov = sum(covs) / max(1, len(covs))
    if mean_cov < 0.5:
        raise SystemExit(f"checkpoint coverage too low ({mean_cov:.0%}) — weights did not "
                         "load; the model would emit noise. Check the checkpoint/config.")
    return model, TextCleaner()


def g2p_en(text: str) -> str:
    from misaki import en
    g2p = en.G2P(trf=False, british=False, fallback=None)
    out = g2p(text)
    ipa = out[0] if isinstance(out, tuple) else out
    return ipa.replace("ʏ", "y")


def synthesize(model, textcleaner, phonemes, voicepack, device):
    import torch
    token_ids = textcleaner(phonemes)
    if not token_ids or len(token_ids) > 510:
        raise ValueError(f"bad token length {len(token_ids)}")
    ref_acoustic = voicepack[:128].unsqueeze(0)
    ref_prosodic = voicepack[128:].unsqueeze(0)
    with torch.no_grad():
        input_ids = torch.LongTensor([[0, *token_ids, 0]]).to(device)
        input_lengths = torch.LongTensor([input_ids.shape[-1]]).to(device)
        text_mask = torch.gt(
            torch.arange(input_lengths.max()).unsqueeze(0).expand(1, -1)
            .type_as(input_lengths) + 1, input_lengths.unsqueeze(1)).to(device)

        bert_dur = model.bert(input_ids, attention_mask=(~text_mask).int())
        d_en = model.bert_encoder(bert_dur).transpose(-1, -2)
        d = model.predictor.text_encoder(d_en, ref_prosodic, input_lengths, text_mask)
        x, _ = model.predictor.lstm(d)
        duration = torch.sigmoid(model.predictor.duration_proj(x)).sum(axis=-1)
        pred_dur = torch.round(duration.squeeze()).clamp(min=1).long()
        if pred_dur.dim() == 0:
            pred_dur = pred_dur.unsqueeze(0)

        n_tokens = input_ids.shape[1]
        total = int(pred_dur.sum().item())
        aln = torch.zeros(n_tokens, total).to(device)
        c = 0
        for i in range(n_tokens):
            di = int(pred_dur[i].item())
            aln[i, c:c + di] = 1
            c += di
        aln = aln.unsqueeze(0)

        en = d.transpose(-1, -2) @ aln
        F0_pred, N_pred = model.predictor.F0Ntrain(en, ref_prosodic)
        t_en = model.text_encoder(input_ids, input_lengths, text_mask)
        asr = t_en @ aln
        audio = model.decoder(asr, F0_pred, N_pred, ref_acoustic)
    return audio.squeeze().cpu().numpy()


def write_wav(path, audio, sr):
    import numpy as np
    a = np.asarray(audio, dtype=np.float32).flatten()
    peak = float(np.max(np.abs(a))) if a.size else 0.0
    if peak > 1.0:
        a = a / peak
    pcm = (a * 32767.0).clip(-32768, 32767).astype("<i2")
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    with wave.open(str(path), "wb") as w:
        w.setnchannels(1); w.setsampwidth(2); w.setframerate(sr)
        w.writeframes(pcm.tobytes())


def main():
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--voicepack", required=True, help=f"voicepacks/<voice>__<register>.pt "
                    f"({'/'.join(VOICES)} x {'/'.join(REGISTERS)})")
    ap.add_argument("--text", required=True)
    ap.add_argument("--out", default="out.wav")
    ap.add_argument("--ckpt", default=str(HERE / "susurro.pth"))
    ap.add_argument("--config", default=str(HERE / "config.yml"))
    ap.add_argument("--sr", type=int, default=24000)
    ap.add_argument("--device", default="")
    args = ap.parse_args()

    _patch_torch_load()
    import torch
    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")

    model, tc = load_model(args.config, args.ckpt, device)
    vp = torch.load(args.voicepack, map_location=device, weights_only=False).squeeze()
    if vp.numel() != 256:
        raise SystemExit(f"voicepack must be 256-d, got {tuple(vp.shape)}")

    audio = synthesize(model, tc, g2p_en(args.text), vp.to(device), device)
    write_wav(args.out, audio, args.sr)
    print(f"[infer] '{args.text[:50]}' -> {args.out} "
          f"({len(audio)/args.sr:.2f}s, {Path(args.voicepack).stem})")


if __name__ == "__main__":
    main()