#!/usr/bin/env python3
"""Export DDColor paper_tiny (Apache-2.0, piddnad/DDColor) to the OptiMix ONNX contract.

Weights: huggingface.co/piddnad/ddcolor_paper_tiny @ cf9fd99c1d7472689ec7413441c1b799a51866a3, pytorch_model.bin
         sha256 8a1277bc90a1bfbb6d2d83933a9a6bc821931879ca93e26e4fcec12165d41fce
Code:    github.com/piddnad/DDColor @ 2adb63f2656ac41cbdf7b894cddd94121a3faf13 (basicsr/archs/ddcolor_arch.py)
Contract: input  `input`  float32 [1,3,512,512] RGB 0..1, the L channel rendered as grey RGB
          output `output` float32 [1,2,512,512] Lab a,b (OpenCV float Lab scale)
"""
import hashlib, json, sys
from pathlib import Path
import cv2, numpy as np, onnx, onnxruntime as ort, torch

# Usage: export_ddcolor_onnx.py DDCOLOR_REPO_DIR WEIGHTS_DIR OUT.onnx TEST_IMAGE [PREVIEW.jpg]
#   DDCOLOR_REPO_DIR: a checkout of github.com/piddnad/DDColor (commit above)
#   WEIGHTS_DIR: holds config.json and pytorch_model.bin from piddnad/ddcolor_paper_tiny (commit above)
REPO = Path(sys.argv[1]).resolve(); sys.path.insert(0, str(REPO))
from basicsr.archs.ddcolor_arch import DDColor

WEIGHTS = Path(sys.argv[2])
OUT = Path(sys.argv[3])
TEST_IMAGE = sys.argv[4]
PREVIEW = sys.argv[5] if len(sys.argv) > 5 else None
OPSET = 17
cfg = json.loads((WEIGHTS / "config.json").read_text())
model = DDColor(**cfg)
state = torch.load(WEIGHTS / "pytorch_model.bin", map_location="cpu", weights_only=True)
state = state.get("params", state)
missing, unexpected = model.load_state_dict(state, strict=False)
print("missing", missing, "unexpected", unexpected)
assert not unexpected, unexpected
assert all(k.endswith("num_batches_tracked") for k in missing), missing
model.eval()

# Reference = the official model exactly as loaded (spectral-norm hooks intact).
img = cv2.imread(TEST_IMAGE)

def preprocess(img_bgr):
    img = (img_bgr / 255.0).astype(np.float32)
    r = cv2.resize(img, (512, 512))
    l = cv2.cvtColor(r, cv2.COLOR_BGR2Lab)[:, :, :1]
    lab = np.concatenate((l, np.zeros_like(l), np.zeros_like(l)), axis=-1)
    rgb = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
    return rgb.transpose(2, 0, 1)[None].astype(np.float32)

x = preprocess(img)
with torch.no_grad():
    ref = model(torch.from_numpy(x)).numpy()

# Bake spectral norm into plain weights (W / sigma with the stored u, v; eval mode
# runs no power iteration), so the graph does not carry weight_orig plus a reshaped copy.
from torch.nn.utils import remove_spectral_norm
baked = 0
for mod in model.modules():
    if hasattr(mod, "weight_orig"):
        remove_spectral_norm(mod); baked += 1
print("spectral norms baked", baked)
with torch.no_grad():
    torch.onnx.export(model, torch.rand(1, 3, 512, 512), str(OUT), opset_version=OPSET,
                      input_names=["input"], output_names=["output"], dynamo=False,
                      do_constant_folding=True)
m = onnx.load(str(OUT)); onnx.checker.check_model(m)
print("opset", [(o.domain, o.version) for o in m.opset_import])
print("inputs", [(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim]) for i in m.graph.input])
print("outputs", [(o.name, [d.dim_value for d in o.type.tensor_type.shape.dim]) for o in m.graph.output])

with torch.no_grad():
    baked_ref = model(torch.from_numpy(x)).numpy()
print("baked vs official torch max abs diff", float(np.abs(baked_ref - ref).max()))
sess = ort.InferenceSession(str(OUT), providers=["CPUExecutionProvider"])
got = sess.run(["output"], {"input": x})[0]
print("shape", got.shape, "ref range", float(ref.min()), float(ref.max()))
print("max abs diff", float(np.abs(got - ref).max()), "mean abs diff", float(np.abs(got - ref).mean()))
# Colourise the test image with the ONNX output, as DDColor's pipeline does.
h, w = img.shape[:2]
orig_l = cv2.cvtColor((img / 255.0).astype(np.float32), cv2.COLOR_BGR2Lab)[:, :, :1]
ab = cv2.resize(got[0].transpose(1, 2, 0), (w, h), interpolation=cv2.INTER_NEAREST)
out = cv2.cvtColor(np.concatenate((orig_l, ab), -1), cv2.COLOR_LAB2BGR)
if PREVIEW:
    cv2.imwrite(PREVIEW, (out * 255).round().clip(0, 255).astype(np.uint8))
print("size", OUT.stat().st_size, "sha256", hashlib.sha256(OUT.read_bytes()).hexdigest())