#!/usr/bin/env python3
"""Export Nemotron-3-Diarization (Sortformer) to CoreML via forward_for_export.

The export boundary is NVIDIA's own streaming graph:

    (chunk, chunk_lengths, spkcache, spkcache_lengths, fifo, fifo_lengths)
        -> (spkcache_fifo_chunk_preds, chunk_pre_encode_embs, chunk_pre_encode_lengths)

Mel features are deliberately NOT part of the graph: the host computes them, so
front-end drift stays isolated from conversion error (gate 1 vs gates 2-3).
The speaker cache / FIFO are host-owned state passed in per chunk, which keeps
the CoreML model stateless and matches how FluidAudio drives Sortformer today.
"""
from __future__ import annotations

import argparse
import warnings
from pathlib import Path

import numpy as np
import torch

warnings.filterwarnings("ignore")

import coremltools as ct  # noqa: E402
from nemo.collections.asr.models import SortformerEncLabelModel  # noqa: E402

import coreml_patches  # noqa: E402

coreml_patches.apply()

HERE = Path(__file__).resolve().parent
NEMO_PATH = HERE.parent / "Nemotron-3-Diarization-preview.nemo"
# The final release ships bare weights; its graph is identical to the preview's.
FINAL_WEIGHTS = Path.home() / "Nemotron-3-Diarization-safetensors" / "model.safetensors"
# Auxiliary 3-class activity head: feeds only the training loss upstream, never the speaker preds.
TRAINING_ONLY_PREFIX = "sortformer_modules.activity_head."

# Latency profiles from the model card. Frames are 80 ms; the input-buffer
# latency is (chunk_len + chunk_right_context) * 80 ms. spkcache_len is 264 in
# every profile, so only the chunk/FIFO geometry changes.
PROFILES = {
    # 30.4 s buffer -- batch use, one call covers half a minute of audio.
    "offline": dict(chunk_len=340, chunk_right_context=40, fifo_len=40,
                    spkcache_update_period=300),
    # 1.04 s buffer -- the usable streaming profile.
    "low": dict(chunk_len=9, chunk_right_context=4, fifo_len=264,
                spkcache_update_period=222),
    # 0.64 s buffer.
    "verylow": dict(chunk_len=6, chunk_right_context=2, fifo_len=264,
                    spkcache_update_period=222),
    # 0.32 s buffer -- lowest latency the card documents.
    "ultralow": dict(chunk_len=3, chunk_right_context=1, fifo_len=264,
                     spkcache_update_period=222),
}
PROFILE_SECONDS = {
    name: (cfg["chunk_len"] + cfg["chunk_right_context"]) * 0.08
    for name, cfg in PROFILES.items()
}

OFFLINE_PROFILE = PROFILES["offline"]

OUTPUT_NAMES = ["spkcache_fifo_chunk_preds", "chunk_pre_encode_embs",
                "chunk_pre_encode_lengths"]


class StreamingWrapper(torch.nn.Module):
    """Trace-friendly shim.

    Lengths cross the CoreML boundary as int32 (CoreML has no int64 inputs) and
    are widened back to int64 here, where the NeMo graph expects them.
    """

    def __init__(self, model: SortformerEncLabelModel):
        super().__init__()
        self.model = model

    def forward(self, chunk, chunk_lengths, spkcache, spkcache_lengths, fifo, fifo_lengths):
        return self.model.forward_for_export(
            chunk,
            chunk_lengths.to(torch.int64),
            spkcache,
            spkcache_lengths.to(torch.int64),
            fifo,
            fifo_lengths.to(torch.int64),
        )


def load_model(profile: dict, weights: str = "preview") -> SortformerEncLabelModel:
    model = SortformerEncLabelModel.restore_from(
        restore_path=str(NEMO_PATH), map_location="cpu", strict=False
    )
    if weights == "final":
        from safetensors.torch import load_file

        state = {
            key: value.float()
            for key, value in load_file(str(FINAL_WEIGHTS)).items()
            if not key.startswith(TRAINING_ONLY_PREFIX)
        }
        missing, unexpected = model.load_state_dict(state, strict=False)
        if missing or unexpected:
            raise SystemExit(f"weight mismatch: missing={missing[:5]} unexpected={unexpected[:5]}")
    model.eval()
    for key, value in profile.items():
        setattr(model.sortformer_modules, key, value)
    return model


def build_examples(model: SortformerEncLabelModel):
    """Model-sized examples, but with lengths as int32 to match the CoreML iface."""
    chunk, chunk_len, spkcache, spkcache_len, fifo, fifo_len = model.streaming_input_examples(1)
    return (
        chunk,
        chunk_len.to(torch.int32),
        spkcache,
        spkcache_len.to(torch.int32),
        fifo,
        fifo_len.to(torch.int32),
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp16")
    parser.add_argument("--profile", choices=sorted(PROFILES), default="offline")
    parser.add_argument("--weights", choices=["preview", "final"], default="preview")
    parser.add_argument("--out", type=Path, default=None)
    args = parser.parse_args()

    suffix = "" if args.profile == "offline" else f"_{args.profile}"
    prefix = "nemotron_diar" if args.weights == "preview" else "nemotron_diar_final"
    out_path = args.out or (HERE / "artifacts" /
                            f"{prefix}_{args.precision}{suffix}.mlpackage")
    out_path.parent.mkdir(parents=True, exist_ok=True)

    profile = PROFILES[args.profile]
    print(f"profile: {args.profile} ({PROFILE_SECONDS[args.profile]:.2f}s buffer) {profile}")
    model = load_model(profile, args.weights)
    wrapper = StreamingWrapper(model).eval()
    examples = build_examples(model)

    for name, tensor in zip(
        ["chunk", "chunk_lengths", "spkcache", "spkcache_lengths", "fifo", "fifo_lengths"],
        examples,
    ):
        print(f"  input {name:<18} {tuple(tensor.shape)} {tensor.dtype}")

    with torch.no_grad():
        reference = wrapper(*examples)
    for name, tensor in zip(OUTPUT_NAMES, reference):
        print(f"  output {name:<28} {tuple(tensor.shape)} {tensor.dtype}")

    print("\ntracing...")
    with torch.no_grad():
        traced = torch.jit.trace(wrapper, examples, strict=False)
    # Every shape in this graph is static (batch 1, fixed chunk/cache/FIFO), but the
    # trace still carries ~200 aten::Int shape computations. Freezing constant-folds
    # them; coremltools chokes on the ones that carry (1,)-shaped tensors.
    traced = torch.jit.freeze(traced)
    traced = torch.jit.run_frozen_optimizations(traced) or traced

    print(f"converting to CoreML ({args.precision})...")
    precision = ct.precision.FLOAT16 if args.precision == "fp16" else ct.precision.FLOAT32
    mlmodel = ct.convert(
        traced,
        convert_to="mlprogram",
        inputs=[
            ct.TensorType(name="chunk", shape=examples[0].shape, dtype=np.float32),
            ct.TensorType(name="chunk_lengths", shape=examples[1].shape, dtype=np.int32),
            ct.TensorType(name="spkcache", shape=examples[2].shape, dtype=np.float32),
            ct.TensorType(name="spkcache_lengths", shape=examples[3].shape, dtype=np.int32),
            ct.TensorType(name="fifo", shape=examples[4].shape, dtype=np.float32),
            ct.TensorType(name="fifo_lengths", shape=examples[5].shape, dtype=np.int32),
        ],
        outputs=[ct.TensorType(name=name) for name in OUTPUT_NAMES],
        compute_precision=precision,
        minimum_deployment_target=ct.target.macOS14,
    )

    mlmodel.save(str(out_path))
    size_mb = sum(f.stat().st_size for f in out_path.rglob("*") if f.is_file()) / 1e6
    print(f"\nsaved {out_path}  ({size_mb:.1f} MB)")


if __name__ == "__main__":
    main()