"""Convert lama_large_512px.ckpt to a self-contained TorchScript module.

    python to_torchscript.py \
        --ckpt lama_large_512px.ckpt \
        --out anime_manga_lama.pt

The exported module has no dependency on lama_ffc.py -- consumers only need
`torch` of a compatible version (see requirements.txt):

    import torch
    m = torch.jit.load("anime_manga_lama.pt").eval()
    out = m(image, mask)            # see the contract below

Contract of the exported `forward(image, mask)`:

    image : (N, 3, H, W) float32, RGB, values in [0, 1], H and W multiples of 8
    mask  : (N, 1, H, W) float32, values in [0, 1]; >= 0.5 marks pixels to erase
    ->      (N, 3, H, W) float32, values in [0, 1]

Baked in: mask binarisation, masked-image construction, 4-channel assembly,
output clamp and the `mask * predicted + (1 - mask) * image` composite.
NOT baked in: padding to a multiple of 8 (do it before the call and crop the
result back -- see inference.py), colour-space conversion, uint8 normalisation.
"""

import argparse
import json
import time

import torch
import torch.nn as nn

from lama_ffc import BIG_LAMA_KWARGS, load_generator

SOURCE_CKPT = "dreMaz/AnimeMangaInpainting :: lama_large_512px.ckpt"
UPSTREAM = "https://github.com/advimman/lama (Apache-2.0)"


class AnimeMangaLama(nn.Module):
    """LaMa generator + baked pre/post processing. See module docstring."""

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

    def forward(self, image: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        mask = (mask >= 0.5).to(image.dtype)
        masked = image * (1.0 - mask)
        x = torch.cat([masked, mask], dim=1)
        predicted = self.generator(x).clamp(0.0, 1.0)
        return predicted * mask + image * (1.0 - mask)


def build_module(ckpt_path: str) -> AnimeMangaLama:
    gen = load_generator(ckpt_path)
    module = AnimeMangaLama(gen).eval()
    for p in module.parameters():
        p.requires_grad_(False)
    return module


@torch.inference_mode()
def export(ckpt_path: str, out_path: str) -> None:
    module = build_module(ckpt_path)

    example = (torch.rand(1, 3, 256, 256), (torch.rand(1, 1, 256, 256) > 0.5).float())
    traced = torch.jit.trace(module, example, check_trace=False)
    traced = torch.jit.freeze(traced.eval())

    validate(module, traced)

    meta = {
        "name": "anime-manga-lama-torchscript",
        "architecture": "LaMa FFCResNetGenerator (big-lama)",
        "generator_kwargs": _jsonable(BIG_LAMA_KWARGS),
        "source_checkpoint": SOURCE_CKPT,
        "upstream_architecture": UPSTREAM,
        "license": "MIT (checkpoint) / Apache-2.0 (architecture code)",
        "torch_version": torch.__version__,
        "exported_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "forward": "forward(image[N,3,H,W] f32 RGB [0,1], mask[N,1,H,W] f32 [0,1]) -> [N,3,H,W] f32 [0,1]",
        "padding": "caller pads H,W to a multiple of 8 and crops the result back",
    }
    torch.jit.save(traced, out_path, _extra_files={"meta.json": json.dumps(meta, indent=2)})
    print(f"wrote {out_path}")
    print(json.dumps(meta, indent=2))

    write_config(meta, out_path)


def write_config(meta: dict, out_path: str, path: str = "config.json") -> None:
    """A Hub-readable sidecar. Also the query file HF uses for download counts."""
    cfg = {
        "model_type": "lama",
        "architecture": "FFCResNetGenerator",
        "framework": "pytorch",
        "format": "torchscript",
        "artifact": out_path,
        "task": "image-inpainting",
        "generator_kwargs": meta["generator_kwargs"],
        "input": {
            "image": "float32 (N,3,H,W) RGB in [0,1], H and W multiples of 8",
            "mask": "float32 (N,1,H,W) in [0,1]; >= 0.5 marks pixels to erase",
        },
        "output": "float32 (N,3,H,W) RGB in [0,1] (mask*predicted + (1-mask)*image)",
        "baked_in": ["mask binarisation", "masked-image assembly", "4-channel concat",
                     "output clamp", "composite"],
        "not_baked_in": ["padding to multiple of 8", "colour-space conversion",
                         "uint8 normalisation"],
        "source_checkpoint": meta["source_checkpoint"],
        "source_checkpoint_repo": "dreMaz/AnimeMangaInpainting",
        "upstream_architecture": meta["upstream_architecture"],
        "exported_with_torch": meta["torch_version"],
        "license": "MIT (checkpoint) / Apache-2.0 (architecture code)",
    }
    with open(path, "w") as f:
        f.write(json.dumps(cfg, indent=2) + "\n")
    print(f"wrote {path}")


def _jsonable(d):
    return {k: (_jsonable(v) if isinstance(v, dict) else v) for k, v in d.items()}


@torch.inference_mode()
def validate(eager: nn.Module, scripted: nn.Module) -> None:
    torch.manual_seed(0)
    cases = [(1, 256, 256), (1, 512, 512), (2, 512, 384), (1, 1024, 768), (4, 320, 320)]
    worst = 0.0
    for n, h, w in cases:
        img = torch.rand(n, 3, h, w)
        mask = (torch.rand(n, 1, h, w) > 0.7).float()
        a = eager(img, mask)
        b = scripted(img, mask)
        d = (a - b).abs().max().item()
        worst = max(worst, d)
        assert a.shape == (n, 3, h, w)
        assert b.min() >= 0.0 and b.max() <= 1.0
        print(f"  case N{n} {h}x{w}: max|eager-scripted| = {d:.2e}")
    assert worst < 1e-3, f"scripted output diverges from eager (max {worst:.2e})"
    print(f"  OK (worst {worst:.2e})")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="lama_large_512px.ckpt")
    ap.add_argument("--out", default="anime_manga_lama.pt")
    args = ap.parse_args()
    export(args.ckpt, args.out)