"""Load and export PRIMA trainable state without changing parameter names."""
from pathlib import Path
import torch


PREFIX = 'base_model.model.'
REQUIRED_GROUPS = (
    'model.layers.', 'model.grounding_encoder.mask_decoder.',
    'model.query_tokens', 'model.residual_query_tokens',
    'model.residual_extractor.', 'model.mm_projector.', 'model.text_hidden_fcs.',
)


def load_checkpoint(model, path, *, require_complete=True):
    payload = torch.load(path, map_location='cpu')
    state = payload.get('state_dict', payload) if isinstance(payload, dict) else payload
    if not isinstance(state, dict) or not state:
        raise ValueError('Checkpoint must contain a nonempty tensor state dictionary.')
    canonical = {}
    for key, value in state.items():
        name = key[len(PREFIX):] if key.startswith(PREFIX) else key
        if name in canonical:
            raise ValueError(f'Checkpoint prefix normalization produced a duplicate: {name}')
        if not torch.is_tensor(value):
            raise ValueError(f'Checkpoint entry is not a tensor: {name}')
        canonical[name] = value
    current = model.state_dict()
    wrapped = any(key.startswith(PREFIX) for key in current)
    normalized = {(PREFIX + key if wrapped else key): value for key, value in canonical.items()}
    # These are recomputable buffers from supported dependency versions.
    ignored = [key for key in normalized if key not in current and
               ('rotary_emb.inv_freq' in key or key.endswith('position_ids'))]
    for key in ignored:
        normalized.pop(key)
    unexpected = set(normalized) - set(current)
    if unexpected:
        raise ValueError('Unexpected checkpoint tensors: ' + ', '.join(sorted(unexpected)))
    if require_complete:
        missing_groups = [group for group in REQUIRED_GROUPS
                          if not any(key.startswith(group) for key in canonical)]
        required = {key for key in current
                    if any((key[len(PREFIX):] if wrapped else key).startswith(group)
                           for group in REQUIRED_GROUPS)
                    and ('lora_' in key or 'model.layers.' not in key)
                    and 'rotary_emb.inv_freq' not in key}
        missing = required - set(normalized)
        if missing_groups or missing:
            raise ValueError('Incomplete PRIMA checkpoint. Missing groups: '
                             + ', '.join(missing_groups) + '; missing tensors: '
                             + ', '.join(sorted(missing)))
    model.load_state_dict(normalized, strict=False)
    return {'loaded_tensors': len(normalized), 'ignored_buffers': ignored}


def export_checkpoint(model, path):
    """Export trainable parameters plus persistent buffers, with compatible keys."""
    state = {}
    for key, value in model.state_dict().items():
        module_path, _, leaf = key.rpartition('.')
        parent = model.get_submodule(module_path) if module_path else model
        parameter = parent._parameters.get(leaf)
        if parameter is not None and not parameter.requires_grad:
            continue
        name = key[len(PREFIX):] if key.startswith(PREFIX) else key
        if name in state:
            raise ValueError(f'Export normalization produced a duplicate: {name}')
        state[name] = value.detach().cpu().clone()
    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_name(destination.name + '.tmp')
    torch.save(state, temporary)
    temporary.replace(destination)