"""Split a JointRadiomapCIRModel ckpt into two single-task ckpts suitable for
existing V2.2a / D2.2 visualization scripts.

Usage:
    python split_ckpt_for_viz.py \
        --joint-ckpt <path>/best.pt \
        --out-radiomap-ckpt <path>/best_rm.pt \
        --out-csi-ckpt <path>/best_csi.pt
"""
from __future__ import annotations

import argparse
from pathlib import Path

import torch


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--joint-ckpt", required=True)
    p.add_argument("--out-radiomap-ckpt", required=True)
    p.add_argument("--out-csi-ckpt", required=True)
    args = p.parse_args()

    ckpt = torch.load(args.joint_ckpt, map_location="cpu", weights_only=False)
    sd = ckpt["model_state_dict"]

    # --- Split 1: radiomap-only (compatible with RadiomapModelV15 = V2.2a ckpt format) ---
    # RadiomapModelV15 keys: backbone.*, tx_proj.*, tx_embed.freqs,
    #   txrx_query_encoder.*, ray_gather.*, cross_blocks.*, self_blocks.*,
    #   down_proj.*, spatial_stem.*, spatial_output.*
    rm_sd = {k: v for k, v in sd.items() if not k.startswith("csi_head.")}
    rm_ckpt = {
        "epoch": ckpt.get("epoch", 0),
        "step": ckpt.get("step", 0),
        "tag": "split_from_joint",
        "metrics": ckpt.get("metrics", {}),
        "cli_args": ckpt.get("cli_args", {}),
        "model_state_dict": rm_sd,
    }
    Path(args.out_radiomap_ckpt).parent.mkdir(parents=True, exist_ok=True)
    torch.save(rm_ckpt, args.out_radiomap_ckpt)
    print(f"✓ Radiomap ckpt: {len(rm_sd)} keys → {args.out_radiomap_ckpt}")

    # --- Split 2: CIR-only (compatible with SparseCsiDetrModel = D2.2 ckpt format) ---
    # SparseCsiDetrModel keys: backbone.*, tx_proj.*, tx_embed.freqs, head.*
    csi_sd = {}
    for k, v in sd.items():
        if k.startswith("csi_head."):
            csi_sd[k.replace("csi_head.", "head.", 1)] = v
        elif k.startswith("backbone.") or k.startswith("tx_proj.") or k.startswith("tx_embed."):
            csi_sd[k] = v
        # drop radiomap-only keys
    # Infer model_config for D2.2-style loader (visualize_checkpoint.py reads 'model_config')
    # Detect num_decoder_layers and head arch from csi state_dict
    num_layers = 0
    for k in csi_sd:
        if k.startswith("head.decoder.layers."):
            idx = int(k.split(".")[3])
            num_layers = max(num_layers, idx + 1)
    csi_ckpt = {
        "epoch": ckpt.get("epoch", 0),
        "step": ckpt.get("step", 0),
        "tag": "split_from_joint",
        "metrics": ckpt.get("metrics", {}),
        "cli_args": ckpt.get("cli_args", {}),
        "model_state_dict": csi_sd,
        "model_config": {
            "channels": 512,
            "num_queries": 8,
            "num_decoder_layers": num_layers,
            "db_low": -115.0,
            "db_high": -7.5,
            "delay_max_ns": 15.0,
            "backbone_kind": "trellis2",
        },
        "trainer_config": {
            "matcher_backend": "scipy",
            "no_object_exists_weight": 5.0,
        },
    }
    Path(args.out_csi_ckpt).parent.mkdir(parents=True, exist_ok=True)
    torch.save(csi_ckpt, args.out_csi_ckpt)
    print(f"✓ CIR ckpt: {len(csi_sd)} keys (head.* renamed from csi_head.*) → {args.out_csi_ckpt}")


if __name__ == "__main__":
    main()