import torch
import torch
from collections import OrderedDict

def load_t5_from_3dllm_ckpt_V2(ckpt_path, t5, device="cpu"):
    obj = torch.load(ckpt_path, map_location=device, weights_only=False)
    sd = obj.get("model", obj)

    def fit_vocab(src, tgt_shape):
        if list(src.shape) == list(tgt_shape):
            return src
        Vsrc, D = src.shape
        Vtgt, D2 = tgt_shape
        assert D == D2, f"embed dim mismatch {D} vs {D2}"
        if Vsrc > Vtgt:
            raise ValueError(f"Checkpoint vocab {Vsrc} > model vocab {Vtgt}")
        pad = torch.zeros((Vtgt - Vsrc, D), dtype=src.dtype, device=src.device)
        return torch.cat([src, pad], dim=0)

    new_sd = OrderedDict()
    for name, param in t5.state_dict().items():
        ckpt_name = name if name in sd else f"t5_model.{name}"
        if ckpt_name not in sd:
            continue
        w = sd[ckpt_name]
        if "embed_tokens.weight" in name or "shared.weight" in name or "lm_head.weight" in name:
            w = fit_vocab(w, param.shape)
        new_sd[name] = w

    t5.load_state_dict(new_sd, strict=False)
    t5.tie_weights()
    return t5


def load_t5_from_3dllm_ckpt(ckpt_path, t5, device="cpu"):
    
    obj = torch.load(ckpt_path, map_location=device, weights_only=False)
    sd = obj.get("model", obj)

    def take(key):
        # 3d llm keys start with : 't5_model.xxx'
        k = key if key in sd else f"t5_model.{key}"
        return sd[k]
    def fit_vocab(src, tgt_shape):  # (Vsrc, D) -> (Vtgt, D)
        if list(src.shape) == list(tgt_shape):
            return src
        Vsrc, D = src.shape
        Vtgt, D2 = tgt_shape
        assert D == D2, f"embed dim mismatch {D} vs {D2}"
        if Vsrc >= Vtgt:  #
            raise Exception('the source(checkpoint) must have smaller vocab size')
        pad = torch.zeros((Vtgt - Vsrc, D), dtype=src.dtype)
        return torch.cat([src, pad], dim=0)

    with torch.no_grad():
        # shared / embed_tokens / lm_head
        sw = take("shared.weight")
        t5.shared.weight.copy_(fit_vocab(sw, t5.shared.weight.shape))

        ew = take("encoder.embed_tokens.weight")
        t5.encoder.embed_tokens.weight.copy_(fit_vocab(ew, t5.encoder.embed_tokens.weight.shape))

        dw = take("decoder.embed_tokens.weight")
        t5.decoder.embed_tokens.weight.copy_(fit_vocab(dw, t5.decoder.embed_tokens.weight.shape))

        lw = take("lm_head.weight")
        t5.lm_head.weight.copy_(fit_vocab(lw, t5.lm_head.weight.shape))

    t5.tie_weights()
    return t5
    
def load_pretrained(cfg, model, logger=None, phase="train"):
    if logger is not None:
        logger.info(f"Loading pretrain model from {cfg.TRAIN.PRETRAINED}")

    if phase == "train":
        ckpt_path = cfg.TRAIN.PRETRAINED
    elif phase == "test":
        ckpt_path = cfg.TEST.CHECKPOINTS

    state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False)["state_dict"]

    ignore_prefixes = []

    # dangerous modification
    if cfg.TRAIN.get('LOAD_CKPT_WO_TP_ENC', False):
        ignore_prefixes.append("three_points_proj.")

    # GRPO validation: exclude ref_model weights from checkpoint
    if cfg.TRAIN.get('GRPO_VAL', False):
        ignore_prefixes.append("ref_model.")

    if ignore_prefixes:
        state_dict = {k: v for k, v in state_dict.items()
                      if not any(k.startswith(p) for p in ignore_prefixes)}
        if logger is not None:
            logger.info(f"Filtered out keys with prefixes: {ignore_prefixes}")
        model.load_state_dict(state_dict, strict=False)
    else:
        model.load_state_dict(state_dict, strict=False)
    return model


def load_pretrained_vae(cfg, model, logger=None):
    state_dict = torch.load(cfg.TRAIN.PRETRAINED_VAE,
                            map_location="cpu")['state_dict']
    if logger is not None:
        logger.info(f"Loading pretrain vae from {cfg.TRAIN.PRETRAINED_VAE}")
        
    # Extract encoder/decoder
    from collections import OrderedDict
    vae_dict = OrderedDict()
    for k, v in state_dict.items():
        if "motion_vae" in k:
            name = k.replace("motion_vae.", "")
            vae_dict[name] = v
        elif "vae" in k:
            name = k.replace("vae.", "")
            vae_dict[name] = v
    if hasattr(model, 'vae'):
        model.vae.load_state_dict(vae_dict, strict=True)
    else:
        model.motion_vae.load_state_dict(vae_dict, strict=True)
    
    return model


def load_pretrained_grpo_ref(cfg, model, logger=None):
    """Load GRPO reference model from checkpoint.

    Reference model is co-located with the main model on each rank's device.
    The legacy ``GRPO_REF_DEVICE`` setting is intentionally ignored: with DDP,
    Lightning recursively moves child submodules to the per-rank device, so any
    fixed device string would go stale and trigger cross-device tensor moves.
    """
    ref_path = cfg.TRAIN.GRPO_REF_MODEL_PATH
    ref_device = model.device
    if logger is not None:
        logger.info(f"Loading GRPO reference model from {ref_path}")
        logger.info(f"Reference model will be placed on {ref_device} (co-located with main model)")

    from mGPT.models.build_model import build_model

    # Create ref model with same architecture as training model
    ref_model = build_model(cfg, model.datamodule)

    # Add same special tokens as training model
    if cfg.model.params.get("dummy_motion", False):
        ref_model.lm.init_dummy_motion_token(
            cfg.model.params.p_dummy,
            cfg.model.params.per_token,
            cfg.model.params.get('use_single_span', False),
            cfg.model.params.get('span_pairs', 6),
            cfg.model.params.get('p_span', 1.0)
        )
    if cfg.model.params.get("tp_special_token", False):
        ref_model.lm.init_tp_special_token()
    if cfg.LOSS.ABLATION.get("TRAIN_ALL_AT_ONCE", False):
        ref_model.lm.init_text_special_token()
    if cfg.TRAIN.get("CHAIN_OF_THOUGHT", False):
        ref_model.lm.init_cot_special_token()

    # Load state dict (vocab size now matches checkpoint)
    state_dict = torch.load(ref_path, map_location="cpu", weights_only=False)["state_dict"]
    # If the ckpt itself was a GRPO save (contains its own nested ref_model.*
    # subtree), strip that prefix — those keys would otherwise be unexpected
    # in a fresh MotionGPT (= ref_model) which has no nested ref_model. Same
    # for three_points_proj when LOAD_CKPT_WO_TP_ENC is set.
    ignore_prefixes = ["ref_model."]
    if cfg.TRAIN.get('LOAD_CKPT_WO_TP_ENC', False):
        ignore_prefixes.append("three_points_proj.")
    state_dict = {k: v for k, v in state_dict.items()
                  if not any(k.startswith(p) for p in ignore_prefixes)}
    # strict=False to also tolerate any benign add/remove (e.g. peft wrapping
    # differences when the ckpt was full-FT but config wraps with LoRA).
    ref_model.load_state_dict(state_dict, strict=False)

    ref_model = ref_model.to(ref_device)
    ref_model.eval()

    # Freeze all parameters
    for param in ref_model.parameters():
        param.requires_grad = False

    model.ref_model = ref_model

    return model