# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Adapter that loads the upstream HaMeR Lightning checkpoint into the
self-contained torch reference shipped in :mod:`hamer.reference.hamer`.

The upstream ``mano_head.transformer.transformer.layers.{i}.{0|1|2}.{norm|fn}``
naming (and the dropout-padded ``Sequential`` indices) are remapped onto the
reference module's flatter ``head.transformer.layers[i].{norm_sa|sa|...}``
layout.  No new dependencies — pure ``torch`` + ``numpy``.
"""
from __future__ import annotations

import re
from pathlib import Path
from typing import Dict

import numpy as np
import torch


def _remap_key(k: str) -> str | None:
    """Translate one upstream key to the reference module's name; return
    ``None`` for keys we don't model (e.g. discriminator)."""
    if k.startswith("backbone."):
        return k
    m = re.match(r"mano_head\.transformer\.transformer\.layers\.(\d+)\.(\d+)\.(.+)", k)
    if m:
        block, sub, rest = int(m.group(1)), int(m.group(2)), m.group(3)
        slot = ("sa", "ca", "ff")[sub]
        if rest.startswith("norm."):
            slot_norm = ("norm_sa", "norm_ca", "norm_ff")[sub]
            return f"head.transformer.layers.{block}.{slot_norm}." + rest[len("norm.") :]
        if rest.startswith("fn."):
            inner = rest[len("fn.") :]
            # FFN: net.0/net.3 (with dropouts) → net.0/net.2 (no dropouts)
            if slot == "ff":
                inner = inner.replace("net.3.", "net.2.")
            # Linear-in-Sequential out projection: to_out.0.* → to_out.*
            inner = inner.replace("to_out.0.", "to_out.")
            return f"head.transformer.layers.{block}.{slot}." + inner
        return None
    if k.startswith("mano_head.transformer.to_token_embedding."):
        return "head.transformer.to_token_embedding." + k.split(".", 3)[3]
    if k == "mano_head.transformer.pos_embedding":
        return "head.transformer.pos_embedding"
    if k.startswith("mano_head.decpose.") or k.startswith("mano_head.decshape.") or k.startswith("mano_head.deccam."):
        return "head." + k.split(".", 1)[1]
    if k.startswith("mano_head.init_"):
        return "head." + k.split(".", 1)[1]
    return None


def load_hamer_checkpoint(
    ref_module,
    ckpt_path: Path | str,
    mean_params_path: Path | str | None = None,
    weights_only: bool = True,
) -> dict:
    """Load the official HaMeR checkpoint into the reference ``Hamer`` module.

    Returns a dict reporting how many tensors were mapped, missing, and extra.
    Also overwrites the ``init_*`` buffers with the upstream MANO mean params
    if ``mean_params_path`` is provided.

    ``weights_only=True`` (default) refuses arbitrary pickle globals.  The
    released ``hamer.ckpt`` only references ``collections.OrderedDict`` and
    torch storage/rebuild helpers, so it loads under the safe loader; pass
    ``weights_only=False`` only for a checkpoint you trust that does not.
    """
    ckpt = torch.load(str(ckpt_path), map_location="cpu", weights_only=weights_only)
    src = ckpt["state_dict"]

    remapped: Dict[str, torch.Tensor] = {}
    skipped = []
    for k, v in src.items():
        new_k = _remap_key(k)
        if new_k is None:
            skipped.append(k)
            continue
        remapped[new_k] = v

    missing, unexpected = ref_module.load_state_dict(remapped, strict=False)

    if mean_params_path is not None:
        mp = np.load(str(mean_params_path))
        with torch.no_grad():
            ref_module.head.init_hand_pose.copy_(torch.from_numpy(mp["pose"].astype("float32")).unsqueeze(0))
            ref_module.head.init_betas.copy_(torch.from_numpy(mp["shape"].astype("float32")).unsqueeze(0))
            ref_module.head.init_cam.copy_(torch.from_numpy(mp["cam"].astype("float32")).unsqueeze(0))

    return {
        "loaded": len(remapped),
        "skipped_upstream": len(skipped),
        "missing_in_ref": len(missing),
        "unexpected_in_ref": len(unexpected),
        "missing_examples": missing[:5],
        "unexpected_examples": unexpected[:5],
    }