"""Load auxiliary heads from an smerkd/arrow-heads snapshot, without a backbone."""
from pathlib import Path
import json
import numpy as np
from safetensors.torch import load_file
from arrow_train.heads import TokenHeads

ROOT = Path(__file__).resolve().parents[1]

def load_token_heads(variant='E1', device='cpu'):
    if variant not in ('E0b', 'E1'):
        raise ValueError('Token-head variant must be E0b or E1')
    folder = ROOT / 'weights' / variant
    cfg = json.loads((folder / 'head_config.json').read_text())
    model = TokenHeads(tuple(cfg['names']), cfg['hidden'], cfg['descriptor_dim'])
    model.load_state_dict(load_file(str(folder / 'heads.safetensors')), strict=True)
    return model.eval().to(device)

def load_gaussian_head(device='cpu'):
    from arrow_train.gaussians import GaussianHead, GaussianConfig
    from arrow_train.language import LanguageDecoder
    folder = ROOT / 'weights/G1d'
    cfg = json.loads((folder / 'head_config.json').read_text())
    model = GaussianHead(cfg=GaussianConfig(**cfg))
    with np.load(ROOT / 'assets/label_embeddings.npz') as z:
        model.lang = LanguageDecoder(cfg['lang_dim'], z['targets'])
    model.load_state_dict(load_file(str(folder / 'gaussian_head.safetensors')), strict=True)
    return model.eval().to(device)