"""Load released SimpleP2G prediction heads without an image encoder."""
from __future__ import annotations

import json
from pathlib import Path
from types import SimpleNamespace

import torch
from safetensors.torch import load_file

from src.simpleP2G import SimpleP2G


DATASETS = ("brain", "prad", "her2st", "lung", "ccrcc")
PANELS = ("hmhvg", "hvg")
SETTINGS = ("single-patch", "cross-patch")


def load_pretrained(
    repo_id_or_path="king159/SimpleP2G",
    *,
    dataset="brain",
    panel="hmhvg",
    setting="single-patch",
    device="cpu",
    revision=None,
    cache_dir=None,
    local_files_only=False,
):
    """Return an eval-mode model, ordered gene names and checkpoint metadata.

    ``repo_id_or_path`` is a Hub model ID or a downloaded repository directory.
    Inputs are precomputed image embeddings; no image-encoder weights are loaded.
    Predictions follow the order in ``genes`` and use log2(count + 1) units.
    """
    dataset, panel = dataset.lower(), panel.lower()
    if dataset not in DATASETS or panel not in PANELS or setting not in SETTINGS:
        raise ValueError(f"Unsupported model selection: {dataset}/{panel}/{setting}")
    folder = f"models/{setting}/{dataset}/{panel}"
    local = Path(repo_id_or_path).expanduser()

    def resolve(filename):
        if local.is_dir():
            path = local / folder / filename
            if not path.is_file():
                raise FileNotFoundError(path)
            return path
        from huggingface_hub import hf_hub_download

        return Path(hf_hub_download(
            repo_id=str(repo_id_or_path), filename=f"{folder}/{filename}",
            revision=revision, cache_dir=cache_dir, local_files_only=local_files_only,
        ))

    metadata = json.loads(resolve("config.json").read_text())
    if metadata.get("format_version") != 1:
        raise ValueError("Unsupported checkpoint metadata version")
    expected = {"dataset": dataset, "panel": panel, "setting": setting}
    if any(metadata.get(key) != value for key, value in expected.items()):
        raise ValueError("Checkpoint metadata does not match the requested model")
    genes = json.loads(resolve("genes.json").read_text())
    config = SimpleNamespace(**metadata["model_config"])
    if len(genes) != config.num_genes or len(set(genes)) != len(genes):
        raise ValueError("The ordered gene list must contain one unique name per output")
    if setting == "single-patch":
        model = SimpleP2G(config)
    else:
        from src.simpleP2G_neigh import SimpleP2GNeigh

        model = SimpleP2GNeigh(config)
    state = load_file(str(resolve("model.safetensors")), device="cpu")
    model.load_state_dict(state, strict=True)
    model.to(device=device).eval().requires_grad_(False)
    return model, genes, metadata