#!/usr/bin/env python3
"""
Export the HTDemucs split models (4-stem and 6-stem) as fp16 CoreML mlpackages.

This mirrors convert_htdemucs_6s_split_coreml.py but:
  * parameterises the model name, so the same wrapper drives both htdemucs
    (4 sources) and htdemucs_6s (6 sources) - the wrapper already derives the
    output channel counts from len(model.sources);
  * emits fp16 weight/compute precision instead of fp32.

STFT/iSTFT stay outside the CoreML graph, exactly as in the fp32 exports, so
the two extra outputs keep the same interface the app already expects:

  inputs : spectral_magnitude [1, 4, 2048, 336]
           audio_waveform     [1, 2, 343980]
  outputs: freq_output        [1, stems*4, 2048, 336]
           time_output        [1, stems*2, 343980]

Usage:
  python convert_split_fp16_coreml.py --outdir /path/to/output
"""

from __future__ import annotations

import argparse
from pathlib import Path

import coremltools as ct
import torch
from demucs.pretrained import get_model
from einops import rearrange


SEGMENT_SAMPLES = 343_980
NUM_BINS = 2_048
NUM_FRAMES = 336

# (demucs model name, output mlpackage basename, source order description)
MODELS = [
    (
        "htdemucs",
        "HTDemucs_SourceSeparation_F16",
        "drums, bass, other, vocals",
    ),
    (
        "htdemucs_6s",
        "HTDemucs_6S_Split_F16",
        "drums, bass, other, vocals, guitar, piano",
    ),
]


class HTDemucsSplitExport(torch.nn.Module):
    """Waveform/spectral split wrapper - identical logic to the fp32 6-stem
    export, with the stem count taken from the model itself."""

    def __init__(self, model: torch.nn.Module):
        super().__init__()
        self.model = model

    def forward(
        self,
        spectral_magnitude: torch.Tensor,
        audio_waveform: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        model = self.model
        x = spectral_magnitude
        xt = audio_waveform

        batch = x.shape[0]
        spec_channels = x.shape[1]
        freq_bins = x.shape[2]
        frames = x.shape[3]

        mean = x.mean(dim=(1, 2, 3), keepdim=True)
        std = x.std(dim=(1, 2, 3), keepdim=True)
        x = (x - mean) / (1e-5 + std)

        time_mean = xt.mean(dim=(1, 2), keepdim=True)
        time_std = xt.std(dim=(1, 2), keepdim=True)
        xt = (xt - time_mean) / (1e-5 + time_std)

        saved = []
        saved_t = []
        lengths = []
        lengths_t = []

        for idx, encode in enumerate(model.encoder):
            lengths.append(x.shape[-1])
            inject = None

            if idx < len(model.tencoder):
                lengths_t.append(xt.shape[-1])
                tenc = model.tencoder[idx]
                xt = tenc(xt)

                if not tenc.empty:
                    saved_t.append(xt)
                else:
                    inject = xt

            x = encode(x, inject)

            if idx == 0 and model.freq_emb is not None:
                freqs = torch.arange(x.shape[-2], device=x.device)
                emb = model.freq_emb(freqs).t()[None, :, :, None].expand_as(x)
                x = x + model.freq_emb_scale * emb

            saved.append(x)

        if model.crosstransformer:
            if model.bottom_channels:
                _, _, freq, _ = x.shape
                x = rearrange(x, "b c f t -> b c (f t)")
                x = model.channel_upsampler(x)
                x = rearrange(x, "b c (f t) -> b c f t", f=freq)
                xt = model.channel_upsampler_t(xt)

            x, xt = model.crosstransformer(x, xt)

            if model.bottom_channels:
                _, _, freq, _ = x.shape
                x = rearrange(x, "b c f t -> b c (f t)")
                x = model.channel_downsampler(x)
                x = rearrange(x, "b c (f t) -> b c f t", f=freq)
                xt = model.channel_downsampler_t(xt)

        for idx, decode in enumerate(model.decoder):
            skip = saved.pop(-1)
            x, pre = decode(x, skip, lengths.pop(-1))

            offset = model.depth - len(model.tdecoder)
            if idx >= offset:
                tdec = model.tdecoder[idx - offset]
                length_t = lengths_t.pop(-1)

                if tdec.empty:
                    pre = pre[:, :, 0]
                    xt, _ = tdec(pre, None, length_t)
                else:
                    skip = saved_t.pop(-1)
                    xt, _ = tdec(xt, skip, length_t)

        stem_count = len(model.sources)

        freq_out = x.view(batch, stem_count, -1, freq_bins, frames)
        freq_out = freq_out * std[:, None] + mean[:, None]
        freq_out = freq_out.reshape(batch, stem_count * spec_channels, freq_bins, frames)

        time_out = xt.view(batch, stem_count, -1, audio_waveform.shape[-1])
        time_out = time_out * time_std[:, None] + time_mean[:, None]
        time_out = time_out.reshape(batch, stem_count * 2, audio_waveform.shape[-1])

        return freq_out, time_out


def convert(model_name: str, out_name: str, source_order: str, outdir: Path) -> Path:
    print(f"[convert] {model_name} -> {out_name} (fp16)")
    bag = get_model(model_name)
    model = bag.models[0] if hasattr(bag, "models") else bag
    model.eval()

    stem_count = len(model.sources)
    print(f"  sources: {model.sources}")
    print(f"  samplerate: {model.samplerate}  segment: {int(model.segment * model.samplerate)}")

    wrapper = HTDemucsSplitExport(model).eval()
    spectral = torch.randn(1, 4, NUM_BINS, NUM_FRAMES)
    waveform = torch.randn(1, 2, SEGMENT_SAMPLES)

    with torch.no_grad():
        freq_out, time_out = wrapper(spectral, waveform)
    print(f"  freq_output: {tuple(freq_out.shape)}  time_output: {tuple(time_out.shape)}")

    traced = torch.jit.trace(
        wrapper,
        (spectral, waveform),
        strict=False,
        check_trace=False,
    )

    mlmodel = ct.convert(
        traced,
        inputs=[
            ct.TensorType(name="spectral_magnitude", shape=spectral.shape),
            ct.TensorType(name="audio_waveform", shape=waveform.shape),
        ],
        outputs=[
            ct.TensorType(name="freq_output"),
            ct.TensorType(name="time_output"),
        ],
        minimum_deployment_target=ct.target.iOS16,
        convert_to="mlprogram",
        compute_precision=ct.precision.FLOAT16,
    )

    mlmodel.author = "Meta Research (Demucs)"
    mlmodel.license = "MIT License"
    mlmodel.short_description = (
        f"HTDemucs {stem_count}-stem split CoreML export (fp16). STFT/iSTFT are "
        "computed by the app; the model runs the HTDemucs neural network body."
    )
    mlmodel.input_description["spectral_magnitude"] = (
        "Stereo complex STFT [1,4,2048,336]: L.real, L.imag, R.real, R.imag."
    )
    mlmodel.input_description["audio_waveform"] = (
        "Stereo waveform [1,2,343980] at 44100 Hz."
    )
    mlmodel.output_description["freq_output"] = (
        f"Frequency branch [1,{stem_count * 4},2048,336]. Source order: {source_order}."
    )
    mlmodel.output_description["time_output"] = (
        f"Time branch [1,{stem_count * 2},343980]. Source order: {source_order}."
    )

    out_path = outdir / f"{out_name}.mlpackage"
    if out_path.exists():
        import shutil
        shutil.rmtree(out_path)
    mlmodel.save(str(out_path))
    return out_path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--outdir",
        type=Path,
        required=True,
        help="Directory to write the .mlpackage outputs into.",
    )
    parser.add_argument(
        "--only",
        choices=[m[0] for m in MODELS],
        help="Convert just one model (default: both).",
    )
    args = parser.parse_args()
    args.outdir.mkdir(parents=True, exist_ok=True)

    for model_name, out_name, source_order in MODELS:
        if args.only and model_name != args.only:
            continue
        path = convert(model_name, out_name, source_order, args.outdir)
        size = sum(f.stat().st_size for f in path.rglob("*") if f.is_file())
        print(f"[done] {path}  ({size / 1024 / 1024:.1f} MB)")


if __name__ == "__main__":
    main()