"""Load model weights and PCA components."""

import json
from pathlib import Path


def shared_resources(directory, resources_directory):
    if resources_directory is not None:
        return Path(resources_directory)
    for parent in directory.parents:
        resources = parent / "public-resources"
        if resources.is_dir():
            return resources
    raise FileNotFoundError(
        "Shared public-resources directory is missing; supply resources_directory explicitly"
    )


def load_model(model_directory, adata=None, resources_directory=None):
    directory = Path(model_directory).resolve()
    if (directory / "encoder.ckpt").is_file():
        from _native_scimilarity import load_scimilarity

        return load_scimilarity(directory)
    config = json.loads((directory / "config.json").read_text())
    if config.get("model_type") == "bert":
        from transformers import BertForMaskedLM

        model, info = BertForMaskedLM.from_pretrained(
            directory, local_files_only=True, output_loading_info=True,
            attn_implementation="eager",
        )
        if any(info[key] for key in (
            "missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs",
        )):
            raise ValueError(f"Geneformer strict loading failed: {info}")
        model.eval()
        return model
    family = config.get("family")
    if family == "PCA":
        import numpy as np

        components = np.load(directory / "components.npy", allow_pickle=False)
        if list(components.shape) != config["shape"] or str(components.dtype) != config["dtype"]:
            raise ValueError("PCA components disagree with the checkpoint config")
        if components.dtype.kind != "f" or not np.isfinite(components).all():
            raise ValueError("PCA components are not finite numeric data")
        return components
    from safetensors.torch import load_file

    weights = load_file(str(directory / "model.safetensors"), device="cpu")
    if family == "SSL":
        from _native_ssl import native_ssl_class

        resources = shared_resources(directory, resources_directory)
        model = native_ssl_class(resources)(**config["constructor"])
        model.load_state_dict(weights, strict=True)
        model.eval()
        return model
    if family == "scVI":
        import scvi

        if adata is None:
            raise ValueError("scVI requires AnnData with raw counts in X and the exact gene order")
        if list(adata.var_names) != config["gene_names"]:
            raise ValueError("AnnData gene order does not match the checkpoint config")
        scvi.model.SCVI.setup_anndata(adata, **config["setup_anndata"])
        model = scvi.model.SCVI(adata, **config["constructor"])
        model.module.load_state_dict(weights, strict=True)
        model.module.eval()
        model.is_trained = True
        model.to_device("cpu")
        return model
    raise ValueError(f"Unsupported portable model family: {family}")