#!/usr/bin/env python3
"""CAFE-25lang — reference inference (language supplied by the caller).

This model routes by language. English is transcribed by the bare trunk; every other language
additionally needs the language-conditional adapters, which are shipped alongside the checkpoint
in `adapters.safetensors` and injected at inference time.

    python inference.py --lang en    audio_en.wav
    python inference.py --lang other audio_de.wav audio_pt.wav

`--lang en` runs the trunk on its own; `--lang other` injects the adapters. Passing the wrong one
does not fail, it just costs accuracy, so route deliberately. If you do not know the language of
the audio, use `inference_gated.py`, which decides per utterance.

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

import argparse
import os

import torch
import torch.nn as nn
from safetensors.torch import load_file

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


class Bottleneck(nn.Module):
    """Residual adapter: x -> x + up(gelu(down(x)))."""

    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)))


def attach_adapters(model, path: str, device, dtype=torch.bfloat16):
    """Inject the language-conditional adapters as residual forward hooks.

    The .nemo export carries the trunk only, so the adapters have to be re-attached here; without
    them the model is on its English path.
    """
    sd = load_file(path)
    idxs = sorted({int(k.split(".")[1]) for k in sd})
    rank, dim = sd["adapter.0.down.weight"].shape
    if len(idxs) != len(ADAPTER_LAYERS):
        raise SystemExit(f"{len(idxs)} adapters for {len(ADAPTER_LAYERS)} layers")
    handles, mods = [], []
    for i, layer in zip(idxs, ADAPTER_LAYERS):
        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"]})
        a = a.to(device=device, dtype=dtype).eval()
        mods.append(a)

        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                      # shape guard, never corrupt the graph
                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)))
    # keep references alive and verifiable
    model._adapters, model._adapter_handles = mods, handles
    print(f"[CAFE] adapters injected at encoder layers {list(ADAPTER_LAYERS)} "
          f"(dim={dim}, rank={rank})")


def safe_decoding(model, cuda_graph: bool = False):
    """Greedy batched decoding with the CUDA-graph path off by default.

    The graph decoder is only a throughput optimisation and is sensitive to driver / NeMo
    versions; transcripts are identical either way.
    """
    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)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("audio", nargs="+", help="16 kHz mono wav files")
    ap.add_argument("--lang", choices=["en", "other"], required=True,
                    help="'en' = bare trunk; 'other' = inject the adapters")
    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("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--batch_size", type=int, default=8)
    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)
    if a.lang == "other":
        attach_adapters(model, a.adapters, device, dtype)
    else:
        print("[CAFE] English path: bare trunk, adapters not attached")

    with torch.inference_mode():
        out = model.transcribe(list(a.audio), batch_size=a.batch_size, verbose=False)
    if isinstance(out, tuple):
        out = out[0]
    for f, h in zip(a.audio, out):
        print(f"{f}\t{getattr(h, 'text', h)}")


if __name__ == "__main__":
    main()