"""
Cohere Transcribe OpenVINO inference.

Usage:
    python inference.py --model_dir . --audio sample.wav --device CPU
"""

import argparse
import json
import time
from pathlib import Path

import numpy as np
import openvino as ov
from transformers import AutoProcessor

META_JSON = "ov_cohere_transcribe_kvcache.json"


def transcribe(model_dir: Path, audio: str, device: str, max_new_tokens: int = 256):
    import librosa

    meta = json.loads((model_dir / META_JSON).read_text())
    num_layers, sr = meta["num_layers"], meta["sampling_rate"]
    eos = meta["eos_token_id"]
    eos_set = set(eos) if isinstance(eos, (list, tuple)) else {eos}

    processor = AutoProcessor.from_pretrained(model_dir)
    core = ov.Core()
    encoder = core.compile_model(model_dir / meta["encoder_ir"], device)
    prefill = core.compile_model(model_dir / meta["decoder_ir"], device)
    decode = core.compile_model(model_dir / meta["decoder_with_past_ir"], device)

    wav, _ = librosa.load(audio, sr=sr, mono=True)
    inp = processor(wav, sampling_rate=sr, language="en", return_tensors="np")
    feat = np.asarray(inp["input_features"]).astype(np.float32)
    amask = np.asarray(inp["attention_mask"]).astype(bool)
    prompt = np.asarray(inp["decoder_input_ids"]).astype(np.int64)

    t0 = time.perf_counter()
    enc = encoder({"input_features": feat, "attention_mask": amask})
    ehs, emask = enc["encoder_hidden_states"], enc["encoder_attention_mask"]

    pf = {k.get_any_name(): v for k, v in prefill(
        {"decoder_input_ids": prompt, "encoder_hidden_states": ehs, "encoder_attention_mask": emask}
    ).items()}
    self_kv = [(pf[f"present.{i}.self.key"], pf[f"present.{i}.self.value"]) for i in range(num_layers)]
    cross_kv = [(pf[f"present.{i}.cross.key"], pf[f"present.{i}.cross.value"]) for i in range(num_layers)]
    next_id = int(pf["logits"][0, -1].argmax())
    generated = [next_id]

    for _ in range(max_new_tokens):
        if next_id in eos_set:
            break
        feed = {
            "decoder_input_ids": np.array([[next_id]], dtype=np.int64),
            "encoder_hidden_states": ehs,
            "encoder_attention_mask": emask,
            "self_attention_mask": np.ones((1, self_kv[0][0].shape[2] + 1), dtype=np.int64),
        }
        for i in range(num_layers):
            feed[f"past.{i}.self.key"], feed[f"past.{i}.self.value"] = self_kv[i]
            feed[f"past.{i}.cross.key"], feed[f"past.{i}.cross.value"] = cross_kv[i]
        step = {k.get_any_name(): v for k, v in decode(feed).items()}
        self_kv = [(step[f"present.{i}.self.key"], step[f"present.{i}.self.value"]) for i in range(num_layers)]
        next_id = int(step["logits"][0, -1].argmax())
        generated.append(next_id)

    elapsed = time.perf_counter() - t0
    text = processor.batch_decode([generated], skip_special_tokens=True)[0].strip()
    return text, elapsed, len(generated)


def main():
    here = Path(__file__).parent
    default_dir = here / "models" if (here / "models").is_dir() else here
    ap = argparse.ArgumentParser(description="Transcribe audio with Cohere Transcribe OpenVINO IR")
    ap.add_argument("--model_dir", type=Path, default=default_dir, help="Directory with the OpenVINO IR")
    ap.add_argument("--audio", required=True, help="Path to input audio (.wav)")
    ap.add_argument("--device", default="CPU", help="OpenVINO device: CPU / GPU / NPU / AUTO")
    ap.add_argument("--max_new_tokens", type=int, default=256)
    args = ap.parse_args()

    text, elapsed, n_tokens = transcribe(args.model_dir.resolve(), args.audio, args.device, args.max_new_tokens)
    print(f"\nDevice     : {args.device}")
    print(f"Tokens     : {n_tokens}")
    print(f"Time       : {elapsed:.3f} s  ({n_tokens / elapsed:.1f} tok/s)")
    print(f"Transcript : {text}")


if __name__ == "__main__":
    main()