#!/usr/bin/env python3
"""CAFE-25lang — inference with the automatic language gate (no language flag).

`inference.py` needs to be told whether the audio is English. This variant decides per utterance
with the small gate head shipped in `language_gate.safetensors`, so the caller supplies nothing:

    python inference_gated.py audio_en.wav audio_de.wav audio_pt.wav

How it works: a first pass runs the bare trunk up to the gate's tap layer and pools those features
into p(non-English); utterances above the threshold are then transcribed with the adapters
attached, the rest on the bare trunk. Two passes, one set of weights. `--threshold` overrides the
packaged threshold.

Requirements: nemo_toolkit[asr], torch, safetensors, soundfile. Audio must be 16 kHz mono.
"""
from __future__ import annotations

import argparse
import json
import os

import soundfile as sf
import torch
import torch.nn as nn
from safetensors import safe_open
from safetensors.torch import load_file

HERE = os.path.dirname(os.path.abspath(__file__))
ADAPTER_LAYERS = (3, 7, 11, 15, 19, 23)
SR = 16000


class Bottleneck(nn.Module):
    def __init__(self, dim: int, rank: int):
        super().__init__()
        self.down = nn.Linear(dim, rank)
        self.act = nn.GELU()
        self.up = nn.Linear(rank, dim)

    def forward(self, x):
        return self.up(self.act(self.down(x)))


class AttentiveStatsPool(nn.Module):
    """Attention-weighted mean and standard deviation over time, with a length mask."""

    def __init__(self, dim: int, attn_dim: int):
        super().__init__()
        self.proj = nn.Sequential(nn.Linear(dim, attn_dim), nn.Tanh(), nn.Linear(attn_dim, dim))

    def forward(self, x, lens):
        t = x.shape[1]
        mask = (torch.arange(t, device=x.device)[None, :] < lens[:, None]).unsqueeze(-1)
        w = self.proj(x).masked_fill(~mask, float("-inf")).softmax(dim=1)
        mean = (w * x).sum(1)
        var = (w * (x - mean.unsqueeze(1)).pow(2)).sum(1)
        return torch.cat([mean, var.clamp_min(1e-8).sqrt()], dim=-1)


class LanguageGate(nn.Module):
    """Pooling -> LayerNorm -> MLP -> one logit; positive means non-English."""

    def __init__(self, dim=640, attn_dim=128, hidden=256):
        super().__init__()
        self.pool = AttentiveStatsPool(dim, attn_dim)
        self.norm = nn.LayerNorm(2 * dim)
        self.mlp = nn.Sequential(nn.Linear(2 * dim, hidden), nn.GELU(),
                                 nn.Dropout(0.1), nn.Linear(hidden, 1))

    def forward(self, x, lens):
        return self.mlp(self.norm(self.pool(x, lens))).squeeze(-1)


def load_gate(path: str, device):
    with safe_open(path, framework="pt") as f:
        meta = f.metadata() or {}
    tap = int(meta.get("tap_encoder_layer", 11))
    attn = int(meta.get("attn_dim", 128))
    hidden = int(meta.get("hidden", 256))
    thr = float(meta.get("threshold", 0.8582))
    g = LanguageGate(640, attn, hidden)
    g.load_state_dict(load_file(path))
    return g.to(device).eval(), tap, thr


def make_adapters(path: str, device, dtype):
    sd = load_file(path)
    rank, dim = sd["adapter.0.down.weight"].shape
    mods = []
    for i in sorted({int(k.split(".")[1]) for k in sd}):
        a = Bottleneck(dim, rank)
        a.load_state_dict({"down.weight": sd[f"adapter.{i}.down.weight"],
                           "down.bias": sd[f"adapter.{i}.down.bias"],
                           "up.weight": sd[f"adapter.{i}.up.weight"],
                           "up.bias": sd[f"adapter.{i}.up.bias"]})
        mods.append(a.to(device=device, dtype=dtype).eval())
    return mods


def attach(model, mods):
    handles = []
    for a, layer in zip(mods, ADAPTER_LAYERS):
        def mk(a):
            def hook(module, inp, out):
                x = out[0] if isinstance(out, tuple) else out
                if x.dim() != 3 or x.shape[-1] != a.down.in_features:
                    return out
                y = x + a(x)
                return (y,) + tuple(out[1:]) if isinstance(out, tuple) else y
            return hook
        handles.append(model.encoder.layers[layer].register_forward_hook(mk(a)))
    return handles


def safe_decoding(model, cuda_graph: bool = False):
    from omegaconf import open_dict

    cfg = model.cfg.decoding
    if cfg.strategy != "beam":
        cfg.strategy = "greedy_batch"
    with open_dict(cfg):
        if hasattr(cfg, "greedy"):
            cfg.greedy.use_cuda_graph_decoder = bool(cuda_graph)
            cfg.greedy.loop_labels = True
        cfg.fused_batch_size = -1
    model.change_decoding_strategy(cfg)


@torch.no_grad()
def gate_probs(model, gate, tap: int, files, device, max_len_s: float = 20.0):
    """First pass: bare-trunk features at the tap layer -> p(non-English) per file."""
    buf = {}

    def hook(module, inp, out):
        x = out[0] if isinstance(out, tuple) else out
        if x.dim() == 3:
            buf["x"] = x
        return out

    h = model.encoder.layers[tap].register_forward_hook(hook)
    probs = []
    try:
        for f in files:
            x, sr = sf.read(f, dtype="float32")
            if x.ndim > 1:
                x = x.mean(axis=1)
            if sr != SR:
                raise SystemExit(f"{f}: expected {SR} Hz mono, got {sr}")
            if max_len_s and len(x) > int(max_len_s * SR):
                x = x[: int(max_len_s * SR)]        # a few seconds are enough to decide
            audio = torch.from_numpy(x)[None].to(device)
            alen = torch.tensor([audio.shape[-1]], device=device)
            _, enc_len = model.forward(input_signal=audio.float(), input_signal_length=alen)
            probs.append(float(torch.sigmoid(gate(buf["x"].float(), enc_len).float())[0]))
            buf.clear()
    finally:
        h.remove()
    return probs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("audio", nargs="+", help="16 kHz mono wav files")
    ap.add_argument("--model", default=os.path.join(HERE, "CAFE-25lang.nemo"))
    ap.add_argument("--adapters", default=os.path.join(HERE, "adapters.safetensors"))
    ap.add_argument("--gate", default=os.path.join(HERE, "language_gate.safetensors"))
    ap.add_argument("--threshold", type=float, default=None, help="overrides the packaged value")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--batch_size", type=int, default=8)
    ap.add_argument("--json", action="store_true", help="emit one JSON object per line")
    ap.add_argument("--cuda-graph", dest="cuda_graph", action="store_true")
    a = ap.parse_args()

    import nemo.collections.asr as nemo_asr

    device = torch.device(a.device)
    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
    model = nemo_asr.models.ASRModel.restore_from(a.model, map_location=device)
    model.to(dtype).eval()
    safe_decoding(model, a.cuda_graph)
    gate, tap, thr = load_gate(a.gate, device)
    if a.threshold is not None:
        thr = a.threshold
    print(f"[CAFE] gate on encoder layer {tap}, threshold {thr:.4f}")

    probs = gate_probs(model, gate, tap, a.audio, device)
    non_en = [f for f, p in zip(a.audio, probs) if p >= thr]
    en = [f for f, p in zip(a.audio, probs) if p < thr]
    print(f"[CAFE] {len(en)} English (bare trunk), {len(non_en)} non-English (adapters)")

    res = {}
    with torch.inference_mode():
        if en:
            out = model.transcribe(en, batch_size=a.batch_size, verbose=False)
            res.update(zip(en, out[0] if isinstance(out, tuple) else out))
        if non_en:
            handles = attach(model, make_adapters(a.adapters, device, dtype))
            try:
                out = model.transcribe(non_en, batch_size=a.batch_size, verbose=False)
                res.update(zip(non_en, out[0] if isinstance(out, tuple) else out))
            finally:
                for h in handles:
                    h.remove()                     # leave the model as restored
    for f, p in zip(a.audio, probs):
        text = getattr(res[f], "text", res[f])
        if a.json:
            print(json.dumps({"audio": f, "p_non_english": round(p, 4),
                              "path": "adapters" if p >= thr else "trunk",
                              "text": text}, ensure_ascii=False))
        else:
            print(f"{f}\tp_non_en={p:.4f}\t{'ML' if p >= thr else 'EN'}\t{text}")


if __name__ == "__main__":
    main()