# SPDX-License-Identifier: LGPL-3.0-or-later
"""The model that takes the coordinates, cell and atom types as input
and predicts some property. The models are automatically generated from
atomic models by the `deepmd.dpmodel.make_model` method.

The `make_model` method does the reduction, auto-differentiation and
communication of the atomic properties according to output variable
definition `deepmd.dpmodel.OutputVariableDef`.

All models should be inherited from :class:`deepmd.pt.model.model.model.BaseModel`.
Models generated by `make_model` have already done it.
"""

import copy
import json
from typing import (
    Any,
)

import numpy as np

from deepmd.pt.model.atomic_model import (
    DPAtomicModel,
    PairTabAtomicModel,
)
from deepmd.pt.model.descriptor.base_descriptor import (
    BaseDescriptor,
)
from deepmd.pt.model.task import (
    BaseFitting,
)
from deepmd.pt.model.task.sezm_ener import (
    SeZMEnergyFittingNet,
    _resolve_auto_neuron,
)
from deepmd.pt.utils.multi_task import (
    preprocess_shared_params,
)
from deepmd.utils.bridging import (
    DEFAULT_FRACTION_INNER,
    DEFAULT_FRACTION_OUTER,
    concise_bridging_window,
    expand_bridging_method,
    resolve_bridging_window,
    route_canonical_learned_options,
    validate_tabulated_zbl,
)
from deepmd.utils.spin import (
    Spin,
    env_protection_descriptors,
)

from .density_model import (
    GridDensityModel,
)
from .dipole_model import (
    DipoleModel,
)
from .dos_model import (
    DOSModel,
)
from .dp_linear_model import (
    LinearEnergyModel,
    normalize_linear_model_type_map,
    validate_linear_shared_descriptor_type_maps,
)
from .dp_model import (
    DPModelCommon,
)
from .dp_zbl_model import (
    DPZBLModel,
)
from .ener_model import (
    EnergyModel,
)
from .frozen import (
    FrozenModel,
)
from .make_hessian_model import (
    make_hessian_model,
)
from .make_model import (
    make_model,
)
from .model import (
    BaseModel,
)
from .polar_model import (
    PolarModel,
)
from .population_model import (
    PopulationModel,
)
from .property_model import (
    PropertyModel,
)
from .sezm_model import (
    SeZMModel,
)
from .sezm_native_spin_model import (
    SeZMNativeSpinModel,
)
from .sezm_property_model import (
    SeZMPropertyModel,
)
from .spin_model import (
    SpinEnergyModel,
    SpinModel,
)


def _reject_long_range(model_params: dict) -> None:
    """Reject a long-range module, which only the pt_expt backend serves."""
    if model_params.get("long_range") is not None:
        raise NotImplementedError(
            "long-range modules (`long_range`) are supported by the pt_expt "
            "backend only"
        )


def _get_standard_model_components(model_params: dict, ntypes: int) -> tuple:
    if "type_embedding" in model_params:
        raise ValueError(
            "In the PyTorch backend, type_embedding is not at the model level, but within the descriptor. See type embedding documentation for details."
        )
    # descriptor
    model_params["descriptor"]["ntypes"] = ntypes
    model_params["descriptor"]["type_map"] = copy.deepcopy(model_params["type_map"])
    descriptor = BaseDescriptor(**model_params["descriptor"])
    # fitting
    fitting_net = model_params.get("fitting_net", {})
    fitting_net["type"] = fitting_net.get("type", "ener")
    fitting_net["ntypes"] = descriptor.get_ntypes()
    fitting_net["type_map"] = copy.deepcopy(model_params["type_map"])
    fitting_net["mixed_types"] = descriptor.mixed_types()
    if fitting_net["type"] in ["dipole", "polar"]:
        fitting_net["embedding_width"] = descriptor.get_dim_emb()
    fitting_net["dim_descrpt"] = descriptor.get_dim_out()
    grad_force = "direct" not in fitting_net["type"]
    if not grad_force:
        fitting_net["out_dim"] = descriptor.get_dim_emb()
        if "ener" in fitting_net["type"]:
            fitting_net["return_energy"] = True
    fitting = BaseFitting(**fitting_net)
    return descriptor, fitting, fitting_net["type"]


def _normalize_spin_use_spin(model_params: dict) -> None:
    """Normalize ``spin.use_spin`` to a per-type boolean list.

    Three equivalent forms are accepted: a per-type boolean list, a list of
    magnetic type indices, or a list of magnetic element symbols. The index and
    symbol forms are expanded against ``type_map``, so a large type map only needs
    its magnetic species named.

    Raises
    ------
    ValueError
        If a symbol in ``use_spin`` is absent from ``type_map``.
    """
    use_spin = model_params["spin"]["use_spin"]
    if use_spin and isinstance(use_spin[0], str):
        type_index = {name: idx for idx, name in enumerate(model_params["type_map"])}
        unknown = [name for name in use_spin if name not in type_index]
        if unknown:
            raise ValueError(
                f"spin.use_spin references element(s) {unknown} absent from type_map."
            )
        use_spin = [type_index[name] for name in use_spin]
    # ``bool`` is a subclass of ``int``; an already-boolean list is left untouched
    # while an index list is scattered into a per-type mask.
    if not use_spin or not isinstance(use_spin[0], bool):
        mask = np.full(len(model_params["type_map"]), False, dtype=bool)
        mask[use_spin] = True
        model_params["spin"]["use_spin"] = mask.tolist()


def get_spin_model(model_params: dict) -> SpinModel:
    model_params = copy.deepcopy(model_params)
    if model_params["spin"].get("allow_missing_label", False):
        raise ValueError(
            "spin.allow_missing_label is supported only by the SeZM/DPA4 spin model "
            "(model.type='dpa4'/'sezm'), where a zero spin reduces the descriptor to "
            "its spin-free form; the virtual-atom expansion here has no such limit."
        )
    _normalize_spin_use_spin(model_params)
    # include virtual spin and placeholder types
    model_params["type_map"] += [item + "_spin" for item in model_params["type_map"]]
    spin = Spin(
        use_spin=model_params["spin"]["use_spin"],
        virtual_scale=model_params["spin"]["virtual_scale"],
    )
    pair_exclude_types = spin.get_pair_exclude_types(
        exclude_types=model_params.get("pair_exclude_types", None)
    )
    model_params["pair_exclude_types"] = pair_exclude_types
    # for descriptor data stat
    model_params["descriptor"]["exclude_types"] = pair_exclude_types
    atom_exclude_types = spin.get_atom_exclude_types(
        exclude_types=model_params.get("atom_exclude_types", None)
    )
    model_params["atom_exclude_types"] = atom_exclude_types
    for descriptor in env_protection_descriptors(model_params["descriptor"]):
        if "env_protection" not in descriptor or descriptor["env_protection"] == 0.0:
            descriptor["env_protection"] = 0.01
    if model_params["descriptor"]["type"] in ["se_e2_a"]:
        # only expand sel for se_e2_a
        model_params["descriptor"]["sel"] += model_params["descriptor"]["sel"]
    backbone_model = get_standard_model(model_params)
    return SpinEnergyModel(backbone_model=backbone_model, spin=spin)


def get_linear_model(model_params: dict) -> BaseModel:
    for sub in model_params.get("models", []):
        if str(sub.get("bridging_method", "none")).lower() not in ("none", ""):
            raise ValueError(
                "`bridging_method` is not supported on a linear_ener "
                "sub-model: add an `inner_potential` sub-model to the "
                "composition instead."
            )
        _reject_long_range(sub)
    if any(
        sub.get("type") == "inner_potential" for sub in model_params.get("models", [])
    ):
        return _get_bridged_linear_model(model_params)
    model_params = copy.deepcopy(model_params)
    weights = model_params.get("weights", "mean")
    shared_links = None
    if "shared_dict" in model_params:
        shared_config = {
            "model_dict": {
                f"model_{idx}": sub_model
                for idx, sub_model in enumerate(model_params["models"])
            },
            "shared_dict": model_params.get("shared_dict", {}),
        }
        if "type_map" in model_params:
            shared_config["type_map"] = copy.deepcopy(model_params["type_map"])
        shared_config, shared_links = preprocess_shared_params(
            shared_config,
            require_shared_type_map=False,
        )
        model_params["models"] = list(shared_config["model_dict"].values())
        normalize_linear_model_type_map(model_params)
        validate_linear_shared_descriptor_type_maps(
            model_params["models"],
            shared_links,
        )

    list_of_models = []
    for sub_model_params in model_params["models"]:
        if "type_map" not in sub_model_params:
            sub_model_params["type_map"] = model_params["type_map"]
        if "descriptor" in sub_model_params:
            # descriptor
            sub_ntypes = len(sub_model_params["type_map"])
            sub_model_params["descriptor"]["ntypes"] = sub_ntypes
            descriptor, fitting, _ = _get_standard_model_components(
                sub_model_params, sub_ntypes
            )
            list_of_models.append(
                DPAtomicModel(
                    descriptor,
                    fitting,
                    type_map=copy.deepcopy(sub_model_params["type_map"]),
                )
            )

        else:  # must be pairtab
            assert (
                "type" in sub_model_params and sub_model_params["type"] == "pairtab"
            ), "Sub-models in LinearEnergyModel must be a DPModel or a PairTable Model"
            list_of_models.append(
                PairTabAtomicModel(
                    sub_model_params["tab_file"],
                    sub_model_params["rcut"],
                    sub_model_params["sel"],
                    type_map=copy.deepcopy(sub_model_params["type_map"]),
                )
            )

    atom_exclude_types = model_params.get("atom_exclude_types", [])
    pair_exclude_types = model_params.get("pair_exclude_types", [])
    model = LinearEnergyModel(
        models=list_of_models,
        type_map=model_params["type_map"],
        weights=weights,
        atom_exclude_types=atom_exclude_types,
        pair_exclude_types=pair_exclude_types,
        preset_out_bias=model_params.get("preset_out_bias"),
    )
    model.shared_links = shared_links
    if shared_links:
        model.share_params(shared_links, resume=True)
    return model


def _get_bridged_linear_model(model_params: dict) -> BaseModel:
    """Realize a bridged ``linear_ener`` composition in the pt backend.

    The pt backend implements analytical bridging inside ``SeZMModel``
    (the ``InnerPotential`` term is a sub-module of the learned model),
    so a canonical ``linear_ener`` composition over
    ``[learned, inner_potential]`` maps onto the ``SeZMModel``
    constructor arguments. The physics and the checkpoint format are
    identical to the legacy ``bridging_method`` flag form.

    Parameters
    ----------
    model_params : dict
        A ``linear_ener`` config with exactly one ``inner_potential``
        sub-model and exactly one learned (descriptor-bearing) sub-model.

    Raises
    ------
    ValueError
        If the composition shape is not the bridging one (child counts,
        ``weights``, ``shared_dict``).
    NotImplementedError
        If the learned sibling is not of the DPA4/SeZM family: the pt
        backend has no bridging implementation for other descriptors.
    """
    model_params = copy.deepcopy(model_params)
    children = model_params.get("models", [])
    inner_cfgs = [sub for sub in children if sub.get("type") == "inner_potential"]
    learned_cfgs = [sub for sub in children if "descriptor" in sub]
    if len(inner_cfgs) > 1:
        raise ValueError(
            "A linear_ener composition supports at most one "
            "`inner_potential` sub-model."
        )
    if len(learned_cfgs) != 1 or len(children) != 2:
        raise ValueError(
            "An `inner_potential` sub-model bridges exactly one learned "
            "sibling: expected a linear_ener composition over "
            "[learned, inner_potential]."
        )
    if str(model_params.get("weights", "mean")) != "sum":
        raise ValueError('A bridged linear_ener composition requires `weights: "sum"`.')
    if model_params.get("shared_dict"):
        raise NotImplementedError(
            "`shared_dict` is not supported with an `inner_potential` sub-model."
        )
    learned = copy.deepcopy(learned_cfgs[0])
    descriptor_type = str(learned.get("descriptor", {}).get("type", "dpa4"))
    if descriptor_type not in ("dpa4", "DPA4", "sezm", "SeZM"):
        raise NotImplementedError(
            "The pt backend implements `inner_potential` bridging only for "
            f"the DPA4/SeZM descriptor family, but got {descriptor_type!r}."
        )
    if learned.get("lora") is not None:
        raise NotImplementedError(
            "`lora` on the learned child of a bridged linear_ener "
            "composition is not supported: the pt trainer reads `lora` "
            "from the top-level model section only. Use the concise "
            '`type: "dpa4"` form with top-level `lora` and '
            "`bridging_method` instead."
        )
    learned_type_map = learned.get("type_map", model_params["type_map"])
    if learned_type_map != model_params["type_map"]:
        raise NotImplementedError(
            "The pt backend requires the learned child's `type_map` to "
            "match the bridged linear_ener composition's `type_map`."
        )
    route_canonical_learned_options(model_params, learned)
    validate_tabulated_zbl(learned)
    inner_cfg = inner_cfgs[0]
    learned["type"] = "dpa4"
    learned["type_map"] = copy.deepcopy(model_params["type_map"])
    learned["atom_exclude_types"] = model_params.get("atom_exclude_types", [])
    learned["pair_exclude_types"] = model_params.get("pair_exclude_types", [])
    learned["bridging_method"] = inner_cfg.get("mode", "zbl")
    learned["bridging_fraction_inner"] = inner_cfg.get(
        "fraction_inner", DEFAULT_FRACTION_INNER
    )
    learned["bridging_fraction_outer"] = inner_cfg.get(
        "fraction_outer", DEFAULT_FRACTION_OUTER
    )
    learned["bridging_r_inner"] = inner_cfg.get("r_inner")
    learned["bridging_r_outer"] = inner_cfg.get("r_outer")
    if "spin" in model_params:
        learned["spin"] = model_params["spin"]
        return get_sezm_spin_model(learned)
    return get_sezm_model(learned)


def get_zbl_model(model_params: dict) -> DPZBLModel:
    validate_tabulated_zbl(model_params)
    _reject_long_range(model_params)
    model_params = copy.deepcopy(model_params)
    ntypes = len(model_params["type_map"])
    descriptor, fitting, _ = _get_standard_model_components(model_params, ntypes)
    dp_model = DPAtomicModel(descriptor, fitting, type_map=model_params["type_map"])
    # pairtab
    filepath = model_params["use_srtab"]
    pt_model = PairTabAtomicModel(
        filepath,
        descriptor.get_rcut(),
        descriptor.get_sel(),
        type_map=model_params["type_map"],
    )

    rmin = model_params["sw_rmin"]
    rmax = model_params["sw_rmax"]
    atom_exclude_types = model_params.get("atom_exclude_types", [])
    pair_exclude_types = model_params.get("pair_exclude_types", [])
    model = DPZBLModel(
        dp_model,
        pt_model,
        rmin,
        rmax,
        type_map=model_params["type_map"],
        atom_exclude_types=atom_exclude_types,
        pair_exclude_types=pair_exclude_types,
        preset_out_bias=model_params.get("preset_out_bias"),
    )
    model.model_def_script = json.dumps(model_params)
    return model


def get_standard_model(model_params: dict) -> BaseModel:
    bridging_method = str(model_params.get("bridging_method", "none"))
    if bridging_method.lower() not in ("none", ""):
        raise ValueError(
            "`bridging_method` is not supported for a standard model: "
            "analytical bridging builds a linear composition, not a "
            "standard model. Route the config through `get_model` (which "
            "expands the flag), or spell the composition explicitly with "
            '`type: "linear_ener"` and an `inner_potential` sub-model.'
        )
    _reject_long_range(model_params)
    model_params_old = model_params
    model_params = copy.deepcopy(model_params)
    ntypes = len(model_params["type_map"])
    descriptor, fitting, fitting_net_type = _get_standard_model_components(
        model_params, ntypes
    )
    atom_exclude_types = model_params.get("atom_exclude_types", [])
    pair_exclude_types = model_params.get("pair_exclude_types", [])
    preset_out_bias = model_params.get("preset_out_bias")
    data_stat_protect = model_params.get("data_stat_protect", 1e-2)

    if fitting_net_type == "dipole":
        modelcls = DipoleModel
    elif fitting_net_type == "polar":
        modelcls = PolarModel
    elif fitting_net_type == "dos":
        modelcls = DOSModel
    elif fitting_net_type in ["ener", "direct_force_ener"]:
        modelcls = EnergyModel
    elif fitting_net_type == "density":
        modelcls = GridDensityModel
    elif fitting_net_type == "property":
        modelcls = PropertyModel
    elif fitting_net_type == "population":
        modelcls = PopulationModel
    else:
        raise RuntimeError(f"Unknown fitting type: {fitting_net_type}")

    model = modelcls(
        descriptor=descriptor,
        fitting=fitting,
        type_map=model_params["type_map"],
        atom_exclude_types=atom_exclude_types,
        pair_exclude_types=pair_exclude_types,
        preset_out_bias=preset_out_bias,
        data_stat_protect=data_stat_protect,
    )
    if model_params.get("hessian_mode"):
        model.enable_hessian()
    model.model_def_script = json.dumps(model_params_old)
    return model


def _bridging_window_options(model_params: dict) -> dict:
    """
    Resolve the bridging window of a SeZM model section.

    Parameters
    ----------
    model_params : dict
        A DPA4/SeZM model section carrying the concise bridging keys.

    Returns
    -------
    dict
        The window options of the learned descriptor, as
        :func:`deepmd.utils.bridging.resolve_bridging_window` spells them.
    """
    return resolve_bridging_window(concise_bridging_window(model_params))


def _reconcile_sezm_pair_exclude_types(model_params: dict) -> list[list[int]]:
    """Reconcile ``pair_exclude_types`` with ``descriptor.exclude_types``.

    A DPA4/SeZM config may spell pair exclusions at the model level
    (``pair_exclude_types``) or on the descriptor (``exclude_types``).
    Every SeZM builder (plain, native spin, virtual spin) resolves the two
    through this ONE helper so that a mismatch always fails fast instead
    of one spelling silently overwriting the other.

    Parameters
    ----------
    model_params : dict
        The DPA4/SeZM model config; ``model_params["descriptor"]`` must
        exist.

    Returns
    -------
    list[list[int]]
        The reconciled real-type pair exclusion list.

    Raises
    ------
    ValueError
        If both spellings are given and differ.
    """
    descriptor_exclude_types = [
        list(pair) for pair in (model_params["descriptor"].get("exclude_types") or [])
    ]
    if "pair_exclude_types" in model_params:
        pair_exclude_types = [
            list(pair) for pair in (model_params["pair_exclude_types"] or [])
        ]
        if descriptor_exclude_types and descriptor_exclude_types != pair_exclude_types:
            raise ValueError(
                "SeZM `pair_exclude_types` and `descriptor.exclude_types` must match "
                "when both are provided."
            )
    else:
        pair_exclude_types = descriptor_exclude_types
    return pair_exclude_types


def get_sezm_model(model_params: dict) -> BaseModel:
    model_params_old = model_params
    model_params = copy.deepcopy(model_params)
    model_params.setdefault("descriptor", {})
    model_params.setdefault("fitting_net", {})
    model_params["descriptor"].setdefault("type", "dpa4")

    ntypes = len(model_params["type_map"])
    model_params["descriptor"]["ntypes"] = ntypes
    model_params["descriptor"]["type_map"] = copy.deepcopy(model_params["type_map"])
    pair_exclude_types = _reconcile_sezm_pair_exclude_types(model_params)
    model_params["pair_exclude_types"] = pair_exclude_types
    model_params["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types)

    # === Bridging parameters ===
    bridging_method = str(model_params.get("bridging_method", "none")).upper()
    bridging_window = _bridging_window_options(model_params)
    # Only inject bridging parameters when bridging is enabled.
    if bridging_method != "NONE":
        model_params["descriptor"].update(bridging_window)

    descriptor = BaseDescriptor(**model_params["descriptor"])

    fitting_net = copy.deepcopy(model_params["fitting_net"])
    fitting_net.setdefault("type", "dpa4_ener")
    fitting_net["ntypes"] = descriptor.get_ntypes()
    fitting_net["type_map"] = copy.deepcopy(model_params["type_map"])
    fitting_net["mixed_types"] = descriptor.mixed_types()
    fitting_net["dim_descrpt"] = descriptor.get_dim_out()
    if fitting_net["type"] in ("dpa4_ener", "sezm_ener"):
        fitting = BaseFitting(**fitting_net)
        modelcls = SeZMModel
    elif fitting_net["type"] == "property":
        if bridging_method != "NONE":
            raise ValueError(
                "DPA4/SeZM property fitting does not support analytical bridging "
                "potentials; set `bridging_method` to `none`."
            )
        # Share the SeZM auto-width convention
        fitting_net["neuron"] = _resolve_auto_neuron(
            fitting_net.get("neuron"),
            dim_descrpt=fitting_net["dim_descrpt"],
            numb_fparam=fitting_net.get("numb_fparam", 0),
            numb_aparam=fitting_net.get("numb_aparam", 0),
            dim_case_embd=fitting_net.get("dim_case_embd", 0),
            case_film_embd=fitting_net.get("case_film_embd", False),
            use_aparam_as_mask=fitting_net.get("use_aparam_as_mask", False),
        )
        fitting = BaseFitting(**fitting_net)
        modelcls = SeZMPropertyModel
    else:
        raise ValueError(
            "DPA4/SeZM model supports `dpa4_ener`, `sezm_ener`, or `property` "
            f"fitting, but got `{fitting_net['type']}`."
        )
    atom_exclude_types = model_params.get("atom_exclude_types", [])
    preset_out_bias = model_params.get("preset_out_bias")
    data_stat_protect = model_params.get("data_stat_protect", 1e-2)
    use_compile = bool(model_params.get("use_compile", False))
    enable_tf32 = bool(model_params.get("enable_tf32", True))

    model = modelcls(
        descriptor=descriptor,
        fitting=fitting,
        type_map=model_params["type_map"],
        atom_exclude_types=atom_exclude_types,
        pair_exclude_types=pair_exclude_types,
        preset_out_bias=preset_out_bias,
        data_stat_protect=data_stat_protect,
        use_compile=use_compile,
        enable_tf32=enable_tf32,
        bridging_method=bridging_method,
    )
    model.model_def_script = json.dumps(model_params_old)
    return model


def get_sezm_spin_model(model_params: dict) -> BaseModel:
    """Build the native spin DPA4/SeZM model.

    The spin vector enters the descriptor as an equivariant feature, so the
    type map, neighbor selection and type count stay at the real-system sizes;
    only the descriptor gains the spin-embedding branch.
    """
    if str(model_params.get("spin", {}).get("scheme", "")).lower() != "native":
        raise ValueError("DPA4/SeZM spin models require model.spin.scheme='native'.")
    model_params_old = model_params
    model_params = copy.deepcopy(model_params)
    model_params.setdefault("descriptor", {})
    model_params.setdefault("fitting_net", {})
    model_params["descriptor"].setdefault("type", "dpa4")
    _normalize_spin_use_spin(model_params)

    use_spin = [bool(flag) for flag in model_params["spin"]["use_spin"]]
    spin = Spin(
        use_spin=use_spin,
        allow_missing_label=model_params["spin"].get("allow_missing_label", False),
    )

    ntypes = len(model_params["type_map"])
    model_params["descriptor"]["ntypes"] = ntypes
    model_params["descriptor"]["type_map"] = copy.deepcopy(model_params["type_map"])
    model_params["descriptor"]["use_spin"] = use_spin

    pair_exclude_types = _reconcile_sezm_pair_exclude_types(model_params)
    model_params["pair_exclude_types"] = pair_exclude_types
    if pair_exclude_types:
        model_params["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types)

    # === Bridging parameters ===
    bridging_method = str(model_params.get("bridging_method", "none")).upper()
    bridging_window = _bridging_window_options(model_params)
    if bridging_method != "NONE":
        model_params["descriptor"].update(bridging_window)

    descriptor = BaseDescriptor(**model_params["descriptor"])

    fitting_net = copy.deepcopy(model_params["fitting_net"])
    fitting_net_type = fitting_net.get("type", "dpa4_ener")
    if fitting_net_type not in ("dpa4_ener", "sezm_ener"):
        raise ValueError(
            "Native spin DPA4/SeZM currently supports only `dpa4_ener` or "
            f"`sezm_ener` fitting, but got `{fitting_net_type}`."
        )
    fitting_net.pop("type", None)
    fitting_net["ntypes"] = descriptor.get_ntypes()
    fitting_net["type_map"] = copy.deepcopy(model_params["type_map"])
    fitting_net["mixed_types"] = descriptor.mixed_types()
    fitting_net["dim_descrpt"] = descriptor.get_dim_out()
    fitting = SeZMEnergyFittingNet(**fitting_net)

    preset_out_bias = model_params.get("preset_out_bias")
    data_stat_protect = model_params.get("data_stat_protect", 1e-2)
    use_compile = bool(model_params.get("use_compile", False))
    enable_tf32 = bool(model_params.get("enable_tf32", True))

    model = SeZMNativeSpinModel(
        descriptor=descriptor,
        fitting=fitting,
        type_map=model_params["type_map"],
        atom_exclude_types=model_params.get("atom_exclude_types", []),
        pair_exclude_types=pair_exclude_types,
        preset_out_bias=preset_out_bias,
        data_stat_protect=data_stat_protect,
        use_compile=use_compile,
        enable_tf32=enable_tf32,
        bridging_method=bridging_method,
        spin=spin,
    )
    model.model_def_script = json.dumps(model_params_old)
    return model


def get_model(model_params: dict) -> Any:
    model_params = expand_bridging_method(model_params)
    model_type = model_params.get("type", "standard")
    if model_type == "standard":
        if "spin" in model_params:
            return get_spin_model(model_params)
        elif "use_srtab" in model_params:
            return get_zbl_model(model_params)
        else:
            return get_standard_model(model_params)
    elif model_type == "linear_ener":
        return get_linear_model(model_params)
    elif model_type in ("SeZM", "sezm", "DPA4", "dpa4"):
        if "spin" in model_params:
            return get_sezm_spin_model(model_params)
        return get_sezm_model(model_params)
    else:
        return BaseModel.get_class_by_type(model_type).get_model(model_params)


__all__ = [
    "BaseModel",
    "DOSModel",
    "DPModelCommon",
    "DPZBLModel",
    "DipoleModel",
    "EnergyModel",
    "FrozenModel",
    "LinearEnergyModel",
    "PolarModel",
    "PopulationModel",
    "SeZMModel",
    "SeZMNativeSpinModel",
    "SeZMPropertyModel",
    "SpinEnergyModel",
    "SpinModel",
    "get_model",
    "get_sezm_spin_model",
    "make_hessian_model",
    "make_model",
]