#!/usr/bin/env python3
"""MODNet portrait matting -> ExecuTorch WebGPU.

The official code is on GitHub (ZHKKKe/MODNet) and the official checkpoints on Google
Drive; neither is reachable from the export host and the author has no Hugging Face
account. This pins a community Hugging Face mirror that carries a copy of the upstream
source tree and the photographic checkpoint (sha256 checked; four independent mirrors
publish the same checkpoint bytes). The webcam checkpoint comes from a second mirror.

Input: [1, 3, 256, 256] float32 RGB normalized to [-1, 1] ((x / 255 - 0.5) / 0.5).
Output: [1, 1, 256, 256] alpha matte in [0, 1] (inference=True, pred_matte only).
"""

import argparse
import importlib
import sys
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parent))
import _common  # noqa: E402

SOURCE = {
    "repo": "boopathiraj/MODNet",
    "revision": "a967175f8fd63de67a1c161d8daac966d2b57e3a",
    "files": {
        "modnet/__init__.py": None,
        "modnet/modnet.py": None,
        "modnet/src/__init__.py": None,
        "modnet/src/models/__init__.py": None,
        "modnet/src/models/backbones/__init__.py": None,
        "modnet/src/models/backbones/wrapper.py": None,
        "modnet/src/models/backbones/mobilenetv2.py": None,
        "modnet/LICENSE": None,
    },
}
CHECKPOINTS = {
    "photographic": {
        "repo": "boopathiraj/MODNet",
        "revision": "a967175f8fd63de67a1c161d8daac966d2b57e3a",
        "file": "modnet/pretrained/modnet_photographic_portrait_matting.ckpt",
        "sha256": "7c22235f0925deba15d4d63e53afcb654c47055bbcd98f56e393ab2584007ed8",
    },
    "webcam": {
        "repo": "DavG25/modnet-pretrained-models",
        "revision": "903cc06311b3b12071edfa1b42b534e4a31c718a",
        "file": "models/modnet_webcam_portrait_matting.ckpt",
        "sha256": "913b82b66558db39b6286c150f809017d7528c872b156eb14333c9c6cb52108b",
    },
}


class InstanceNorm2dDecomposed(torch.nn.Module):
    """InstanceNorm2d(affine=False) as mean/sub/mul/sqrt/div.

    The Vulkan/WebGPU partitioner does not take aten._native_batch_norm_legit.no_stats,
    so the stock module would leave 16 CPU kernels splitting the graph into 13 delegates.
    """

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

    def forward(self, x):
        d = x - x.mean(dim=(2, 3), keepdim=True)
        var = (d * d).mean(dim=(2, 3), keepdim=True)
        return d / torch.sqrt(var + self.eps)


def build_model(variant):
    root = None
    files = {}
    for f in SOURCE["files"]:
        path, digest = _common.hf_file(SOURCE["repo"], f, SOURCE["revision"])
        files[f] = digest
        root = Path(str(path)[: -len(f)])
    sys.path.insert(0, str(root))
    MODNet = importlib.import_module("modnet.modnet").MODNet
    sys.path.remove(str(root))

    ck = CHECKPOINTS[variant]
    path, _ = _common.hf_file(ck["repo"], ck["file"], ck["revision"], ck["sha256"])
    state = torch.load(path, map_location="cpu", weights_only=True)
    # Checkpoints were saved from nn.DataParallel.
    state = {k.removeprefix("module."): v for k, v in state.items()}
    model = MODNet(backbone_pretrained=False)
    model.load_state_dict(state, strict=True)
    for parent in list(model.modules()):
        for name, child in parent.named_children():
            if isinstance(child, torch.nn.InstanceNorm2d):
                assert not child.affine and not child.track_running_stats
                setattr(parent, name, InstanceNorm2dDecomposed(child.eps))

    class Matte(torch.nn.Module):
        def __init__(self, m):
            super().__init__()
            self.m = m

        def forward(self, img):
            return self.m(img, True)[2]

    return Matte(model).eval(), files


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--out", required=True, help="output directory")
    ap.add_argument("--size", type=int, nargs=2, default=[256, 256], metavar=("H", "W"))
    ap.add_argument("--variant", choices=sorted(CHECKPOINTS), default="photographic")
    ap.add_argument(
        "--golden",
        help="optional goldens.manifest.json to compare against (key: bgremove)",
    )
    args = ap.parse_args()

    model, source_files = build_model(args.variant)
    golden = _common.load_golden(args.golden, "bgremove") if args.golden else None
    _common.export_vision(
        name="modnet",
        model=model,
        example_input=torch.rand(1, 3, *args.size) * 2 - 1,
        out_dir=args.out,
        upstream={
            "model": f"ZHKKKe/MODNet ({args.variant} portrait matting checkpoint)",
            "weights": {
                **CHECKPOINTS[args.variant],
                "official": False,
                "note": "community mirror; official checkpoints are on Google Drive, linked from GitHub, neither reachable from the export host",
            },
            "architecture_source": {
                "repo": SOURCE["repo"],
                "revision": SOURCE["revision"],
                "files_sha256": source_files,
                "official": False,
                "note": "copy of the upstream src/models tree with modnet.py moved to the package root",
            },
            "license": "Apache-2.0",
        },
        extra_recipe={
            "variant": args.variant,
            "preprocessing_outside_graph": "RGB / 255, then (x - 0.5) / 0.5",
            "postprocessing_outside_graph": "alpha in [0, 1]",
        },
        golden=golden,
    )


if __name__ == "__main__":
    main()