"""Small registry mapping photon count -> which pretrained 12 keV checkpoint
and model class to use, so the demo scripts don't repeat this three times.
"""
from pathlib import Path
import torch
from models import get_model_class, get_double_photon_model_class, get_triple_photon_model_class

MODELS_DIR = Path(__file__).resolve().parent.parent / 'models'

REGISTRY = {
    1: dict(version='260511', checkpoint=MODELS_DIR / 'singlePhoton_12keV.pth',
            sample_size=3, get_class=get_model_class),
    2: dict(version='260610', checkpoint=MODELS_DIR / 'doublePhoton_12keV.pth',
            sample_size=6, get_class=get_double_photon_model_class),
    3: dict(version='260611', checkpoint=MODELS_DIR / 'triplePhoton_12keV.pth',
            sample_size=9, get_class=get_triple_photon_model_class),
}


def load_pretrained(n_photons, device='cuda'):
    info = REGISTRY[n_photons]
    model = info['get_class'](info['version'])().to(device)
    model.load_state_dict(torch.load(info['checkpoint'], map_location=device, weights_only=True))
    model.eval()
    return model


def new_model(n_photons, device='cuda'):
    """Same architecture as load_pretrained(), but randomly initialized (for training from scratch)."""
    info = REGISTRY[n_photons]
    return info['get_class'](info['version'])().to(device)