"""Run Rockchip's exporter using its shipped CLIP embeddings for the dry run."""
import os
from pathlib import Path
import runpy
import sys
from tempfile import NamedTemporaryFile

import numpy as np
import torch
import mmdet.apis
from mmengine.config import Config

root = Path(__file__).resolve().parent
upstream = root / "upstream"
zoo = root.parent.parent.parent / "rknn_model_zoo"
assets = zoo / "examples/yolo_world/model"
embeddings = torch.from_numpy(np.load(assets / "coco_text_outp.npy")).reshape(1, 80, 512)
config_name = "yolo_world_v2_m_vlpan_bn_noeinsum_2e-3_100e_4x8gpus_obj365v1_goldg_train_lvis_minival.py"
output = root / "onnx/yolo_world"
output.mkdir(parents=True, exist_ok=True)
os.chdir(upstream)
sys.path[:0] = [str(upstream / "deploy"), str(upstream)]

# The detector takes text embeddings as a separate input. The official dry run
# needs their shape, so reuse Rockchip's COCO embeddings and skip fetching CLIP.
init_detector = mmdet.apis.init_detector


def init_with_local_text(config_path, checkpoint, device):
    source = Path(config_path)
    fallback = root / "mmyolo-source"
    base = "configs/yolov8/yolov8_m_syncbn_fast_8xb16-500e_coco.py"
    if (upstream / "third_party/mmyolo" / base).is_file():
        config = Config.fromfile(source)
    elif (fallback / base).is_file():
        with NamedTemporaryFile(mode="w", suffix=".py", dir=source.parent) as copy:
            copy.write(source.read_text().replace("../../third_party/mmyolo/", f"{fallback}/"))
            copy.flush()
            config = Config.fromfile(copy.name)
    else:
        raise FileNotFoundError("Missing the pinned MMYOLO configuration")
    config.model.backbone.with_text_model = False
    # init_detector reads dataset metadata even though export uses no images
    # or labels from the evaluation dataset.
    config.test_dataloader.dataset = dict(type="mmdet.CocoDataset", lazy_init=True)
    model = init_detector(config, checkpoint, device=device)
    model.reparameterize = lambda texts: setattr(model, "text_feats", embeddings.to(device))
    return model


mmdet.apis.init_detector = init_with_local_text
script = str(upstream / "deploy/export_onnx.py")
sys.argv = [script, str(upstream / "configs/pretrain" / config_name),
            str(root / "weights/yolo_world_v2m.pth"), "--custom-text",
            str(upstream / "data/texts/coco_class_texts.json"), "--opset", "11",
            "--model-only", "--work-dir", str(output)]
runpy.run_path(script, run_name="__main__")