###########################################################################################
# The ASE Calculator for MACE
# Authors: Ilyes Batatia, David Kovacs
# This program is distributed under the MIT License (see MIT.md)
###########################################################################################

import logging
import collections

# pylint: disable=wrong-import-position
import os
from glob import glob
from pathlib import Path
from typing import List, Union

os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"

import numpy as np
import torch
from ase.calculators.calculator import Calculator, all_changes
from ase.stress import full_3x3_to_voigt_6_stress
from macer.externals.e3nn_legacy import o3

from macer.externals.mace_bundled import data
from macer.externals.mace_bundled.modules.utils import extract_invariant
from macer.externals.mace_bundled.tools import torch_geometric, torch_tools, utils
from macer.externals.mace_bundled.tools.compile import prepare
from macer.externals.mace_bundled.tools.scripts_utils import extract_model

try:
    from macer.externals.mace_bundled.cli.convert_e3nn_cueq import run as run_e3nn_to_cueq

    CUEQQ_AVAILABLE = True
except (ImportError, ModuleNotFoundError):
    CUEQQ_AVAILABLE = False
    run_e3nn_to_cueq = None

try:
    from macer.externals.mace_bundled.cli.convert_e3nn_oeq import run as run_e3nn_to_oeq

    OEQ_AVAILABLE = True
except (ImportError, ModuleNotFoundError):
    OEQ_AVAILABLE = False
    run_e3nn_to_oeq = None

try:
    import intel_extension_for_pytorch as ipex

    has_ipex = True
except ImportError:
    has_ipex = False


def get_model_dtype(model: torch.nn.Module) -> torch.dtype:
    """Get the dtype of the model"""
    mode_dtype = next(model.parameters()).dtype
    if mode_dtype == torch.float64:
        return "float64"
    if mode_dtype == torch.float32:
        return "float32"
    raise ValueError(f"Unknown dtype {mode_dtype}")


import sys
from contextlib import contextmanager

@contextmanager
def _mace_shim():
    """Shim to allow torch.load to find bundled mace and e3nn and their submodules"""
    import macer.externals.e3nn_legacy
    import macer.externals.mace_bundled
    
    # Store original modules
    originals = {}
    
    # We need to swap 'e3nn', 'mace' and all their submodules
    to_shim = {
        "e3nn": macer.externals.e3nn_legacy,
        "mace": macer.externals.mace_bundled
    }
    
    # Find all currently loaded submodules that we need to shim
    for prefix in list(to_shim.keys()):
        for mod_name in list(sys.modules.keys()):
            if mod_name == prefix or mod_name.startswith(prefix + "."):
                originals[mod_name] = sys.modules.get(mod_name)
    
    # Apply shims
    # First, the main modules
    sys.modules["e3nn"] = macer.externals.e3nn_legacy
    sys.modules["mace"] = macer.externals.mace_bundled
    
    # Then submodules. We need to map global e3nn.xxx to macer.externals.e3nn_legacy.xxx
    # However, since they are already vendorized, they might already exist in sys.modules
    # as macer.externals.e3nn_legacy.xxx
    for mod_name in list(sys.modules.keys()):
        if mod_name.startswith("macer.externals.e3nn_legacy"):
            target_name = mod_name.replace("macer.externals.e3nn_legacy", "e3nn")
            if target_name not in originals: # Don't overwrite if we already saved it
                 originals[target_name] = sys.modules.get(target_name)
            sys.modules[target_name] = sys.modules[mod_name]
        elif mod_name.startswith("macer.externals.mace_bundled"):
            target_name = mod_name.replace("macer.externals.mace_bundled", "mace")
            if target_name not in originals:
                originals[target_name] = sys.modules.get(target_name)
            sys.modules[target_name] = sys.modules[mod_name]
    
    try:
        yield
    finally:
        # Restore original modules
        # 1. Remove any 'e3nn' or 'mace' entries we might have added that weren't there
        for mod_name in list(sys.modules.keys()):
            if mod_name == "e3nn" or mod_name.startswith("e3nn.") or \
               mod_name == "mace" or mod_name.startswith("mace."):
                if mod_name not in originals:
                    sys.modules.pop(mod_name, None)
        
        # 2. Restore the original ones
        for mod_name, mod_obj in originals.items():
            if mod_obj is not None:
                sys.modules[mod_name] = mod_obj
            else:
                sys.modules.pop(mod_name, None)


class MACECalculator(Calculator):
    """MACE ASE Calculator
    args:
        model_paths: str, path to model or models if a committee is produced
                to make a committee use a wild card notation like mace_*.model
        device: str, device to run on (cuda or cpu or xpu)
        energy_units_to_eV: float, conversion factor from model energy units to eV
        length_units_to_A: float, conversion factor from model length units to Angstroms
        default_dtype: str, default dtype of model
        charges_key: str, Array field of atoms object where atomic charges are stored
        model_type: str, type of model to load
                    Options: [MACE, DipoleMACE, EnergyDipoleMACE]

    Dipoles are returned in units of Debye
    """

    def __init__(
        self,
        model_paths: Union[list, str, None] = None,
        models: Union[List[torch.nn.Module], torch.nn.Module, None] = None,
        device: str = "cpu",
        energy_units_to_eV: float = 1.0,
        length_units_to_A: float = 1.0,
        default_dtype="",
        charges_key="Qs",
        info_keys=None,
        arrays_keys=None,
        model_type="MACE",
        compile_mode=None,
        fullgraph=True,
        enable_cueq=False,
        enable_oeq=False,
        **kwargs,
    ):
        Calculator.__init__(self, **kwargs)
        if enable_cueq or enable_oeq:
            assert model_type == "MACE", "CuEq only supports MACE models"
            if compile_mode is not None:
                logging.warning(
                    "CuEq or Oeq does not support torch.compile, setting compile_mode to None"
                )
                compile_mode = None
        if enable_cueq and enable_oeq:
            raise ValueError(
                "CuEq and OEq cannot be used together, please choose one of them"
            )
        if enable_cueq and not CUEQQ_AVAILABLE:
            raise ImportError(
                "cuequivariance is not installed so CuEq acceleration cannot be used"
            )
        if enable_oeq and not OEQ_AVAILABLE:
            raise ImportError(
                "openequivariance is not installed so OEq acceleration cannot be used"
            )
        if "model_path" in kwargs:
            deprecation_message = (
                "'model_path' argument is deprecated, please use 'model_paths'"
            )
            if model_paths is None:
                logging.warning(f"{deprecation_message} in the future.")
                model_paths = kwargs["model_path"]
            else:
                raise ValueError(
                    f"both 'model_path' and 'model_paths' given, {deprecation_message} only."
                )

        if (model_paths is None) == (models is None):
            raise ValueError(
                "Exactly one of 'model_paths' or 'models' must be provided"
            )

        self.results = {}
        if info_keys is None:
            info_keys = {"total_spin": "spin", "total_charge": "charge"}
        if arrays_keys is None:
            arrays_keys = {}
        self.info_keys = info_keys
        self.arrays_keys = arrays_keys

        self.model_type = model_type
        self.compute_atomic_stresses = False

        if model_type == "MACE":
            self.implemented_properties = [
                "energy",
                "free_energy",
                "node_energy",
                "forces",
                "stress",
            ]
            if kwargs.get("compute_atomic_stresses", False):
                self.implemented_properties.extend(["stresses", "virials"])
                self.compute_atomic_stresses = True
        elif model_type == "DipoleMACE":
            self.implemented_properties = ["dipole"]
        elif model_type == "DipolePolarizabilityMACE":
            self.implemented_properties = [
                "charges",
                "dipole",
                "polarizability",
                "polarizability_sh",
            ]
        elif model_type == "EnergyDipoleMACE":
            self.implemented_properties = [
                "energy",
                "free_energy",
                "node_energy",
                "forces",
                "stress",
                "dipole",
            ]
        else:
            raise ValueError(
                f"Give a valid model_type: [MACE, DipoleMACE, DipolePolarizabilityMACE, EnergyDipoleMACE], {model_type} not supported"
            )

        if model_paths is not None:
            if isinstance(model_paths, str):
                # Find all models that satisfy the wildcard (e.g. mace_model_*.pt)
                model_paths_glob = glob(model_paths)

                if len(model_paths_glob) == 0:
                    raise ValueError(f"Couldn't find MACE model files: {model_paths}")

                model_paths = model_paths_glob
            elif isinstance(model_paths, Path):
                model_paths = [model_paths]

            if len(model_paths) == 0:
                raise ValueError("No mace file names supplied")
            self.num_models = len(model_paths)

            # Load models from files
            with _mace_shim():
                self.models = []
                for model_path in model_paths:
                    loaded = torch.load(f=model_path, map_location=device)
                    
                    if isinstance(loaded, (dict, collections.OrderedDict)):
                        if "model" in loaded:
                            # Standard MACE checkpoint case
                            m = loaded["model"]
                            # If the nested 'model' is also a state_dict, we have a problem, 
                            # but usually it's the model object.
                            if hasattr(m, "to"):
                                self.models.append(m)
                            else:
                                # Sometimes it's a state_dict even inside the "model" key
                                raise ValueError(f"Found 'model' key in {model_path}, but it's a {type(m)} and lacks '.to()'.")
                        elif "state_dict" in loaded:
                            # Some other training scripts use "state_dict"
                            raise ValueError(f"Found 'state_dict' in {model_path}, but MACECalculator needs the full model object. Please use a compiled .model file or a full-object .pt checkpoint.")
                        else:
                            # It's a dict but we don't know what's inside. 
                            # Check if the dict itself is a state_dict (all values are tensors)
                            is_state_dict = all(isinstance(v, torch.Tensor) for v in loaded.values())
                            if is_state_dict:
                                raise ValueError(f"Loaded file {model_path} appears to be a state_dict (OrderedDict of tensors). MACECalculator requires the full model object (architecture + weights).")
                            else:
                                available_keys = list(loaded.keys())
                                raise ValueError(f"Loaded dictionary from {model_path} has unknown keys: {available_keys}. Expected 'model' key.")
                    else:
                        # Loaded object is not a dict, check if it's a model
                        if hasattr(loaded, "to"):
                            self.models.append(loaded)
                        else:
                            raise ValueError(f"Loaded object from {model_path} is a {type(loaded)} and lacks '.to()' method.")


        elif models is not None:
            if not isinstance(models, list):
                models = [models]

            if len(models) == 0:
                raise ValueError("No models supplied")

            self.models = models
            self.num_models = len(models)

        if self.num_models > 1:
            print(f"Running committee mace with {self.num_models} models")

            if model_type in ["MACE", "EnergyDipoleMACE"]:
                self.implemented_properties.extend(
                    ["energies", "energy_var", "forces_comm", "stress_var"]
                )
            elif model_type == "DipoleMACE":
                self.implemented_properties.extend(["dipole_var"])

        if compile_mode is not None:
            print(f"Torch compile is enabled with mode: {compile_mode}")
            self.models = [
                torch.compile(
                    prepare(extract_model)(model=model, map_location=device),
                    mode=compile_mode,
                    fullgraph=fullgraph,
                )
                for model in self.models
            ]
            self.use_compile = True
        else:
            self.use_compile = False

        # Ensure all models are on the same device
        for model in self.models:
            model.to(device)

        if has_ipex and device == "xpu":
            for model in self.models:
                model = ipex.optimize(model)

        r_maxs = [model.r_max.cpu() for model in self.models]
        r_maxs = np.array(r_maxs)
        if not np.all(r_maxs == r_maxs[0]):
            raise ValueError(f"committee r_max are not all the same {' '.join(r_maxs)}")
        self.r_max = float(r_maxs[0])

        self.device = torch_tools.init_device(device)
        self.energy_units_to_eV = energy_units_to_eV
        self.length_units_to_A = length_units_to_A
        # Ensure atomic numbers are standard python ints to avoid np.int64 issues
        atomic_numbers = []
        for z in self.models[0].atomic_numbers:
            if isinstance(z, torch.Tensor):
                atomic_numbers.append(int(z.item()))
            else:
                atomic_numbers.append(int(z))
        self.z_table = utils.AtomicNumberTable(atomic_numbers)
        self.charges_key = charges_key

        try:
            self.available_heads: List[str] = self.models[0].heads  # type: ignore
        except AttributeError:
            self.available_heads = ["Default"]
        kwarg_head = kwargs.get("head", None)
        if kwarg_head is not None:
            self.head = kwarg_head
        elif len(self.available_heads) == 1:
            self.head = self.available_heads[0]
        else:
            self.head = [
                head for head in self.available_heads if head.lower() == "default"
            ]
            if len(self.head) == 0:
                raise ValueError(
                    "Head keyword was not provided, and no head in the model is 'default'. "
                    "Please provide a head keyword to specify the head you want to use. "
                    f"Available heads are: {self.available_heads}"
                )
            self.head = self.head[0]

        print("Using head", self.head, "out of", self.available_heads)

        model_dtype = get_model_dtype(self.models[0])

        if default_dtype == "":
            print(
                f"No dtype selected, switching to {model_dtype} to match model dtype."
            )
            default_dtype = model_dtype
        if model_dtype != default_dtype:
            print(
                f"Default dtype {default_dtype} does not match model dtype {model_dtype}, converting models to {default_dtype}."
            )
            if default_dtype == "float64":
                self.models = [model.double() for model in self.models]
            elif default_dtype == "float32":
                self.models = [model.float() for model in self.models]
        torch_tools.set_default_dtype(default_dtype)
        if enable_cueq:
            print("Converting models to CuEq for acceleration")
            self.models = [
                run_e3nn_to_cueq(model, device=device).to(device)
                for model in self.models
            ]
        if enable_oeq:
            print("Converting models to OEq for acceleration")
            self.models = [
                run_e3nn_to_oeq(model, device=device).to(device)
                for model in self.models
            ]
        for model in self.models:
            for param in model.parameters():
                param.requires_grad = False

    def _create_result_tensors(
        self, model_type: str, num_models: int, num_atoms: int
    ) -> dict:
        """
        Create tensors to store the results of the committee
        :param model_type: str, type of model to load
            Options: [MACE, DipoleMACE, EnergyDipoleMACE]
        :param num_models: int, number of models in the committee
        :return: tuple of torch tensors
        """
        dict_of_tensors = {}
        if model_type in ["MACE", "EnergyDipoleMACE"]:
            energies = torch.zeros(num_models, device=self.device)
            node_energy = torch.zeros(num_models, num_atoms, device=self.device)
            forces = torch.zeros(num_models, num_atoms, 3, device=self.device)
            stress = torch.zeros(num_models, 3, 3, device=self.device)
            dict_of_tensors.update(
                {
                    "energies": energies,
                    "node_energy": node_energy,
                    "forces": forces,
                    "stress": stress,
                }
            )
        if model_type in ["EnergyDipoleMACE", "DipoleMACE", "DipolePolarizabilityMACE"]:
            dipole = torch.zeros(num_models, 3, device=self.device)
            dict_of_tensors.update({"dipole": dipole})
        if model_type in ["DipolePolarizabilityMACE"]:
            charges = torch.zeros(num_models, num_atoms, device=self.device)
            polarizability = torch.zeros(num_models, 3, 3, device=self.device)
            polarizability_sh = torch.zeros(num_models, 6, device=self.device)
            dict_of_tensors.update(
                {
                    "charges": charges,
                    "polarizability": polarizability,
                    "polarizability_sh": polarizability_sh,
                }
            )
        return dict_of_tensors

    def _atoms_to_batch(self, atoms):
        self.arrays_keys.update({self.charges_key: "charges"})
        keyspec = data.KeySpecification(
            info_keys=self.info_keys, arrays_keys=self.arrays_keys
        )
        config = data.config_from_atoms(
            atoms, key_specification=keyspec, head_name=self.head
        )
        data_loader = torch_geometric.dataloader.DataLoader(
            dataset=[
                data.AtomicData.from_config(
                    config,
                    z_table=self.z_table,
                    cutoff=self.r_max,
                    heads=self.available_heads,
                )
            ],
            batch_size=1,
            shuffle=False,
            drop_last=False,
        )
        batch = next(iter(data_loader)).to(self.device)
        return batch

    def _clone_batch(self, batch):
        batch_clone = batch.clone()
        if self.use_compile:
            batch_clone["node_attrs"].requires_grad_(True)
            batch_clone["positions"].requires_grad_(True)
        return batch_clone

    # pylint: disable=dangerous-default-value
    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
        """
        Calculate properties.
        :param atoms: ase.Atoms object
        :param properties: [str], properties to be computed, used by ASE internally
        :param system_changes: [str], system changes since last calculation, used by ASE internally
        :return:
        """
        # call to base-class to set atoms attribute
        Calculator.calculate(self, atoms)

        batch_base = self._atoms_to_batch(atoms)

        if self.model_type in ["MACE", "EnergyDipoleMACE"]:
            batch = self._clone_batch(batch_base)
            node_heads = batch["head"][batch["batch"]]
            num_atoms_arange = torch.arange(batch["positions"].shape[0])
            node_e0 = self.models[0].atomic_energies_fn(batch["node_attrs"])[
                num_atoms_arange, node_heads
            ]
            compute_stress = not self.use_compile
        else:
            compute_stress = False

        ret_tensors = self._create_result_tensors(
            self.model_type, self.num_models, len(atoms)
        )
        for i, model in enumerate(self.models):
            batch = self._clone_batch(batch_base)
            out = model(
                batch.to_dict(),
                compute_stress=compute_stress,
                training=self.use_compile,
                compute_edge_forces=self.compute_atomic_stresses,
                compute_atomic_stresses=self.compute_atomic_stresses,
            )
            if self.model_type in ["MACE", "EnergyDipoleMACE"]:
                ret_tensors["energies"][i] = out["energy"].detach()
                ret_tensors["node_energy"][i] = (out["node_energy"] - node_e0).detach()
                ret_tensors["forces"][i] = out["forces"].detach()
                if out["stress"] is not None:
                    ret_tensors["stress"][i] = out["stress"].detach()
            if self.model_type in [
                "DipoleMACE",
                "EnergyDipoleMACE",
                "DipolePolarizabilityMACE",
            ]:
                ret_tensors["dipole"][i] = out["dipole"].detach()
            if self.model_type == "DipolePolarizabilityMACE":
                ret_tensors["charges"][i] = out["charges"].detach()
                ret_tensors["polarizability"][i] = out["polarizability"].detach()
                ret_tensors["polarizability_sh"][i] = out["polarizability_sh"].detach()
            if self.model_type in ["MACE"]:
                if out["atomic_stresses"] is not None:
                    ret_tensors.setdefault("atomic_stresses", []).append(
                        out["atomic_stresses"].detach()
                    )
                if out["atomic_virials"] is not None:
                    ret_tensors.setdefault("atomic_virials", []).append(
                        out["atomic_virials"].detach()
                    )

        self.results = {}
        if self.model_type in ["MACE", "EnergyDipoleMACE"]:
            self.results["energy"] = (
                torch.mean(ret_tensors["energies"], dim=0).cpu().item()
                * self.energy_units_to_eV
            )
            self.results["free_energy"] = self.results["energy"]
            self.results["node_energy"] = (
                torch.mean(ret_tensors["node_energy"], dim=0).cpu().numpy()
            )
            self.results["forces"] = (
                torch.mean(ret_tensors["forces"], dim=0).cpu().numpy()
                * self.energy_units_to_eV
                / self.length_units_to_A
            )
            if self.num_models > 1:
                self.results["energies"] = (
                    ret_tensors["energies"].cpu().numpy() * self.energy_units_to_eV
                )
                self.results["energy_var"] = (
                    torch.var(ret_tensors["energies"], dim=0, unbiased=False)
                    .cpu()
                    .item()
                    * self.energy_units_to_eV
                )
                self.results["forces_comm"] = (
                    ret_tensors["forces"].cpu().numpy()
                    * self.energy_units_to_eV
                    / self.length_units_to_A
                )
            if out["stress"] is not None:
                self.results["stress"] = full_3x3_to_voigt_6_stress(
                    torch.mean(ret_tensors["stress"], dim=0).cpu().numpy()
                    * self.energy_units_to_eV
                    / self.length_units_to_A**3
                )
                if self.num_models > 1:
                    self.results["stress_var"] = full_3x3_to_voigt_6_stress(
                        torch.var(ret_tensors["stress"], dim=0, unbiased=False)
                        .cpu()
                        .numpy()
                        * self.energy_units_to_eV
                        / self.length_units_to_A**3
                    )
            if "atomic_stresses" in ret_tensors:
                self.results["stresses"] = (
                    torch.mean(torch.stack(ret_tensors["atomic_stresses"]), dim=0)
                    .cpu()
                    .numpy()
                    * self.energy_units_to_eV
                    / self.length_units_to_A**3
                )
            if "atomic_virials" in ret_tensors:
                self.results["virials"] = (
                    torch.mean(torch.stack(ret_tensors["atomic_virials"]), dim=0)
                    .cpu()
                    .numpy()
                    * self.energy_units_to_eV
                )
        if self.model_type in [
            "DipoleMACE",
            "EnergyDipoleMACE",
            "DipolePolarizabilityMACE",
        ]:
            self.results["dipole"] = (
                torch.mean(ret_tensors["dipole"], dim=0).cpu().numpy()
            )
            if self.num_models > 1:
                self.results["dipole_var"] = (
                    torch.var(ret_tensors["dipole"], dim=0, unbiased=False)
                    .cpu()
                    .numpy()
                )
        if self.model_type in [
            "DipolePolarizabilityMACE",
        ]:
            self.results["charges"] = (
                torch.mean(ret_tensors["charges"], dim=0).cpu().numpy()
            )
            self.results["polarizability"] = (
                torch.mean(ret_tensors["polarizability"], dim=0).cpu().numpy()
            )
            self.results["polarizability_sh"] = (
                torch.mean(ret_tensors["polarizability_sh"], dim=0).cpu().numpy()
            )

    def get_dielectric_derivatives(self, atoms=None):
        if atoms is None and self.atoms is None:
            raise ValueError("atoms not set")
        if atoms is None:
            atoms = self.atoms
        if self.model_type not in ["DipoleMACE", "DipolePolarizabilityMACE"]:
            raise NotImplementedError(
                "Only implemented for DipoleMACE or DipolePolarizabilityMACE models"
            )
        batch = self._atoms_to_batch(atoms)
        outputs = [
            model(
                self._clone_batch(batch).to_dict(),
                compute_dielectric_derivatives=True,
                training=self.use_compile,
            )
            for model in self.models
        ]
        dipole_derivatives = [
            output["dmu_dr"].clone().detach().cpu().numpy() for output in outputs
        ]
        if self.models[0].use_polarizability:
            polarizability_derivatives = [
                output["dalpha_dr"].clone().detach().cpu().numpy() for output in outputs
            ]
            if self.num_models == 1:
                dipole_derivatives = dipole_derivatives[0]
                polarizability_derivatives = polarizability_derivatives[0]
            del outputs, batch, atoms
            return dipole_derivatives, polarizability_derivatives
        if self.num_models == 1:
            return dipole_derivatives[0]
        del outputs, batch, atoms
        return dipole_derivatives

    def get_hessian(self, atoms=None):
        if atoms is None and self.atoms is None:
            raise ValueError("atoms not set")
        if atoms is None:
            atoms = self.atoms
        if self.model_type != "MACE":
            raise NotImplementedError("Only implemented for MACE models")
        batch = self._atoms_to_batch(atoms)
        hessians = [
            model(
                self._clone_batch(batch).to_dict(),
                compute_hessian=True,
                compute_stress=False,
                training=self.use_compile,
            )["hessian"]
            for model in self.models
        ]
        hessians = [hessian.detach().cpu().numpy() for hessian in hessians]
        if self.num_models == 1:
            return hessians[0]
        return hessians

    def get_descriptors(self, atoms=None, invariants_only=True, num_layers=-1):
        """Extracts the descriptors from MACE model.
        :param atoms: ase.Atoms object
        :param invariants_only: bool, if True only the invariant descriptors are returned
        :param num_layers: int, number of layers to extract descriptors from, if -1 all layers are used
        :return: np.ndarray (num_atoms, num_interactions, invariant_features) of invariant descriptors if num_models is 1 or list[np.ndarray] otherwise
        """
        if atoms is None and self.atoms is None:
            raise ValueError("atoms not set")
        if atoms is None:
            atoms = self.atoms
        if self.model_type != "MACE":
            raise NotImplementedError("Only implemented for MACE models")
        num_interactions = int(self.models[0].num_interactions)
        if num_layers == -1:
            num_layers = num_interactions
        batch = self._atoms_to_batch(atoms)
        descriptors = [model(batch.to_dict())["node_feats"] for model in self.models]

        irreps_out = o3.Irreps(str(self.models[0].products[0].linear.irreps_out))
        l_max = irreps_out.lmax
        num_invariant_features = irreps_out.dim // (l_max + 1) ** 2
        per_layer_features = [irreps_out.dim for _ in range(num_interactions)]
        per_layer_features[-1] = (
            num_invariant_features  # Equivariant features not created for the last layer
        )

        if invariants_only:
            descriptors = [
                extract_invariant(
                    descriptor,
                    num_layers=num_layers,
                    num_features=num_invariant_features,
                    l_max=l_max,
                )
                for descriptor in descriptors
            ]
        to_keep = np.sum(per_layer_features[:num_layers])
        descriptors = [
            descriptor[:, :to_keep].detach().cpu().numpy() for descriptor in descriptors
        ]

        if self.num_models == 1:
            return descriptors[0]
        return descriptors
