"""Load the pinned ONNX tensors into the same frozen Torch backbone used in evaluation."""

import json
from pathlib import Path

import numpy as np
import onnx
import torch
from onnx import numpy_helper
from transformers import LlamaConfig, LlamaModel
from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding

from .core import projection_weights


def load_backbone(root: Path, device: str) -> LlamaModel:
    source = json.loads((root / "config.json").read_text())
    graph = onnx.load(root / "emma6.onnx", load_external_data=False)
    tensors = {}
    for tensor in graph.graph.initializer:
        if tensor.data_type != onnx.TensorProto.FLOAT:
            continue
        data = {entry.key: entry.value for entry in tensor.external_data}
        if data:
            array = np.memmap(root / data["location"], dtype=np.float32, mode="c",
                              offset=int(data.get("offset", 0)), shape=tuple(tensor.dims))
        else:
            array = numpy_helper.to_array(tensor).copy()
        tensors[tensor.name] = torch.from_numpy(array)
    config = LlamaConfig(
        vocab_size=source["vocab_size"], hidden_size=source["hidden_size"],
        intermediate_size=source["intermediate_size"], num_hidden_layers=source["num_hidden_layers"],
        num_attention_heads=source["num_attention_heads"], num_key_value_heads=source["num_kv_heads"],
        max_position_embeddings=source["context_length"], rms_norm_eps=1e-6,
        rope_parameters={"rope_type": "default", "rope_theta": source["rope_theta"]},
        attention_bias=False, mlp_bias=False, use_cache=False, pad_token_id=0,
        bos_token_id=2, eos_token_id=3,
    )
    config._attn_implementation = "sdpa"
    with torch.device("meta"):
        model = LlamaModel(config)
    state = projection_weights(graph, tensors)
    state.update({"embed_tokens.weight": tensors["token_emb.weight"], "norm.weight": tensors["norm_f.weight"]})
    for i in range(config.num_hidden_layers):
        state[f"layers.{i}.input_layernorm.weight"] = tensors[f"blocks.{i}.norm1.weight"]
        state[f"layers.{i}.post_attention_layernorm.weight"] = tensors[f"blocks.{i}.norm2.weight"]
    model.load_state_dict(state, strict=True, assign=True)
    model.rotary_emb = LlamaRotaryEmbedding(config)
    return model.eval().requires_grad_(False).to(device)