#!/usr/bin/env python3
"""Convert Cosentino's pinned 8 kHz / two-source checkpoint for local Core ML.

Original: https://zenodo.org/records/3873572 (CC BY-SA 3.0).
Architecture follows the MIT-licensed Asteroid ConvTasNet implementation.
This file is build tooling, not an iPhone dependency. No audio leaves this Mac.
Usage: PYTHONPATH=<coremltools-directory> python3 scripts/convert_voice_separation.py
"""
from pathlib import Path
import hashlib
import json
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
import coremltools as ct
from coremltools.models.neural_network.quantization_utils import quantize_weights

ROOT = Path(__file__).resolve().parent.parent
ORIGINAL = ROOT / "output/models/convtasnet-libri2mix-original.pth"
DESTINATION = ROOT / "output/models/FieldRelayTwoVoices-v1.mlmodel"
EXPECTED_MD5 = "07edec1e2ef18bad689acc5f3f06b976"


class FixedConvTasNet(nn.Module):
    """Fixed 4-second shape; all trained tensors loaded, none fine-tuned."""
    def __init__(self, state):
        super().__init__()
        self.names = {}
        self.used = set()
        for index, (key, value) in enumerate(state.items()):
            name = f"trained_{index}"
            self.names[key] = name
            self.register_buffer(name, value)

    def weight(self, name):
        self.used.add(name)
        return getattr(self, self.names[name])

    def norm(self, x, prefix):
        mean = x.mean(dim=(1, 2), keepdim=True)
        centered = x - mean
        variance = (centered * centered).mean(dim=(1, 2), keepdim=True)
        return centered / torch.sqrt(variance + 1e-8) * self.weight(prefix + ".gamma").view(1, -1, 1) + self.weight(prefix + ".beta").view(1, -1, 1)

    def conv(self, x, prefix, **kwargs):
        return F.conv1d(x, self.weight(prefix + ".weight"), self.weight(prefix + ".bias"), **kwargs)

    def forward(self, mixture):
        encoded = F.conv1d(mixture.view(1, 1, 32000), self.weight("encoder.filterbank._filters"), stride=8)
        current = self.conv(self.norm(encoded, "masker.bottleneck.0"), "masker.bottleneck.1")
        skip_sum = 0
        for index in range(24):
            prefix = f"masker.TCN.{index}"
            shared = self.conv(current, prefix + ".shared_block.0")
            shared = F.prelu(shared, self.weight(prefix + ".shared_block.1.weight"))
            shared = self.norm(shared, prefix + ".shared_block.2")
            dilation = 2 ** (index % 8)
            shared = self.conv(shared, prefix + ".shared_block.3", padding=dilation, dilation=dilation, groups=512)
            shared = F.prelu(shared, self.weight(prefix + ".shared_block.4.weight"))
            shared = self.norm(shared, prefix + ".shared_block.5")
            current = current + self.conv(shared, prefix + ".res_conv")
            skip_sum = skip_sum + self.conv(shared, prefix + ".skip_conv")
        masks = self.conv(F.prelu(skip_sum, self.weight("masker.mask_net.0.weight")), "masker.mask_net.1")
        masked = encoded.unsqueeze(1) * F.relu(masks.view(1, 2, 512, 3999))
        decoded = F.conv_transpose1d(masked.view(2, 512, 3999), self.weight("decoder.filterbank._filters"), stride=8)
        unscaled_waves = decoded.view(1, 2, 32000)
        # SI-SDR training leaves output gain unconstrained. Calibrate both voices
        # together against their mixture; never independently amplify a quiet voice.
        combined = unscaled_waves.sum(dim=1, keepdim=True)
        gain = (combined * mixture.view(1, 1, 32000)).sum() / ((combined * combined).sum() + 1e-8)
        calibrated = unscaled_waves * gain
        return calibrated / torch.clamp(calibrated.abs().max(), min=1.0)


def main():
    raw = ORIGINAL.read_bytes()
    if hashlib.md5(raw).hexdigest() != EXPECTED_MD5:
        raise SystemExit("Original checkpoint checksum mismatch")
    # Permit only NumPy scalar metadata; never disable the restricted weights loader.
    torch.serialization.add_safe_globals([(np._core.multiarray.scalar, "numpy.core.multiarray.scalar"), np.dtype, np.dtypes.Float64DType])
    checkpoint = torch.load(ORIGINAL, map_location="cpu", weights_only=True)
    if checkpoint["model_name"] != "ConvTasNet" or checkpoint["model_args"]["n_src"] != 2:
        raise SystemExit("Unexpected checkpoint architecture")
    torch.set_num_threads(2)
    torch.manual_seed(7)
    model = FixedConvTasNet(checkpoint["state_dict"]).eval()
    example = torch.randn(1, 32000) * 0.08
    with torch.inference_mode():
        reference = model(example).numpy()
        if model.used != set(checkpoint["state_dict"]):
            raise SystemExit("Some trained tensors were not used")
        traced = torch.jit.trace(model, example)
    converted = ct.convert(traced, inputs=[ct.TensorType(name="mixture", shape=(1, 32000))],
        outputs=[ct.TensorType(name="sources")], convert_to="neuralnetwork", minimum_deployment_target=ct.target.iOS14)
    converted = quantize_weights(converted, nbits=16)
    spec = converted.get_spec()
    spec.description.output[0].type.multiArrayType.shape.extend([1, 2, 32000])
    converted = ct.models.MLModel(spec)
    converted.author = "Cosentino Joris; Core ML conversion for FieldRelay"
    converted.license = "CC BY-SA 3.0; https://creativecommons.org/licenses/by-sa/3.0/"
    converted.short_description = "Experimental two-voice separation, 8 kHz, four-second windows; not speaker identification."
    converted.user_defined_metadata["original"] = "https://zenodo.org/records/3873572"
    converted.user_defined_metadata["modifications"] = "Fixed four-second input; common mixture gain calibration and peak limit; FP16 weight storage; no retraining."
    converted.save(str(DESTINATION))
    predicted = converted.predict({"mixture": example.numpy()})["sources"]
    difference = float(np.max(np.abs(predicted - reference)))
    relative_rmse = float(np.sqrt(np.mean((predicted - reference) ** 2)) / max(1e-8, np.sqrt(np.mean(reference ** 2))))
    # SI-SDR-trained separators have arbitrary output gain. Check relative parity,
    # not an absolute waveform threshold; playback gain is handled by the app.
    if not np.all(np.isfinite(predicted)) or relative_rmse > 0.005:
        raise SystemExit(f"Core ML conversion parity failed: {difference}")
    report = {"originalURL": "https://zenodo.org/records/3873572/files/model.pth", "originalMD5": EXPECTED_MD5,
        "originalBytes": len(raw), "modelBytes": DESTINATION.stat().st_size,
        "sha256": hashlib.sha256(DESTINATION.read_bytes()).hexdigest(),
        "license": "CC BY-SA 3.0", "coremltools": ct.__version__, "torch": torch.__version__,
        "sampleRate": 8000, "inputShape": [1, 32000], "outputShape": [1, 2, 32000],
        "weightStorage": "FP16", "maxAbsoluteConversionError": difference,
        "relativeConversionRMSE": relative_rmse,
        "validation": "Synthetic random input on Mac; not iPhone or real mixed-language conversation."}
    (ROOT / "output/verification/voice-separation-conversion.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()