import os, sys, json
import numpy as np
import torch
from omegaconf import OmegaConf, open_dict
import nemo.collections.asr as nemo_asr

ATT = [int(x) for x in sys.argv[1].split(",")]   # e.g. "70,6"
OUT = sys.argv[2]
os.makedirs(OUT, exist_ok=True)

LATENCY = {(70,0):80,(70,1):160,(70,6):560,(70,13):1120}

print(f"[export] att_context_size={ATT} -> {OUT}", flush=True)

m = nemo_asr.models.ASRModel.from_pretrained("nvidia/nemotron-speech-streaming-en-0.6b")
m = m.to("cpu")
m.eval()

# 1) Select chunk size (latency) -- bakes it into the graph
m.encoder.set_default_att_context_size(ATT)
print(f"[export] actual att_context_size={m.encoder.att_context_size}", flush=True)

# 2) Greedy RNNT decode strategy (informational for export; matches runtime MAX_SYMBOLS=10)
decoding_cfg = m.cfg.decoding
with open_dict(decoding_cfg):
    decoding_cfg.strategy = "greedy"
    decoding_cfg.preserve_alignments = False
    if hasattr(m, "joint"):
        decoding_cfg.greedy.max_symbols = 10
        decoding_cfg.fused_batch_size = -1
m.change_decoding_strategy(decoding_cfg)
m.eval()

# 3) Cache-aware streaming export MUST be on (NeMo defaults off)
m.set_export_config({"cache_support": "True"})

# 4) Export. NeMo auto-splits RNNT into encoder + decoder_joint. Opset 17.
m.export(os.path.join(OUT, "model.onnx"), onnx_opset_version=17, check_trace=False)

# 5) Dump shared support files straight from the checkpoint
sc = m.encoder.streaming_cfg
chunk_mel = int(sc.chunk_size[1])           # verified: chunk_size[1] == chunk_mel_frames

# tokens.txt : SPACE-separated "piece id", 1024 lines (ids 0..1023); blank id 1024 ABSENT
tok = m.tokenizer
vocab = tok.vocab
assert len(vocab) == 1024, f"unexpected vocab size {len(vocab)}"
with open(os.path.join(OUT, "tokens.txt"), "w", encoding="utf-8") as f:
    for i, piece in enumerate(vocab):
        f.write(f"{piece} {i}\n")

# filterbank.bin : [1,128,257] float32 band-major -> EXACTLY 131584 bytes
fb = m.preprocessor.filter_banks.detach().cpu().numpy().astype(np.float32)
assert fb.shape == (1, 128, 257), f"unexpected filterbank shape {fb.shape}"
fb.tofile(os.path.join(OUT, "filterbank.bin"))
with open(os.path.join(OUT, "filterbank.meta"), "w") as f:
    f.write("shape=1x128x257")

# preprocessor.config : YAML; must show normalize=NA (CMVN OFF)
pp = OmegaConf.to_container(m.cfg.preprocessor, resolve=True)
with open(os.path.join(OUT, "preprocessor.config"), "w") as f:
    f.write(OmegaConf.to_yaml(OmegaConf.create({"preprocessor": pp})))

# config.json : match the app's exact schema. Static fields constant across sizes;
# per-size fields derived from streaming_cfg (verified against the reference table).
cfg = {
    "model_type": "fastconformer_rnnt",
    "source_model": "nvidia/nemotron-speech-streaming-en-0.6b",
    "audio": {"sample_rate": 16000, "sample_format": "S16_LE", "bytes_per_second": 32000},
    "preprocessor": {
        "n_mels": 128, "n_fft": 512, "hop_length": 160, "win_length": 400, "window": "hann",
        "preemph": 0.97, "dither": 1e-05, "normalize": None, "pad_to": 0,
        "mel_norm": "slaney", "mel_layout": "band_major",
    },
    "encoder": {
        "model_file": "encoder_model.onnx", "layers": 24, "dim": 1024,
        "chunk_mel_frames": chunk_mel, "pre_encode_cache_frames": 9,
        "total_input_frames": chunk_mel + 9,
        "cache_last_channel_shape": [1, 24, 70, 1024],
        "cache_last_time_shape": [1, 24, 1024, 8],
    },
    "decoder": {
        "model_file": "decoder_model.onnx", "prediction_layers": 2, "prediction_hidden": 640,
        "vocab_size": 1025, "blank_id": 1024, "max_symbols_per_frame": 10,
    },
    "streaming": {"chunk_duration_ms": LATENCY[tuple(ATT)], "chunk_audio_samples": chunk_mel * 160},
}
with open(os.path.join(OUT, "config.json"), "w") as f:
    json.dump(cfg, f, indent=2)

print(f"[export] DONE att={ATT} chunk_mel={chunk_mel} total_in={chunk_mel+9} samples={chunk_mel*160}", flush=True)