"""Load trained LCP checkpoints with backward compatibility."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import torch
from transformers import AutoTokenizer

from two_head_model import TwoHeadModel, add_tgt_tokens
from utils import ENCODING_SPAN_MARK, MODELS, POOLING_CLS_CONCAT, POOLING_SPAN


def load_model_from_checkpoint(ckpt_path: str | Path, device: torch.device | None = None):
    ckpt_path = Path(ckpt_path)
    device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)

    model_key = ckpt.get("model_key", "deberta")
    encoding = ckpt.get("encoding", ENCODING_SPAN_MARK)
    pooling = ckpt.get("pooling", POOLING_CLS_CONCAT)  # legacy checkpoints
    use_linguistic = ckpt.get("use_linguistic_features", False)
    level_only = ckpt.get("level_only", False)

    tokenizer = AutoTokenizer.from_pretrained(MODELS[model_key])
    tgt_id, tgt_end_id = add_tgt_tokens(tokenizer)

    model = TwoHeadModel(
        model_name=MODELS[model_key],
        pooling_mode=pooling,
        use_linguistic_features=use_linguistic,
    )
    model.encoder.resize_token_embeddings(len(tokenizer))
    model.load_state_dict(ckpt["model_state"])
    model.set_tgt_token_ids(ckpt.get("tgt_id", tgt_id), ckpt.get("tgt_end_id", tgt_end_id))
    model.to(device)
    model.eval()

    feat_mean = np.array(ckpt["feat_mean"]) if ckpt.get("feat_mean") is not None else None
    feat_std = np.array(ckpt["feat_std"]) if ckpt.get("feat_std") is not None else None

    meta = {
        "model_key": model_key,
        "encoding": encoding,
        "pooling": pooling,
        "use_linguistic_features": use_linguistic,
        "level_only": level_only,
        "tgt_id": ckpt.get("tgt_id", tgt_id),
        "tgt_end_id": ckpt.get("tgt_end_id", tgt_end_id),
        "feat_mean": feat_mean,
        "feat_std": feat_std,
    }
    return model, tokenizer, meta