"""Converts Alibaba's MossFormer2_SS_16K (ClearerVoice-Studio) to Core ML, then checks the Core ML model against the
PyTorch one on the same input.

    pip install clearvoice coremltools
    python convert.py

Writes MossFormer2_SS_16K.mlpackage. Converted with clearvoice 0.1.2, torch 2.15 and coremltools 9.0 on macOS 26.
"""
import time

import numpy as np
import torch
import torch.backends.mps

torch.backends.mps.is_available = lambda: False
import coremltools as ct
from clearvoice import ClearVoice

# 4,096 encoder steps (kernel 16, stride 8): 16 whole attention groups of 256, so attention needs no padding mask,
# which Core ML can't convert.
WINDOW = 32776

network = ClearVoice(task="speech_separation", model_names=["MossFormer2_SS_16K"]).models[0].model.eval().cpu()

# The token shift pads by -1 at one end, and Core ML's pad takes only non-negative amounts: cropping where an amount
# is negative, then padding the rest, gives the same tensor.
_pad = torch.nn.functional.pad


def pad(x, amounts, mode="constant", value=None):
    if all(a >= 0 for a in amounts):
        return _pad(x, amounts, mode=mode, value=value)
    keep = [slice(None)] * x.dim()
    positive = []
    for i in range(0, len(amounts), 2):
        dim, before, after = x.dim() - 1 - i // 2, amounts[i], amounts[i + 1]
        keep[dim] = slice(-before if before < 0 else 0, x.shape[dim] + after if after < 0 else x.shape[dim])
        positive += [max(before, 0), max(after, 0)]
    return _pad(x[tuple(keep)], positive, mode=mode, value=value)


# Core ML converts only some einsum forms; each one MossFormer2 uses is a plain or a matrix product.
_einsum = torch.einsum


def einsum(equation, *operands):
    a, b = operands[0], operands[1] if len(operands) > 1 else None
    form = equation.replace(" ", "")
    if form == "...d,hd->...hd":
        return a.unsqueeze(-2) * b
    if form == "...id,...jd->...ij":
        return a @ b.transpose(-1, -2)
    if form == "...ij,...jd->...id":
        return a @ b
    if form == "bgnd,bgne->bde":
        batch, groups, n, d = a.shape
        return a.reshape(batch, groups * n, d).transpose(1, 2) @ b.reshape(batch, groups * n, -1)
    if form == "bgnd,bde->bgne":
        return a @ b.unsqueeze(1)
    if form in ("i,j->ij", "...,f->...f"):
        return a.unsqueeze(-1) * b
    raise ValueError(f"einsum form not handled: {equation}")


from clearvoice.models.mossformer2_ss import mossformer2_block
import rotary_embedding_torch.rotary_embedding_torch as rotary

torch.nn.functional.pad = pad
mossformer2_block.einsum = einsum
if hasattr(rotary, "einsum"):
    rotary.einsum = einsum


class Separate(torch.nn.Module):
    """One tensor in and one out: audio [1, WINDOW] -> voices [1, 2, WINDOW]."""

    def __init__(self, network):
        super().__init__()
        self.network = network

    def forward(self, audio):
        return torch.stack(self.network(audio), dim=1)


model = Separate(network).eval()
example = torch.from_numpy(np.random.default_rng(0).standard_normal((1, WINDOW)).astype(np.float32) * 0.05)
with torch.no_grad():
    reference = model(example).numpy()[0]

started = time.time()
traced = torch.jit.trace(model, example, check_trace=False)
mlmodel = ct.convert(
    traced,
    inputs=[ct.TensorType(name="audio", shape=(1, WINDOW), dtype=np.float32)],
    outputs=[ct.TensorType(name="voices", dtype=np.float32)],
    convert_to="mlprogram",
    compute_precision=ct.precision.FLOAT16,
    minimum_deployment_target=ct.target.macOS15,
)
mlmodel.short_description = "MossFormer2 (ClearerVoice-Studio) speech separation, 16 kHz, two voices, 2.05 s windows"
mlmodel.license = "Apache-2.0"
mlmodel.save("MossFormer2_SS_16K.mlpackage")
print(f"converted in {time.time() - started:.0f} s")


def si_sdr(estimate, target):
    estimate, target = estimate - estimate.mean(), target - target.mean()
    projection = np.dot(estimate, target) / np.dot(target, target) * target
    return 10 * np.log10(np.sum(projection ** 2) / np.sum((estimate - projection) ** 2))


loaded = ct.models.MLModel("MossFormer2_SS_16K.mlpackage", compute_units=ct.ComputeUnit.CPU_AND_GPU)
voices = loaded.predict({"audio": example.numpy()})["voices"][0]
for k in range(2):
    print(f"voice {k + 1}: Core ML vs PyTorch {si_sdr(voices[k], reference[k]):.1f} dB SI-SDR")