"""Load a released checkpoint into the model defined by the code release's methods.py.

Place this file in the root directory of the code release, next to methods.py.
"""
import re

import torch
from huggingface_hub import hf_hub_download

import methods

REPO_ID = "headless-start/parameter-efficient-dfer"
CLASS_NAMES = ["angry", "disgust", "fear", "happy", "sad", "surprise"]


def _rename(name):
    # The training runs used older parameter names than the code release.
    name = re.sub(r"\.ssf_gamma$", ".gamma", name)
    name = re.sub(r"\.ssf_beta$", ".beta", name)
    for old, new in ((".lora_q_A.", ".q_a."), (".lora_q_B.", ".q_b."),
                     (".lora_v_A.", ".v_a."), (".lora_v_B.", ".v_b."),
                     (".adaptmlp.", ".adapter.")):
        name = name.replace(old, new)
    return name


def _state(path):
    return torch.load(path, map_location="cpu", weights_only=False)["model"]


def load_stage1():
    model = methods.build_model(pretrained=False)
    model.load_state_dict(_state(hf_hub_download(REPO_ID, "stage1/best.ckpt")))
    return model.eval()


def load_stage2(method, fold):
    """method: full_ft, linear_probe, ssf, lora or adaptformer; fold: 0-9."""
    file = f"stage2/{method}/fold_{fold:02d}/best.ckpt"
    state = {_rename(k): v for k, v in _state(hf_hub_download(REPO_ID, file)).items()}
    model = methods.build_model(pretrained=False)
    if method != "full_ft":
        # Adaptation weights sit on top of the frozen Stage-1 encoder.
        backbone = _state(hf_hub_download(REPO_ID, "stage1/best.ckpt"))
        model.load_state_dict({k: v for k, v in backbone.items()
                               if not k.startswith("head.")}, strict=False)
    methods.configure(model, method, rank=4, reduction=12)
    _, unexpected = model.load_state_dict(state, strict=method == "full_ft")
    if unexpected:
        raise ValueError(f"unexpected keys in {file}: {unexpected[:3]}")
    return model.eval()


if __name__ == "__main__":
    import sys
    from PIL import Image
    from data import build_transform

    model = load_stage2("lora", 0)
    image = build_transform(False, kmufed=True)(Image.open(sys.argv[1]).convert("RGB"))
    with torch.no_grad():
        print(CLASS_NAMES[model(image[None]).argmax(1).item()])