# SPDX-License-Identifier: AGPL-3.0-only
# Safetensors loading added 2026-10-28; inference preprocessing is unchanged.
"""Strict, direct safetensors loader for this pinned RF-DETR Small release."""
from importlib.metadata import version
import json
from pathlib import Path
import warnings

import torch
from safetensors.torch import load_file

HERE = Path(__file__).resolve().parent


def load_model(weights_path=None, config_path=None, device='cpu'):
    """Rebuild Small, load every tensor strictly, and return its RF-DETR facade.

    No checkpoint deserialization, temporary pickle, pretrained download, dtype
    conversion, or inference optimization is used. The CLI supplies the existing
    RGB preparation, tiling, boundary rechecks and NMS via inference_engine.py.
    """
    weights_path = Path(weights_path) if weights_path is not None else HERE / 'model.safetensors'
    config_path = Path(config_path) if config_path is not None else weights_path.parent / 'model_config.json'
    if weights_path.suffix != '.safetensors':
        raise ValueError('Use a .safetensors weights file.')
    config = json.loads(config_path.read_text())
    expected_config = {
        'variant': 'RFDETRSmall', 'rfdetr_version': '1.11.2',
        'resolution': 512, 'num_classes': 1, 'id2label': {'0': 'door'},
    }
    if config != expected_config:
        raise ValueError('model_config.json does not describe this RF-DETR Small release.')
    if version('rfdetr') != config['rfdetr_version']:
        raise RuntimeError('Install the pinned rfdetr==1.11.2 inference dependencies.')
    if device not in ('cpu', 'mps', 'cuda'):
        raise ValueError('Choose cpu, mps or cuda.')
    if device == 'mps' and not torch.backends.mps.is_available():
        raise RuntimeError('MPS unavailable; use --device cpu.')
    if device == 'cuda' and not torch.cuda.is_available():
        raise RuntimeError('CUDA unavailable; use --device cpu.')

    from rfdetr import RFDETRSmall
    from rfdetr.config import PretrainWeightsCompatibilityWarning

    # The pinned Small constructor uses bundled backbone configuration. Patch 16
    # and positional grid 32 create the backbone locally, without DINOv2 weights.
    # Its expected "initialised from scratch" warning is superseded by the strict
    # complete-state load below; no randomly initialized weights remain in use.
    with warnings.catch_warnings():
        warnings.simplefilter('ignore', PretrainWeightsCompatibilityWarning)
        model = RFDETRSmall(pretrain_weights=None, num_classes=config['num_classes'],
                           resolution=config['resolution'], device=device)
    module = model.model.model
    state = load_file(str(weights_path), device='cpu')
    expected = module.state_dict()
    missing = sorted(expected.keys() - state.keys())
    unexpected = sorted(state.keys() - expected.keys())
    if missing or unexpected:
        raise RuntimeError(f'Weight key mismatch: missing={missing}, unexpected={unexpected}')
    for name, tensor in state.items():
        if tensor.shape != expected[name].shape or tensor.dtype != expected[name].dtype:
            raise RuntimeError(f'Weight shape/dtype mismatch: {name}')
    module.load_state_dict(state, strict=True)
    loaded_state = module.state_dict()
    if not all(torch.equal(tensor, loaded_state[name]) for name, tensor in state.items()):
        raise RuntimeError('Loaded model tensors do not exactly match the safetensors file.')
    module.eval()
    model.model.class_names = [config['id2label'][str(i)] for i in range(config['num_classes'])]
    model.model.args.class_names = model.model.class_names
    return model