"""Convert the Haas-ja-v1.0 FastConformer encoder to CoreML (.mlpackage, fp16).

Usage: python convert.py <T_mel> [out.mlpackage]
Input: audio_signal [1,128,T_mel] f32 (zero past length), length [1] i32 -> encoded [1,1024,T_mel/8]
"""
import os, sys
# HAAS_EXPORT_DIR: dir containing the haas_onnx package (scripts/ of Sanami/haas-ja-onnx)
# HAAS_NEMO_DIR: the extracted Haas-ja-v1.0.nemo (model_config.yaml, model_weights.ckpt)
sys.path.insert(0, os.environ.get("HAAS_EXPORT_DIR", "export"))
import numpy as np, torch, coremltools as ct
from haas_onnx.load import load_model
import cmlpatch  # noqa: F401

T = int(sys.argv[1])
out = sys.argv[2] if len(sys.argv) > 2 else f"encoder_T{T}.mlpackage"


class Enc(torch.nn.Module):
    def __init__(self, enc):
        super().__init__()
        self.enc = enc

    def forward(self, audio_signal, length=None):
        if length is None:
            length = torch.full((1,), audio_signal.shape[-1], dtype=torch.long)
        y, _ = self.enc(audio_signal=audio_signal, length=length.long())
        return y


model = load_model(os.environ.get("HAAS_NEMO_DIR", "nemo_x"))
enc = model.encoder.eval()
NL = int(os.environ.get("NL", "0"))  # debug: keep only the first NL layers
if NL:
    enc.layers = enc.layers[:NL]
NOLEN = os.environ.get("NOLEN") == "1"  # debug: no length input (whole window valid)
enc.update_max_seq_length(T, device="cpu") if hasattr(enc, "update_max_seq_length") else None
w = Enc(enc).eval()
x = torch.randn(1, 128, T)
L = torch.tensor([T], dtype=torch.int32)
args = (x,) if NOLEN else (x, L)
with torch.no_grad():
    ref = w(*args).numpy()
    tr = torch.jit.trace(w, args, check_trace=False)
ins = [ct.TensorType("audio_signal", shape=(1, 128, T), dtype=np.float32)]
if not NOLEN:
    ins.append(ct.TensorType("length", shape=(1,), dtype=np.int32))
ml = ct.convert(
    tr,
    inputs=ins,
    outputs=[ct.TensorType("encoded", dtype=np.float32)],
    convert_to="mlprogram", compute_precision=ct.precision.FLOAT16,
    minimum_deployment_target=ct.target.iOS17,
)
ml.save(out)
feed = {"audio_signal": x.numpy()}
if not NOLEN:
    feed["length"] = L.numpy()
y = ml.predict(feed)["encoded"]
print("shape", y.shape, "max abs diff", float(np.abs(y - ref).max()), "ref max", float(np.abs(ref).max()),
      "cos", float((y * ref).sum() / np.linalg.norm(y) / np.linalg.norm(ref)))