"""Export LiquidInbetween to Core ML.

Images in and out (RGB, 512 × 512): pass two pixel buffers and a moment
t ∈ (0, 1), get the frame at that moment. Weights are FP16 (about 8.6 MB).

    python export_coreml.py <weights.pt | teacher.safetensors> <out.mlpackage> [--int8]
"""
import sys
from pathlib import Path

import coremltools as ct
import coremltools.optimize.coreml as cto
import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).parent))
from liquid_model import IFNet, load_teacher, student_from_teacher

SIZE = 512


class Wrapper(torch.nn.Module):
    """Pixel values 0...255 in and out, as Core ML image types expect."""

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

    def forward(self, first, second, t):
        out = self.net(first / 255.0, second / 255.0, t)
        return (out.clamp(0, 1) * 255.0)


weights, out = Path(sys.argv[1]), Path(sys.argv[2])
if weights.suffix == ".safetensors":
    net = student_from_teacher(load_teacher(weights))
else:
    net = IFNet(timestep=True)
    net.load_state_dict(torch.load(weights, map_location="cpu"))
model = Wrapper(net.eval()).eval()

example = (torch.rand(1, 3, SIZE, SIZE) * 255, torch.rand(1, 3, SIZE, SIZE) * 255, torch.full((1, 1, 1, 1), 0.5))
with torch.no_grad():
    traced = torch.jit.trace(model, example)

mlmodel = ct.convert(
    traced,
    inputs=[
        ct.ImageType(name="first", shape=(1, 3, SIZE, SIZE), color_layout=ct.colorlayout.RGB),
        ct.ImageType(name="second", shape=(1, 3, SIZE, SIZE), color_layout=ct.colorlayout.RGB),
        ct.TensorType(name="t", shape=(1, 1, 1, 1), dtype=np.float32),
    ],
    outputs=[ct.ImageType(name="frame", color_layout=ct.colorlayout.RGB)],
    convert_to="mlprogram",
    compute_precision=ct.precision.FLOAT16,
    minimum_deployment_target=ct.target.iOS17,
)
if "--int8" in sys.argv:
    # Experimental: 8-bit weights halve the size, but measured 24.7 dB against
    # PyTorch (FP16 alone: 42.9 dB), visibly worse, so FP16 is the default.
    # The transposed convolutions are the flow and mask heads: a tiny weight
    # error there becomes pixels of motion error, so they stay FP16.
    config = cto.OptimizationConfig(
        global_config=cto.OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int8", granularity="per_channel"),
        op_type_configs={"conv_transpose": None},
    )
    mlmodel = cto.linear_quantize_weights(mlmodel, config)
mlmodel.short_description = "LiquidInbetween: the frame at moment t between two photos."
mlmodel.author = "Haplo LLC (jc_builds). Based on RIFE (Huang et al., ECCV 2022, MIT)."
mlmodel.license = "MIT"
mlmodel.version = "1.0"
mlmodel.input_description["first"] = "The first photo, 512 × 512 RGB."
mlmodel.input_description["second"] = "The second photo, 512 × 512 RGB."
mlmodel.input_description["t"] = "The moment between them, 0...1."
mlmodel.output_description["frame"] = "The frame at moment t."
mlmodel.save(str(out))
print("saved", out)