#!/usr/bin/env python3
# coding=utf-8
"""Convert a ConCor-1 training checkpoint into this Hugging Face model repository.

The research training code stores a checkpoint directory

    step{N}/model.pt          # torch state dict of the training model
    step{N}/train/rank0.pt    # training config + bridge token ids

whose module names follow the internal research code. This script renames them
to the released module names (which follow the paper) and writes a standard HF
repository: `config.json`, `model.safetensors`, tokenizer and processor files.

Renaming performed (research name -> released name):

    backbone.model.*                 -> backbone.*            (`lm_head` is dropped: unused)
    fine_patch_decoder.branches.*    -> vision_segmentation_head.patch_expander.branches.*
    subpatch_decoder.blocks.*        -> vision_segmentation_head.convolutional_decoder.blocks.*
    fine_seg_head.vis_*              -> vision_segmentation_head.mask_predictor.feature_proj.*
    fine_seg_head.bridge_*           -> vision_segmentation_head.mask_predictor.bridge_proj.*
    fine_seg_head.bilinear           -> vision_segmentation_head.mask_predictor.bilinear
    text_head.vis_*                  -> text_segmentation_head.feature_proj.*
    text_head.bridge_*               -> text_segmentation_head.bridge_proj.*
    text_head.bilinear               -> text_segmentation_head.bilinear
    presence_head.*                  -> presence_head.*        (unchanged)

Usage:
    python convert_original_checkpoint.py \
        --checkpoint /path/to/step100000 \
        --qwen3_5_path /path/to/Qwen3.5-0.8B \
        --output_dir .
"""

from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path
from typing import Dict

import torch
from safetensors.torch import save_file
from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config

from configuration_concor1 import ConCor1Config

# Architecture flags of the research config that the released code implements.
# Anything else means the checkpoint is a different ablation and must not be
# silently converted.
EXPECTED_TRAINING_FLAGS = {
    "query_text_grounding": True,
    "presence_loss_type": "raw",
    "presence_mlp": True,
    "fine_patch_decoding": True,
    "use_subpatch_decoder": True,
    "use_pixel_decoder": False,
    "fine_decoder_swiglu": True,
    "add_vit_features": True,
    "add_where": "decoder_output",
    "vit_feature_layers": None,
    "multi_seg_loss": False,
    "bridge_self_attention": False,
    "bbox_bridge_encoding": False,
    "bidirectional_full_attention": True,
    "asymmetric_bridge_attention": False,
    "num_register_tokens": 0,
    "dot_product_similarity": False,
    "reset_position": False,
    "spatial_grid_assignment": False,
    "patch_size": 16,
    "merge_size": 2,
    "grounding_hidden_dim": 256,
}

TOKENIZER_FILES = [
    "tokenizer.json",
    "tokenizer_config.json",
    "vocab.json",
    "merges.txt",
]

_PROJECTION_SUFFIX = {
    "gate_proj.weight": "gate_proj.weight",
    "up_proj.weight": "up_proj.weight",
    "norm.weight": "norm.weight",
}


def rename_scorer_key(key: str, prefix: str, new_prefix: str) -> str:
    """Map a research bilinear-head key onto the released scorer layout."""
    remainder = key[len(prefix) + 1 :]
    if remainder == "bilinear.weight":
        return f"{new_prefix}.bilinear.weight"
    for side, new_side in (("vis", "feature_proj"), ("bridge", "bridge_proj")):
        for old, new in (
            (f"{side}_gate_proj.weight", "gate_proj.weight"),
            (f"{side}_up_proj.weight", "up_proj.weight"),
            (f"{side}_norm.weight", "norm.weight"),
        ):
            if remainder == old:
                return f"{new_prefix}.{new_side}.{new}"
    raise KeyError(f"Unexpected key in {prefix}: {key}")


def convert_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """Rename research module names to the released ones."""
    converted: Dict[str, torch.Tensor] = {}
    dropped = []

    for key, tensor in state_dict.items():
        if key == "backbone.lm_head.weight":
            # ConCor-1 never decodes text, so the tied LM head is not part of the
            # released model (it is also what the paper's parameter count reflects).
            dropped.append(key)
            continue
        if key.startswith("subpatch_decoder.block1.") or key.startswith("subpatch_decoder.block2."):
            # The research module aliases its two blocks both as `block{1,2}` and as
            # `blocks.{0,1}`; keep one copy.
            dropped.append(key)
            continue

        if key.startswith("backbone.model."):
            new_key = "backbone." + key[len("backbone.model.") :]
        elif key.startswith("fine_patch_decoder.branches."):
            new_key = "vision_segmentation_head.patch_expander." + key[len("fine_patch_decoder.") :]
        elif key.startswith("subpatch_decoder.blocks."):
            new_key = (
                "vision_segmentation_head.convolutional_decoder." + key[len("subpatch_decoder.") :]
            )
        elif key.startswith("fine_seg_head."):
            new_key = rename_scorer_key(
                key, "fine_seg_head", "vision_segmentation_head.mask_predictor"
            )
        elif key.startswith("text_head."):
            new_key = rename_scorer_key(key, "text_head", "text_segmentation_head")
        elif key.startswith("presence_head."):
            new_key = key
        else:
            raise KeyError(f"Unhandled checkpoint key: {key}")

        converted[new_key] = tensor

    print(f"Renamed {len(converted)} tensors, dropped {len(dropped)}: {sorted(dropped)[:4]} ...")
    return converted


def build_config(training_config: dict, bridge_token_ids: list, qwen3_5_path: Path, vocab_size: int) -> ConCor1Config:
    """Build the released config from the training config + the Qwen3.5 architecture."""
    for field, expected in EXPECTED_TRAINING_FLAGS.items():
        actual = training_config.get(field, "<missing>")
        if actual != expected:
            raise ValueError(
                f"Training config field {field!r} is {actual!r}, expected {expected!r}. "
                "This checkpoint is not the architecture the released code implements."
            )

    backbone_config = Qwen3_5Config.from_pretrained(qwen3_5_path)
    # The embedding table was grown to host the bridge-token ids.
    backbone_config.vocab_size = vocab_size
    backbone_config.text_config.vocab_size = vocab_size
    # The backbone is used as an encoder, so it is not a Qwen3.5 generator anymore.
    backbone_config.architectures = None

    grid_levels = training_config["multiscale_bridge_levels"]
    num_bridge_tokens = training_config["num_bridges"]
    if grid_levels is None or sum(level ** 2 for level in grid_levels) != num_bridge_tokens:
        raise ValueError(
            f"multiscale_bridge_levels={grid_levels} is inconsistent with num_bridges={num_bridge_tokens}."
        )
    if bridge_token_ids != list(range(bridge_token_ids[0], bridge_token_ids[0] + num_bridge_tokens)):
        raise ValueError("Bridge token ids are expected to be a contiguous range.")

    return ConCor1Config(
        backbone_config=backbone_config,
        num_bridge_tokens=num_bridge_tokens,
        bridge_grid_levels=grid_levels,
        bridge_token_id_start=bridge_token_ids[0],
        correspondence_dim=training_config["grounding_hidden_dim"],
        presence_hidden_dim=training_config["grounding_hidden_dim"],
        patch_size=training_config["patch_size"],
        merge_size=training_config["merge_size"],
        num_mask_upsample_blocks=2,  # SubpatchDecoder: two 2x blocks -> 4 px cells
        fuse_vision_encoder_features=training_config["add_vit_features"],
        bidirectional_full_attention=training_config["bidirectional_full_attention"],
        image_min_pixels=training_config["min_pixels"],
        image_max_pixels=training_config["max_pixels"],
        dtype="bfloat16",
        architectures=["ConCor1ForConceptCorrespondence"],
        auto_map={
            "AutoConfig": "configuration_concor1.ConCor1Config",
            "AutoModel": "modeling_concor1.ConCor1ForConceptCorrespondence",
            "AutoProcessor": "processing_concor1.ConCor1Processor",
        },
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", required=True, type=Path, help="Training checkpoint directory (step{N})")
    parser.add_argument("--qwen3_5_path", required=True, type=Path, help="Qwen3.5-0.8B directory (config + tokenizer)")
    parser.add_argument("--output_dir", default=Path("."), type=Path)
    parser.add_argument(
        "--dtype",
        default="keep",
        choices=["keep", "bfloat16", "float32"],
        help=(
            "dtype of the released weights. 'keep' (default) preserves the checkpoint exactly — "
            "bf16 backbone, fp32 prediction heads — which is what the released model reproduces "
            "via `_keep_in_fp32_modules_strict`. 'bfloat16' halves the head weights' precision "
            "(image masks then agree with the checkpoint at ~0.995 IoU)."
        ),
    )
    args = parser.parse_args()

    train_state_path = args.checkpoint / "train" / "rank0.pt"
    if not train_state_path.exists():
        train_state_path = args.checkpoint / "train_state.pt"
    train_state = torch.load(train_state_path, map_location="cpu", weights_only=False)
    training_config = train_state["config"]
    bridge_token_ids = train_state["bridge_token_ids"]
    print(
        f"Loaded training state from {train_state_path} "
        f"(step {train_state.get('global_step')}, {len(bridge_token_ids)} bridge tokens)"
    )

    state_dict = torch.load(args.checkpoint / "model.pt", map_location="cpu", weights_only=True)
    vocab_size = state_dict["backbone.model.language_model.embed_tokens.weight"].shape[0]
    converted = convert_state_dict(state_dict)

    if args.dtype != "keep":
        target_dtype = getattr(torch, args.dtype)
        converted = {
            key: tensor.to(target_dtype) if tensor.is_floating_point() else tensor
            for key, tensor in converted.items()
        }

    config = build_config(training_config, bridge_token_ids, args.qwen3_5_path, vocab_size)

    args.output_dir.mkdir(parents=True, exist_ok=True)
    config.save_pretrained(args.output_dir)
    save_file(
        {key: tensor.contiguous() for key, tensor in converted.items()},
        args.output_dir / "model.safetensors",
        metadata={"format": "pt"},
    )
    total_parameters = sum(tensor.numel() for tensor in converted.values())
    print(f"Wrote model.safetensors — {len(converted)} tensors, {total_parameters / 1e6:.2f} M parameters")

    # Tokenizer: Qwen3.5's, unchanged. Bridge tokens are unused vocabulary slots
    # with no string form, so the tokenizer needs no new entries.
    for filename in TOKENIZER_FILES:
        source = args.qwen3_5_path / filename
        if source.exists():
            shutil.copy2(source, args.output_dir / filename)
    tokenizer_config_path = args.output_dir / "tokenizer_config.json"
    tokenizer_config = json.loads(tokenizer_config_path.read_text())
    tokenizer_config["processor_class"] = "ConCor1Processor"
    tokenizer_config_path.write_text(json.dumps(tokenizer_config, indent=2, ensure_ascii=False) + "\n")

    # Image processor: Qwen3.5's, with ConCor-1's fixed pixel budget.
    image_processor_config = json.loads((args.qwen3_5_path / "preprocessor_config.json").read_text())
    image_processor_config["processor_class"] = "ConCor1Processor"
    image_processor_config["min_pixels"] = config.image_min_pixels
    image_processor_config["max_pixels"] = config.image_max_pixels
    (args.output_dir / "preprocessor_config.json").write_text(
        json.dumps(image_processor_config, indent=2) + "\n"
    )

    processor_config = {
        "processor_class": "ConCor1Processor",
        "auto_map": {"AutoProcessor": "processing_concor1.ConCor1Processor"},
        "num_bridge_tokens": config.num_bridge_tokens,
        "bridge_token_id_start": config.bridge_token_id_start,
        "patch_size": config.patch_size,
        "merge_size": config.merge_size,
        "num_mask_upsample_blocks": config.num_mask_upsample_blocks,
        "min_pixels": config.image_min_pixels,
        "max_pixels": config.image_max_pixels,
        "presence_threshold": config.presence_threshold,
        "text_threshold": config.text_threshold,
        "image_threshold": config.image_threshold,
        "nms_iou_threshold": config.nms_iou_threshold,
    }
    (args.output_dir / "processor_config.json").write_text(
        json.dumps(processor_config, indent=2) + "\n"
    )
    print(f"Wrote config.json, processor_config.json, preprocessor_config.json and tokenizer files to {args.output_dir}")


if __name__ == "__main__":
    main()