import functools
import itertools
import collections
from abc import ABC, abstractmethod
import copy
import typing as tp
import math

import torch
import torch.nn as nn

from . import constants
from . import mesher
from . import particles
from . import utils
from .mesher import BaseMesh


# Подумать над изменением логики.
def kronecker_product(matrices: list) -> torch.Tensor:
    """Computes the Kronecker product of a list of matrices."""
    return functools.reduce(torch.kron, matrices)


def create_operator(system_particles: list, target_idx: int, matrix: torch.Tensor) -> torch.Tensor:
    """Creates an operator acting on the target particle with identity elsewhere."""
    operator = []
    for i, p in enumerate(system_particles):
        operator.append(matrix if i == target_idx else p.identity)
    return kronecker_product(operator)


def scalar_tensor_multiplication(
    tensor_components_A: torch.Tensor,
    tensor_components_B: torch.Tensor,
    transformation_matrix: torch.Tensor
) -> torch.Tensor:
    """
    Computes the scalar product of two tensor components after applying a transformation.
    Parameters:
        tensor_components_A (torch.Tensor): Input tensor components with shape [..., 3, K, K].
        tensor_components_B (torch.Tensor): Input tensor components with shape [..., 3, K, K].
        transformation_matrix (torch.Tensor): Transformation matrix with shape [..., 3, 3].
    Returns:
        torch.Tensor: Scalar product with shape [..., K, K].
    """
    return torch.einsum(
        '...ij, jnl, ilm->...nm',
        transformation_matrix,
        tensor_components_A,
        tensor_components_B
    )


def transform_tensor_components(tensor_components: torch.Tensor, transformation_matrix: torch.Tensor) -> torch.Tensor:
    """
    Applies a matrix transformation to a collection of tensor components.
    Parameters:
        tensor_components (torch.Tensor): Input tensor components with shape [..., 3, K, K]
            where 3 represents different components (e.g., x,y,z) and K is the system dimension. For example,
            [Sx, Sy, Sz]
        transformation_matrix (torch.Tensor): Transformation matrix with shape [..., 3, 3]. For example, g - tensor
    Returns:
        torch.Tensor: Transformed tensor components with shape [..., 3, K, K]
    """
    return torch.einsum('...ij,...jkl->...ikl', transformation_matrix, tensor_components)


# Возможно, стоит переделать логику работы расчёта тензоров через тенорное произведение. Сделать отдельный тип данных.
# Сейчас каждый спин даёт матрицу [K, K] и расчёт взаимодействией не оптимальный
def init_tensor(
        components: tp.Union[torch.Tensor, tp.Sequence[float], float],  device: torch.device, dtype: torch.dtype
):
    if isinstance(components, torch.Tensor):
        tensor = components.to(device=device, dtype=dtype)
        if tensor.ndim:
            if tensor.shape[-1] == 3:
                return tensor

            elif tensor.shape[-1] == 2:
                axis_val, z_val = tensor[0], tensor[1]
                return torch.stack([axis_val, axis_val, z_val], dim=-1)

            elif tensor.shape[-1] == 1:
                tensor = tensor.squeeze(-1)
                tensor = torch.stack([tensor, tensor, tensor], dim=-1)
                return tensor
            else:
                raise ValueError(f"Tensor must have shape [..., 3] or [1] or [2], got {tensor.shape}")

        else:
            tensor = torch.stack([tensor, tensor, tensor], dim=-1)
            return tensor

    elif isinstance(components, (list, tuple)):
        if len(components) == 1:
            value = components[0]
            return torch.full((3,), value, device=device, dtype=dtype)
        elif len(components) == 2:
            axis_val, z_val = components
            return torch.tensor([axis_val, axis_val, z_val], device=device, dtype=dtype)
        elif len(components) == 3:
            return torch.tensor(components, device=device, dtype=dtype)
        else:
            raise ValueError(f"List must have 1, 2, or 3 elements, got {len(components)}")

    elif isinstance(components, (int, float)):
        return torch.full((3,), components, device=device, dtype=dtype)

    else:
        raise TypeError(f"components must be a tensor, list, tuple, or scalar, got {type(components)}")


def init_de_tensor(
        components: tp.Union[torch.Tensor, tp.Sequence[float], float],  device: torch.device, dtype: torch.dtype
):
    if isinstance(components, torch.Tensor):
        tensor = components.to(device=device, dtype=dtype)
        if tensor.shape:
            if tensor.shape[-1] == 3:
                return tensor

            elif tensor.shape[-1] == 2:
                D, E = tensor[..., 0], tensor[..., 1]
                Dx = - D / 3 + E
                Dy = - D / 3 - E
                Dz = 2 * D / 3
                return torch.stack([Dx, Dy, Dz], dim=-1)

            elif tensor.shape[-1] == 1:
                tensor = tensor.squeeze(-1)
                Dx = - tensor / 3
                Dy = - tensor / 3
                Dz = 2 * tensor / 3
                return torch.stack([Dx, Dy, Dz], dim=-1)

            else:
                raise ValueError(f"Tensor must have shape [..., 3] or [1] or [2], got {tensor.shape}")
        else:
            Dx = - tensor / 3
            Dy = - tensor / 3
            Dz = 2 * tensor / 3
            return torch.stack([Dx, Dy, Dz], dim=-1)

    elif isinstance(components, (list, tuple)):
        if len(components) == 1:
            value = components[0]
            Dx = - value / 3
            Dy = - value / 3
            Dz = 2 * value / 3
            return torch.tensor([Dx, Dy, Dz], device=device, dtype=dtype)

        elif len(components) == 2:
            D, E = components[0], components[1]
            Dx = - D / 3 + E
            Dy = - D / 3 - E
            Dz = 2 * D / 3
            return torch.tensor([Dx, Dy, Dz], device=device, dtype=dtype)

        elif len(components) == 3:
            return torch.tensor(components, device=device, dtype=dtype)
        else:
            raise ValueError(f"List must have 1, 2, or 3 elements, got {len(components)}")

    elif isinstance(components, (int, float)):
        return torch.full((3,), components, device=device, dtype=dtype)

    else:
        raise TypeError(f"components must be a tensor, list, tuple, or scalar, got {type(components)}")


class BaseInteraction(nn.Module, ABC):
    @property
    @abstractmethod
    def tensor(self):
        pass

    @property
    def components(self):
        return self.tensor

    @property
    def strain(self):
        return self.strained_tensor

    @property
    def frame(self):
        return None

    @property
    @abstractmethod
    def strained_tensor(self):
        pass

    @property
    @abstractmethod
    def config_shape(self):
        pass

    @property
    @abstractmethod
    def strain_correlation(self):
        """
        :return: The correlation matrix of strain parameters
        """
        pass

    def __len__(self):
        return len(self.components)

    def __repr__(self):
        is_batched = hasattr(self.components, 'shape') and len(self.components.shape) > 1

        if is_batched:
            batch_size = self.components.shape[0]
            lines = [f"BATCHED (batch_size={batch_size}) - showing first instance:"]

            first_components = self.components.flatten(0, -2)[0]
            if hasattr(first_components, 'tolist'):
                components_str = [f"{val:.2e}" if abs(val) >= 1e4 else f"{val:.4f}"
                                  for val in first_components.tolist()]
            else:
                components_str = [f"{val:.2e}" if abs(val) >= 1e4 else f"{val:.4f}"
                                  for val in first_components]

            lines.append(f"Principal values: [{', '.join(components_str)}]")

            # Handle batched frame
            if self.frame is not None:
                first_frame = self.frame.flatten(0, -2)[0]
                if hasattr(first_frame, 'tolist'):
                    frame_vals = first_frame.tolist()
                else:
                    frame_vals = first_frame

                if len(frame_vals) == 3:  # Euler angles
                    frame_str = f"[α={frame_vals[0]:.3f}, β={frame_vals[1]:.3f}, γ={frame_vals[2]:.3f}] rad"
                    lines.append(f"Frame (Euler angles): {frame_str}")
                else:
                    lines.append(f"Frame: {frame_vals}")
            else:
                lines.append("Frame: None")

            if self.strain is not None:
                first_strain = self.strain.flatten(0, -2)[0]
                if hasattr(first_strain, 'tolist'):
                    strain_vals = first_strain.tolist()
                    strain_str = [f"{val:.2e}" if abs(val) >= 1e4 else f"{val:.4f}"
                                  for val in strain_vals]
                    lines.append(f"Strain: [{', '.join(strain_str)}]")
                else:
                    lines.append(f"Strain: {first_strain}")
            else:
                lines.append("Strain: None")

        else:
            if hasattr(self.components, 'tolist'):
                components_str = [f"{val:.2e}" if abs(val) >= 1e4 else f"{val:.4f}"
                                  for val in self.components.tolist()]
            else:
                components_str = [f"{val:.2e}" if abs(val) >= 1e4 else f"{val:.4f}"
                                  for val in self.components]

            lines = [
                f"Principal values: [{', '.join(components_str)}]",
            ]

            if self.frame is not None:
                if hasattr(self.frame, 'tolist'):
                    frame_vals = self.frame.tolist()
                else:
                    frame_vals = self.frame

                if len(frame_vals) == 3:
                    frame_str = f"[α={frame_vals[0]:.3f}, β={frame_vals[1]:.3f}, γ={frame_vals[2]:.3f}] rad"
                    lines.append(f"Frame (Euler angles): {frame_str}")
                else:
                    lines.append(f"Frame: {frame_vals}")
            else:
                lines.append("Frame: Identity (no rotation)")

            if self.strain is not None:
                if hasattr(self.strain, 'tolist'):
                    strain_vals = self.strain.tolist()
                    strain_str = [f"{val:.2e}" if abs(val) >= 1e4 else f"{val:.4f}"
                                  for val in strain_vals]
                    lines.append(f"Strain: [{', '.join(strain_str)}]")
                else:
                    lines.append(f"Strain: {self.strain}")
            else:
                lines.append("Strain: None")

        return '\n'.join(lines)


class Interaction(BaseInteraction):
    def __init__(self, components: tp.Union[torch.Tensor, tp.Sequence, float],
                 frame: tp.Optional[tp.Union[torch.Tensor, tp.Sequence[float]]] = None,
                 strain: tp.Optional[tp.Union[torch.Tensor, tp.Sequence, float]] = None,
                 device=torch.device("cpu"), dtype=torch.float32):
        """
        :param components:
        torch.Tensor | Sequence[float] | float
            The tensor components, provided in one of the following forms:
              - A scalar (for isotropic interaction).
              - A sequence of two values (axial and z components).
              - A sequence of three values (principal components).
        The possible units are [T, Hz, dimensionless]

        :param frame:
        torch.Tensor | Sequence[float] optional
            Orientation of the tensor. Can be provided as:
              - A 1D tensor of shape (3,) representing Euler angles in ZYZ' convention.
              - A 2D tensor of shape (3, 3) representing a rotation matrix.
            Default is `None`, meaning lab frame.

        :param strain:
        torch.Tensor| Sequence[float] | float, optional
            Parameters describing interaction broadening or distribution.
            Default is `None`.

        If the batched paradigm is used then only torch.Tensors with shape [..., 3] are acceptable.

        :param device:

        :param dtype:
        """
        super().__init__()
        self.register_buffer("_components", init_tensor(components, device=device, dtype=dtype))
        self.shape = self._components.shape
        batch_shape = self._components.shape[:-1]

        self._construct_rot_matrix(frame, batch_shape, device=device, dtype=dtype)

        _strain = init_tensor(strain, device=device, dtype=dtype) if strain is not None else None
        self.register_buffer("_strain", _strain)

        if (self._strain is not None) and (self._strain.shape != self.shape):
            raise ValueError("The strain shape must be equal to shape of the initial components."
                             "Please point it as x, y, z components")

        _strain_correlation = torch.eye(3, device=device, dtype=dtype)
        self.register_buffer("_strain_correlation", _strain_correlation)
        self.to(device)
        self.to(dtype)

    def _construct_rot_matrix(
            self, frame: tp.Optional[tp.Union[torch.Tensor, tp.Sequence[float]]], batch_shape,
            device: torch.device,
            dtype: torch.dtype
    ):
        if frame is None:
            _frame = torch.zeros((*batch_shape, 3), device=device, dtype=dtype)  # alpha, beta, gamma
            _rot_matrix = self.euler_to_rotmat(_frame).to(self.components.dtype)

        else:
            if isinstance(frame, torch.Tensor):
                if frame.shape[-2:] == (3, 3) and not batch_shape:
                    _frame = utils.rotation_matrix_to_euler_angles(frame)
                    _rot_matrix = frame.to(self.components.dtype)

                elif frame.shape == (*batch_shape, 3):
                    _frame = frame.to(dtype)
                    _rot_matrix = self.euler_to_rotmat(_frame).to(self.components.dtype)

                else:
                    raise ValueError(
                        "frame must be either:\n"
                        "  • None (→ identity rotation),\n"
                        "  • a tensor of Euler angles with shape batch×3,\n"
                        "  • or a tensor of rotation matrices with shape batch×3×3."
                    )
            elif isinstance(frame, collections.abc.Sequence):
                if len(frame) != 3:
                    raise ValueError("frame must have exactly 3 values")
                _frame = torch.tensor(frame, dtype=dtype, device=device)
                _rot_matrix = self.euler_to_rotmat(_frame).to(self.components.dtype)
            else:
                raise ValueError("frame must be a Sequence of 3 values, a torch.Tensor, or None.")


        self.register_buffer("_frame", _frame)
        self.register_buffer("_rot_matrix", _rot_matrix)

    def euler_to_rotmat(self, euler_angles: torch.Tensor):
        return utils.euler_angles_to_matrix(euler_angles)

    def _tensor(self):
        """
        :return: the tensor in the spin system axis
        the shape of the returned tensor is [..., 3, 3]
        """
        return utils.apply_single_rotation(self._rot_matrix, torch.diag_embed(self.components))

    def _strained_tensor(self) -> tp.Optional[torch.Tensor]:
        """
        :return: return the None or the tensor with the shape [..., 3, 3, 3] or None
        """
        if self._strain is None:
            return None
        else:
            return self._strain.unsqueeze(-1).unsqueeze(-1) *\
               torch.einsum("...ik, ...jk->...kij", self._rot_matrix, self._rot_matrix)

    @property
    def tensor(self):
        """
        :return: the full tensor of interaction with shape [..., 3, 3] with applied rotation
        """
        return self._tensor()

    @property
    def strain(self):
        """
        :return: None or tensor with shape [..., 3]
        """
        return self._strain

    @property
    def strained_tensor(self):
        """
        :return: None or the tensor with shape [...., 3, 3, 3], where first '3' is x, y, z components
        """
        return self._strained_tensor()

    @property
    def config_shape(self):
        return self.shape[:-1]

    @property
    def components(self):
        """
        :return: tensor with shape [..., 3] - the principle components of a tensor
        """
        return self._components

    @property
    def frame(self):
        """
        :return: angles with ZYZ' notation.
        """
        return self._frame

    @frame.setter
    def frame(self, frame):
        if frame is None:
            self._frame = torch.tensor(
                [0.0, 0.0, 0.0], device=self.components.device, dtype=self.components.dtype
            )  # alpha, beta, gamma
        self._rot_matrix = self.euler_to_rotmat(self._frame)

    @property
    def strain_correlation(self):
        """
        In some cases the components of the interaction can correlate.
        To implement this correlation the strain_correlation matrix is used. For example, in the case of D/E interaction
        strain_correlation = [[-1/3, 1], [-1/3, -1], [2/3, 0]] - the matrix of trnasformation of Dx, Dy, Dz to D and E
        :return:
        """
        return self._strain_correlation

    @strain_correlation.setter
    def strain_correlation(self, correlation_matrix: torch.Tensor):
        self._strain_correlation = correlation_matrix

    def __add__(self, other):
        if not isinstance(other, Interaction):
            raise TypeError("Can only add Interaction objects together")

        if torch.allclose(self.frame, other.frame, atol=1e-6):
            new_frame = self.frame.clone()
            new_components = self.components + other.components
        else:
            tensor_self = self._tensor()
            tensor_other = other._tensor()  # [..., 3, 3]
            combined_tensor = tensor_self + tensor_other

            eigenvalues, eigenvectors = torch.linalg.eigh(combined_tensor)

            sorted_indices = torch.argsort(eigenvalues, dim=-1, descending=True)
            new_components = torch.gather(eigenvalues, -1, sorted_indices)

            new_rot_matrix = eigenvectors[..., sorted_indices].transpose(-2, -1)
            new_frame = utils.rotation_matrix_to_euler_angles(new_rot_matrix)

        if self.strain is not None and other.strain is not None:
            new_strain = self.strain + other.strain
            correlation_matrix = torch.cat((self.strain_correlation, other.strain_correlation), dim=-1)
            #  It is suggested that D,E are not correlated

        elif self.strain is not None:
            new_strain = self.strain.clone()
            correlation_matrix = self.strain_correlation.clone()

        elif other.strain is not None:
            new_strain = other.strain.clone()
            correlation_matrix = other.strain_correlation.clone()

        else:
            new_strain = None
            correlation_matrix = torch.eye(3, device=new_components.device, dtype=new_components.dtype)

        interaction = Interaction(
            components=new_components,
            frame=new_frame,
            strain=new_strain,
            device=self.components.device,
            dtype=self.components.dtype
        )
        interaction.strain_correlation = correlation_matrix
        return interaction


class DEInteraction(Interaction):
    def __init__(self, components: torch.Tensor,
                 frame: torch.Tensor = None, strain: torch.Tensor = None,
                 device=torch.device("cpu"), dtype=torch.float32):
        """
        DEInteraction is given by two components D and E. To transform to x, y, z components the next equation is used:
        Dx = -D * 1/3 + E
        Dy = -D * 1/3 - E
        Dz = D * 2/3

                Note on DE Interaction vs. Simple Interaction
        The DE Interaction is equivalent to simple Interaction in terms of components when
        the trace of the tensor equals zero,
        but they are not equivalent in terms of strains.
        In DE Interaction, the D and E components (or only D) have a distribution,
        whereas in simple interaction, the components Dx, Dy, and Dz are distributed.

        :param components:
        torch.Tensor | Sequence[float] | float
            The tensor components, provided in one of the following forms:
              - A scalar. It is only D value.
              - A sequence of two values (D and E values).
        The possible units are [T, Hz, dimensionless]

        :param frame:
        torch.Tensor | Sequence[float] optional
            Orientation of the tensor. Can be provided as:
              - A 1D tensor of shape (3,) representing Euler angles in ZYZ' convention.
              - A 2D tensor of shape (3, 3) representing a rotation matrix.
            Default is `None`, meaning lab frame.

        :param strain:
        torch.Tensor| Sequence[float] | float, optional
            Parameters describing interaction broadening or distribution.
            Default is `None`.

        :param device: device to compute (cpu / gpu)

        :param dtype: float32 / float64
        """
        components = init_de_tensor(components, device, dtype)
        strain = init_de_tensor(strain, device=device, dtype=dtype) if strain is not None else None

        super().__init__(components, frame, strain, device, dtype)
        self._strain_correlation = torch.tensor([[-1 / 3, 1], [-1 / 3, -1], [2 / 3, 0]], device=device, dtype=dtype)


class MultiOrientedInteraction(BaseInteraction):
    def __init__(self, oriented_tensor, strained_tensor, config_shape, strain_correlation: torch.Tensor,
                 device=torch.device("cpu")):
        super().__init__()
        self.register_buffer("_oriented_tensor", oriented_tensor)
        self.register_buffer("_strained_tensor", strained_tensor)
        self.register_buffer("_strain_correlation", strain_correlation)
        self._config_shape = config_shape

    @property
    def tensor(self):
        return self._oriented_tensor

    @property
    def strained_tensor(self):
        return self._strained_tensor

    @property
    def config_shape(self):
        return self.tensor.shape[:-2]

    @property
    def strain_correlation(self):
        return self._strain_correlation

    @strain_correlation.setter
    def strain_correlation(self, correlation_matrix: torch.Tensor):
        self._strain_correlation = correlation_matrix


class SpinSystem(nn.Module):
    """Represents a spin system with electrons, nuclei, and interactions."""
    def __init__(self, electrons: tp.Union[list[particles.Electron], list[float]],
                 g_tensors: list[BaseInteraction],
                 nuclei: tp.Optional[tp.Union[list[particles.Nucleus], list[str]]] = None,
                 electron_nuclei: list[tuple[int, int, BaseInteraction]] | None = None,
                 electron_electron: list[tuple[int, int, BaseInteraction]] | None = None,
                 nuclei_nuclei: list[tuple[int, int, BaseInteraction]] | None = None,
                 device=torch.device("cpu"), dtype: torch.dtype = torch.float32):

        """
        :param electrons:
        list[Electron] | list[float]
            Electron spins in the system. Can be specified as:
              - A list of `Electron` particle instances.
              - A list of spin quantum numbers (e.g., [0.5, 1.0]).

        :param g_tensors:
        list[BaseInteraction]
            g-tensors corresponding to each electron in `electrons`.
            Each element must be an instance of `BaseInteraction` (e.g., `Interaction`).

        :param nuclei:
        list[Nucleus] | list[str], optional
            Nuclei in the system. Can be given as:
              - A list of `Nucleus` particle instances.
              - A list of isotope symbols (e.g., ["1H", "13C"]).
            Default is `None` (no nuclei).

        :param electron_nuclei:
        list[tuple[int, int, BaseInteraction]], optional
            Hyperfine interactions between electrons and nuclei.
            Each tuple is of the form (electron_index, nucleus_index, interaction_tensor).
            Default is `None`.

        :param electron_electron:
        list[tuple[int, int, BaseInteraction]], optional
            Dipolar or exchange interactions between pairs of electrons.
            Each tuple is of the form (electron_index, electron_index, interaction_tensor).
            Default is `None`.

        :param nuclei_nuclei:
        list[tuple[int, int, BaseInteraction]], optional
            Dipolar or J-coupling interactions between pairs of nuclei.
            Each tuple is of the form (nucleus_index, nucleus_index, interaction_tensor).
            Default is `None`

        :param device: device to compute (cpu / gpu)
        """
        super().__init__()

        complex_dtype = utils.float_to_complex_dtype(dtype)


        self.electrons = self._init_electrons(electrons)
        self.g_tensors = nn.ModuleList(g_tensors)
        self.nuclei = self._init_nuclei(nuclei) if nuclei else []

        self.electron_nuclei_interactions = nn.ModuleList()
        self.electron_electron_interactions = nn.ModuleList()
        self.nuclei_nuclei_interactions = nn.ModuleList()

        if len(self.g_tensors) != len(self.electrons):
            raise ValueError("the number of g tensors must be equal to the number of electrons")

        self.en_indices = []
        self.ee_indices = []
        self.nn_indices = []

        self._register_interactions(electron_nuclei, self.electron_nuclei_interactions, self.en_indices)
        self._register_interactions(electron_electron, self.electron_electron_interactions, self.ee_indices)
        self._register_interactions(nuclei_nuclei, self.nuclei_nuclei_interactions, self.nn_indices)

        _operator_cache = self._precompute_all_operators(device=device, complex_dtype=complex_dtype)
        self.register_buffer("_operator_cache_real",  _operator_cache.real)
        self.register_buffer("_operator_cache_imag", _operator_cache.imag)

        self.to(device)
        self.to(dtype)

    def _register_interactions(self,
                               interactions: list[tuple[int, int, BaseInteraction]] | None,
                               module_list: nn.ModuleList,
                               index_list: list):
        """Helper to register interactions and store indices"""
        if interactions:
            for idx1, idx2, interaction in interactions:
                module_list.append(interaction)
                index_list.append((idx1, idx2))

    def _init_electrons(self, electrons):
        return [particles.Electron(electron) if isinstance(electron, float) else electron for electron in electrons]

    def _init_nuclei(self, nuclei: tp.Union[list[particles.Nucleus], list[str]]):
        return [particles.Nucleus(nucleus) if isinstance(nucleus, str) else nucleus for nucleus in nuclei]

    @property
    def device(self):
        return self.operator_cache.device

    @property
    def dtype(self):
        return self.g_tensors[0].components.dtype

    @property
    def complex_dtype(self):
        return utils.float_to_complex_dtype(self.dtype)

    @property
    def config_shape(self) -> tp.Iterable:
        """shape of the tensor"""
        return self.g_tensors[0].config_shape

    @property
    def spin_system_dim(self) -> int:
        """Dimension of the system's Hilbert space."""
        return math.prod([int(2 * p.spin + 1) for p in itertools.chain(self.electrons, self.nuclei)])

    @property
    def electron_nuclei(self):
        return [(idx[0], idx[1], inter) for idx, inter in zip(self.en_indices, self.electron_nuclei_interactions)]

    @electron_nuclei.setter
    def electron_nuclei(self, interactions: list[tuple[int, int, BaseInteraction]] | None):
        self._register_interactions(interactions, self.electron_nuclei_interactions, self.en_indices)

    @property
    def electron_electron(self):
        return [(idx[0], idx[1], inter) for idx, inter in zip(self.ee_indices, self.electron_electron_interactions)]

    @electron_electron.setter
    def electron_electron(self, interactions: list[tuple[int, int, BaseInteraction]] | None):
        self._register_interactions(interactions, self.electron_electron_interactions, self.ee_indices)

    @property
    def nuclei_nuclei(self):
        return [(idx[0], idx[1], inter) for idx, inter in zip(self.nn_indices, self.nuclei_nuclei_interactions)]

    @nuclei_nuclei.setter
    def nuclei_nuclei(self, interactions: list[tuple[int, int, BaseInteraction]] | None):
        self._register_interactions(interactions, self.nuclei_nuclei_interactions, self.nn_indices)

    @property
    def operator_cache(self):
        return torch.complex(self._operator_cache_real, self._operator_cache_imag)

    #  Нужно передалать функцию, слишком много циклов. Можно улучшить
    def _precompute_all_operators(self, device: torch.device, complex_dtype: torch.dtype):
        """Precompute spin operators for all particles in the full Hilbert space."""
        particels = self.electrons + self.nuclei
        operator_cache = []
        for idx, p in enumerate(particels):
            axis_cache = []
            for axis, mat in zip(['x', 'y', 'z'], p.spin_matrices):
                operator = create_operator(particels, idx, mat)
                axis_cache.append(operator.to(device).to(complex_dtype))
            operator_cache.append(torch.stack(axis_cache, dim=-3))   # Сейчас каждый спин даёт матрицу [K, K] и
                                                                     # расчёт взаимодействией не оптимальный
        return torch.stack(operator_cache, dim=0)

    def get_electron_z_operator(self) -> torch.Tensor:
        """
        Compute the total Sz operator for all electron spins in the system.
        This method sums the individual Sz operators from each electron spin operator
        cache to produce the total spin projection operator along the z-axis.

        :return: The total Sz operator with shape [spin_dim, spin_dim], where spin_dim is the
        total dimension of the spin system Hilbert space.

        Examples
        --------
        For a system with one spin-1/2 electron:
            Returns a 2x2 matrix representing the Pauli Sz operator.
        For a system with two spin-1/2 electrons:
        Returns a 4x4 matrix representing the sum of both Sz operators.
        """
        return sum(self.operator_cache[idx][2, :, :] for idx in range(len(self.electrons)))

    def get_electron_squared_operator(self) -> torch.Tensor:
        """
        Compute the total S² operator for all electron spins in the system.
        This method calculates S² = Sx² + Sy² + Sz² by first summing the individual
        spin vector operators from each electron, then computing the dot product of
        the total spin vector with itself.

        :return: The total S² operator with shape [spin_dim, spin_dim], where spin_dim is the
        total dimension of the spin system Hilbert space.

        Examples
        --------
        For a system with two spin-1/2 electrons:
            Eigenvalues correspond to singlet (S=0) and triplet (S=1) states.
        """
        S_vector = sum(self.operator_cache[idx] for idx in range(len(self.electrons)))
        return torch.matmul(S_vector, S_vector).sum(dim=-3)

    def get_spin_multiplet_basis(self) -> torch.Tensor:
        """
        Compute eigenvector in the |S, M⟩ basis (total spin and projection basis).
        This method diagonalizes a combination of sS² and Sz operators to obtain
        eigenvectors ordered by total spin quantum number S, then by magnetic
        quantum number M (spin projection).

        :return: torch.Tensor
        Matrix of eigenvectors with shape [spin_dim, spin_dim], where each column
        represents an eigenstate. States are ordered in ascending order of S,
        then in ascending order of M within each S manifold.

        Examples
        --------
        For two spin-1/2 electrons:
            Returns basis ordered as: |S=0, M=0⟩, |S=1, M=-1⟩, |S=1, M=0⟩, |S=1, M=1⟩
        """
        C = self.get_electron_squared_operator() + 1j * self.get_electron_z_operator()
        values, vectors = torch.linalg.eig(C)
        sorting_key = values.real * (values.imag.abs().max() + 1) + values.imag
        indices = torch.argsort(sorting_key)
        return vectors[:, indices]

    def get_product_state_basis(self) -> torch.Tensor:
        """
        Return the identity matrix representing the computational product state basis.
        The product state basis is |ms1, ms2, ..., msk, is1, ..., ism⟩ where:
        - ms1, ms2, ..., msk are magnetic quantum numbers for electrons through k
        - is1, is2, ..., ism are magnetic quantum numbers for nuclei through m
        Each state corresponds to a definite spin projection for each particle.
        :return: torch.Tensor
        Identity matrix with shape [spin_system_dim, spin_system_dim]. The identity
        matrix indicates that the current representation is already in the product
        state basis.

        Examples
        --------
        For one spin-1/2 electron and one spin-1/2 nucleus:
            Basis states: |↑, ↑⟩, |↑, ↓⟩, |↓, ↑⟩, |↓, ↓⟩
        """
        return torch.eye(self.spin_system_dim, device=self.device, dtype=self.dtype)

    def get_total_projections(self) -> torch.Tensor:
        """
        Compute the total magnetic quantum number M for each product state.
        This method calculates M = Σmsi + Σisj (sum of all electron and nuclear
        spin projections) for each basis state in the product state representation.

        :return: torch.Tensor
        1D tensor with shape [spin_dim] containing the total spin projection
        (magnetic quantum number) for each product state basis vector.

        Examples
        --------
        For one spin-1/2 electron:
            Returns tensor([0.5, -0.5])

        For one spin-1/2 electron and one spin-1/2 nucleus:
            Returns tensor([1.0, 0.0, 0.0, -1.0])
            Corresponding to states: |↑↑⟩, |↑↓⟩, |↓↑⟩, |↓↓⟩

        For two spin-1/2 electrons:
            Returns tensor([1.0, 0.0, 0.0, -1.0])
            Corresponding to states: |↑↑⟩, |↑↓⟩, |↓↑⟩, |↓↓⟩
        """
        all_projections = []
        for particle in self.electrons:
            s = particle.spin
            num_vals = int(round(2 * s)) + 1
            m_vals = torch.linspace(-s, s, num_vals).tolist()
            all_projections.append(m_vals)

        for particle in self.nuclei:
            s = particle.spin
            num_vals = int(round(2 * s)) + 1
            m_vals = torch.linspace(-s, s, num_vals).tolist()
            all_projections.append(m_vals)

        total_projections = []
        for combination in itertools.product(*all_projections):
            total_projections.append(sum(combination))

        return torch.tensor(total_projections, dtype=self.dtype, device=self.device)

    def get_electron_projections(self) -> torch.Tensor:
        """
        Compute the electron-only spin projection for each product state.

        This method calculates Me = Σmsi (sum of electron spin projections only)
        for each basis state, ignoring nuclear spin contributions.

        :return: torch.Tensor
        1D tensor with shape [spin_dim] containing the total electron spin
        projection for each product state basis vector. Nuclear contributions
        are set to zero.

        Examples
        --------
        For one spin-1/2 electron:
            Returns tensor([0.5, -0.5])

        For one spin-1/2 electron and one spin-1/2 nucleus:
            Returns tensor([0.5, 0.5, -0.5, -0.5])
            Nuclear spins don't contribute, so we get: |↑_e,↑_n⟩, |↑_e,↓_n⟩, |↓_e,↑_n⟩, |↓_e,↓_n⟩

        For two spin-1/2 electrons:
            Returns tensor([1.0, 0.0, 0.0, -1.0])
            Corresponding to: |↑↑⟩, |↑↓⟩, |↓↑⟩, |↓↓⟩
        """
        all_projections = []
        for particle in self.electrons:
            s = particle.spin
            num_vals = int(round(2 * s)) + 1
            m_vals = torch.linspace(-s, s, num_vals).tolist()
            all_projections.append(m_vals)

        for particle in self.nuclei:
            s = particle.spin
            num_vals = int(round(2 * s)) + 1
            m_vals = [0] * num_vals
            all_projections.append(m_vals)

        total_projections = []
        for combination in itertools.product(*all_projections):
            total_projections.append(sum(combination))

        return torch.tensor(total_projections, dtype=self.dtype, device=self.device)

    def update(self,
               g_tensors: list[BaseInteraction] = None,
               electron_nuclei: list[tuple[int, int, BaseInteraction]] | None = None,
               electron_electron: list[tuple[int, int, BaseInteraction]] | None = None,
               nuclei_nuclei: list[tuple[int, int, BaseInteraction]] | None = None):
        """
        Update the parameters of spin system. No recomputation of spin vectors does not occur
        :param g_tensors:
        list[BaseInteraction]
            g-tensors corresponding to each electron in `electrons`.
            Each element must be an instance of `BaseInteraction` (e.g., `Interaction`).

        :param electron_nuclei:
        list[tuple[int, int, BaseInteraction]], optional
            Hyperfine interactions between electrons and nuclei.
            Each tuple is of the form (electron_index, nucleus_index, interaction_tensor).
            Default is `None`.

        :param electron_electron:
        list[tuple[int, int, BaseInteraction]], optional
            Dipolar or exchange interactions between pairs of electrons.
            Each tuple is of the form (electron_index, electron_index, interaction_tensor).
            Default is `None`.

        :param nuclei_nuclei:
        list[tuple[int, int, BaseInteraction]], optional
            Dipolar or J-coupling interactions between pairs of nuclei.
            Each tuple is of the form (nucleus_index, nucleus_index, interaction_tensor).
            Default is `None`
        """

        if g_tensors is not None:
            self.g_tensors = nn.ModuleList(g_tensors)

        if electron_nuclei is not None:
            self.electron_nuclei_interactions = nn.ModuleList()
            self.en_indices = []
            self._register_interactions(electron_nuclei, self.electron_nuclei_interactions, self.en_indices)

        if electron_electron is not None:
            self.electron_electron_interactions = nn.ModuleList()
            self.ee_indices = []
            self._register_interactions(electron_electron, self.electron_electron_interactions, self.ee_indices)

        if nuclei_nuclei is not None:
            self.nuclei_nuclei_interactions = nn.ModuleList()
            self.nn_indices = []
            self._register_interactions(nuclei_nuclei, self.nuclei_nuclei_interactions, self.nn_indices)

        self.to(self.device)
        self.to(self.dtype)

    def __repr__(self):
        lines = ["=" * 60]
        lines.append("SPIN SYSTEM SUMMARY")
        lines.append("=" * 60)

        lines.append("\nPARTICLES:")
        lines.append("-" * 20)

        if self.electrons:
            electron_info = []
            for i, electron in enumerate(self.electrons):
                spin_str = f"S={electron.spin} \n"
                g_gactor_str = str(self.g_tensors[i]).replace('\n', '\n      ')
                spin_str += g_gactor_str
                electron_info.append(f"  e{i}: {spin_str}")

            lines.append(f"Electrons ({len(self.electrons)}):")
            lines.extend(electron_info)
        else:
            lines.append("Electrons: None")

        if self.nuclei:
            lines.append(f"\nNuclei ({len(self.nuclei)}):")
            for i, nucleus in enumerate(self.nuclei):
                nucleus_info = f"  n{i}: "
                if hasattr(nucleus, 'isotope'):
                    nucleus_info += f"{nucleus.isotope}, "
                if hasattr(nucleus, 'spin'):
                    nucleus_info += f"I={nucleus.spin}"
                lines.append(nucleus_info)
        else:
            lines.append("\nNuclei: None")

        lines.append(f"\nSYSTEM PROPERTIES:")
        lines.append("-" * 20)
        lines.append(f"Hilbert space dimension: {self.spin_system_dim}")
        lines.append(f"Configuration shape: {tuple(self.config_shape)}")

        total_interactions = (len(self.electron_nuclei) +
                              len(self.electron_electron) +
                              len(self.nuclei_nuclei))

        if total_interactions > 0:
            lines.append(f"\nINTERACTIONS ({total_interactions} total):")
            lines.append("-" * 30)

            # Electron-nuclei interactions
            if self.electron_nuclei:
                lines.append(f"\nElectron-Nucleus ({len(self.electron_nuclei)}):")
                for i, (e_idx, n_idx, interaction) in enumerate(self.electron_nuclei):
                    lines.append(f"  {i + 1}. e{e_idx} ↔ n{n_idx}:")
                    interaction_str = str(interaction).replace('\n', '\n      ')
                    lines.append(f"      {interaction_str}")

            # Electron-electron interactions
            if self.electron_electron:
                lines.append(f"\nElectron-Electron ({len(self.electron_electron)}):")
                for i, (e1_idx, e2_idx, interaction) in enumerate(self.electron_electron):
                    lines.append(f"  {i + 1}. e{e1_idx} ↔ e{e2_idx}:")
                    interaction_str = str(interaction).replace('\n', '\n      ')
                    lines.append(f"      {interaction_str}")

            # Nucleus-nucleus interactions
            if self.nuclei_nuclei:
                lines.append(f"\nNucleus-Nucleus ({len(self.nuclei_nuclei)}):")
                for i, (n1_idx, n2_idx, interaction) in enumerate(self.nuclei_nuclei):
                    lines.append(f"  {i + 1}. n{n1_idx} ↔ n{n2_idx}:")
                    interaction_str = str(interaction).replace('\n', '\n      ')
                    lines.append(f"      {interaction_str}")
        else:
            lines.append(f"\nINTERACTIONS: None")

        lines.append("\n" + "=" * 60)
        return '\n'.join(lines)


class BaseSample(nn.Module):
    def __init__(self, spin_system: SpinSystem,
                 ham_strain: tp.Optional[tp.Union[torch.Tensor, float]] = None,
                 gauss: tp.Optional[tp.Union[torch.Tensor, float]] = None,
                 lorentz: tp.Optional[tp.Union[torch.Tensor, float]] = None,
                 device=torch.device("cpu"), dtype: torch.dtype = torch.float32,
                 *args, **kwargs):
        """
        :param spin_system:
        SpinSystem
            The spin system describing electrons, nuclei, and their interactions.

        :param ham_strain:
        torch.Tensor, float, optional
            Anisotropic line width, due to the unresolved hyperfine interactions.
            The tensor components, provided in one of the following forms:
              - A scalar (for isotropic interaction).
              - A sequence of two values (axial and z components).
              - A sequence of three values (principal components).

        :param gauss:
        torch.Tensor, float, optional
            Gaussian broadening parameter(s). Defines inhomogeneous linewidth
            contributions (e.g., due to static disorder). Default is `None`.

        :param lorentz:
        torch.Tensor, float, optional
            Lorentzian broadening parameter(s). Defines homogeneous linewidth
            contributions (e.g., due to relaxation). Default is `None`

        :param device: device to compute (cpu / gpu)
        :param dtype: dtype
        :param args:
        :param kwargs:
        """
        super().__init__()
        self.base_spin_system = spin_system
        self.modified_spin_system = copy.deepcopy(spin_system)
        self.register_buffer("_ham_strain", self._init_ham_str(ham_strain, device, dtype))

        self.base_ham_strain = copy.deepcopy(self._ham_strain)
        self.register_buffer("gauss", self._init_gauss_lorentz(gauss, device, dtype))
        self.register_buffer("lorentz", self._init_gauss_lorentz(lorentz, device, dtype))
        self.register_buffer("secular_threshold", torch.tensor(1e-9, device=device, dtype=dtype))

        self.to(device)
        self.to(dtype)

    @property
    def device(self):
        return self.base_spin_system.device

    @property
    def complex_dtype(self):
        return self.base_spin_system.complex_dtype

    @property
    def dtype(self):
        return self.base_spin_system.dtype

    @property
    def spin_system_dim(self):
        return self.modified_spin_system.spin_system_dim

    @property
    def config_shape(self):
        return self.modified_spin_system.config_shape

    def _init_gauss_lorentz(
            self, width: tp.Optional[tp.Union[torch.Tensor, float]], device: torch.device, dtype: torch.dtype):
        if width is None:
            if self.base_spin_system.config_shape:
                width = torch.zeros(
                    (*self.base_spin_system.config_shape, ), device=device, dtype=dtype)
            else:
                width = torch.tensor(0.0, device=device, dtype=dtype)
        else:
            width = torch.tensor(width, device=device, dtype=dtype)
            if width.shape != self.base_spin_system.config_shape:
                raise ValueError(f"width batch shape must be equel to base_spin_system config shape")
        return width

    def _init_ham_str(
            self, ham_strain: tp.Optional[tp.Union[torch.Tensor, float]], device: torch.device, dtype: torch.dtype):
        if ham_strain is None:
            ham_strain = torch.zeros(
                (*self.base_spin_system.config_shape, 3), device=device, dtype=dtype
            )
        else:
            ham_strain = init_tensor(ham_strain, device=device, dtype=dtype)
            if ham_strain.shape[:-1] != self.base_spin_system.config_shape:
                raise ValueError(f"ham_strain batch shape must be equel to base_spin_system config shape")
        return ham_strain

    def update(self,
               g_tensors: list[BaseInteraction] = None,
               electron_nuclei: list[tuple[int, int, BaseInteraction]] | None = None,
               electron_electron: list[tuple[int, int, BaseInteraction]] | None = None,
               nuclei_nuclei: list[tuple[int, int, BaseInteraction]] | None = None,
               ham_strain: tp.Optional[torch.Tensor] = None,
               gauss: tp.Union[torch.Tensor, float] = None,
               lorentz: tp.Union[torch.Tensor, float] = None
               ):
        raise NotImplementedError

    def build_electron_electron(self) -> torch.Tensor:
        """Constructs the zero-field Hamiltonian F."""
        F = torch.zeros((*self.config_shape,
                         self.spin_system_dim, self.spin_system_dim),
                        dtype=self.complex_dtype,
                        device=self.device)
        operator_cache = self.modified_spin_system.operator_cache
        for e_idx_1, e_idx_2, interaction in self.modified_spin_system.electron_electron:
            interaction = interaction.tensor.to(self.complex_dtype)
            F += scalar_tensor_multiplication(
                operator_cache[e_idx_1],
                operator_cache[e_idx_2],
                interaction)
        return F

    def build_electron_nuclei(self) -> torch.Tensor:
        """Constructs the zero-field Hamiltonian F."""
        F = torch.zeros((*self.config_shape,
                         self.spin_system_dim, self.spin_system_dim),
                        dtype=self.complex_dtype,
                        device=self.device)
        operator_cache = self.modified_spin_system.operator_cache
        for e_idx, n_idx, interaction in self.modified_spin_system.electron_nuclei:
            interaction = interaction.tensor.to(self.complex_dtype)
            F += scalar_tensor_multiplication(
                operator_cache[e_idx],
                operator_cache[len(self.modified_spin_system.electrons) + n_idx],
                interaction)
        return F

    def build_nuclei_nuclei(self) -> torch.Tensor:
        """Constructs the zero-field Hamiltonian F."""
        F = torch.zeros((*self.config_shape,
                         self.spin_system_dim, self.spin_system_dim),
                        dtype=self.complex_dtype,
                        device=self.device)
        operator_cache = self.modified_spin_system.operator_cache
        for n_idx_1, n_idx_2, interaction in self.modified_spin_system.nuclei_nuclei:
            interaction = interaction.tensor.to(self.complex_dtype)
            F += scalar_tensor_multiplication(
                operator_cache[len(self.modified_spin_system.electrons) + n_idx_1],
                operator_cache[len(self.modified_spin_system.electrons) + n_idx_2],
                interaction)
        return F

    def build_first_order_interactions(self) -> torch.Tensor:
        """Constructs the zero-field Hamiltonian F of the first order operators"""
        return self.build_nuclei_nuclei() + self.build_electron_nuclei() + self.build_electron_electron()

    def build_zero_field_term(self) -> torch.Tensor:
        """Constructs the zero-field Hamiltonian F."""
        return self.build_first_order_interactions()

    def _build_electron_zeeman_terms(self) -> torch.Tensor:
        """Constructs the Zeeman interaction terms Gx, Gy, Gz. for electron spins with give g-tensors"""
        G = torch.zeros((*self.config_shape, 3,
                         self.spin_system_dim, self.spin_system_dim),
                        dtype=self.complex_dtype,
                        device=self.device)
        operator_cache = self.modified_spin_system.operator_cache
        for idx, g_tensor in enumerate(self.modified_spin_system.g_tensors):
            g = g_tensor.tensor.to(self.complex_dtype)
            G += transform_tensor_components(operator_cache[idx], g)
        G *= (constants.BOHR / constants.PLANCK)
        return G

    def _build_nucleus_zeeman_terms(self) -> torch.Tensor:
        """Constructs the Nucleus interaction terms Gx, Gy, Gz. for nucleus spins"""
        G = torch.zeros((*self.config_shape, 3,
                         self.spin_system_dim,
                         self.spin_system_dim),
                        dtype=self.complex_dtype,
                        device=self.device)
        operator_cache = self.modified_spin_system.operator_cache
        for idx, nucleus in enumerate(self.modified_spin_system.nuclei):
            g = nucleus.g_factor
            G += operator_cache[len(self.modified_spin_system.electrons) + idx] * g
        G *= (constants.NUCLEAR_MAGNETRON / constants.PLANCK)
        return G

    def build_zeeman_terms(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Constructs the Zeeman interaction terms Gx, Gy, Gz. for the system"""
        G = self._build_electron_zeeman_terms() + self._build_nucleus_zeeman_terms()
        return G[..., 0, :, :], G[..., 1, :, :], G[..., 2, :, :]

    def calculate_derivative_max(self):
        """
        Calculate the maximum value of the energy derivatives with respect to magnetic field.
        It is assumed that B has direction along z-axis
        :return: the maximum value of the energy derivatives with respect to magnetic field
        """
        electron_contrib = 0
        for idx, electron in enumerate(self.modified_spin_system.electrons):
            electron_contrib += electron.spin * torch.sum(
                self.modified_spin_system.g_tensors[idx].tensor[..., :, 0], dim=-1, keepdim=True).abs()

        nuclei_contrib = 0
        for idx, nucleus in enumerate(self.modified_spin_system.nuclei):
            nuclei_contrib += nucleus.spin * nucleus.g_factor.abs()
        return (electron_contrib * (constants.BOHR / constants.PLANCK) +
            nuclei_contrib * (constants.NUCLEAR_MAGNETRON / constants.PLANCK)).squeeze(dim=-1)

    def get_hamiltonian_terms(self) -> tuple:
        """
        Returns F, Gx, Gy, Gz.
        F is magnetic field free term
        Gx, Gy, Gz are terms multiplied to Bx, By, Bz respectively

        """
        return self.build_zero_field_term(), *self.build_zeeman_terms()

    def get_hamiltonian_terms_secular(self) -> tuple:
        """
        Returns F0, Gx, Gy, Gz.
        F0 is a part of magnetic field free term, which commutes with Gz
        Gx, Gy, Gz are terms multiplied to Bx, By, Bz respectively

        The next approach for computation is used:

        1) The mask M is created
                m = diag(Sz)  # shape (.., n)
                M[i,j] = 1 if m[i] == m[j], else 0
        2) Then F0 = F * M is used
        """
        Gx, Gy, Gz = self.build_zeeman_terms()
        diag = torch.diagonal(Gz, dim1=-2, dim2=-1).real
        mask = abs(diag[..., :, None] - diag[..., None, :]) < self.secular_threshold
        return self.build_zero_field_term() * mask, *(Gx, Gy, Gz)

    def build_field_dep_strain(self):
        """
        Calculate electron Zeeman field dependant strained part
        :return:
        """
        operator_cache = self.modified_spin_system.operator_cache
        for idx, g_tensor in enumerate(self.modified_spin_system.g_tensors):
            g = g_tensor.strained_tensor
            if g is None:
                pass
            else:
                g = g.to(self.complex_dtype)
                yield (
                    g_tensor.strain_correlation.to(self.complex_dtype),
                    operator_cache[idx],
                    g[..., :, 2, :] * constants.BOHR / constants.PLANCK)

    def build_zero_field_strain(self) -> torch.Tensor:
        """Constructs the zero-field strained part."""
        yield from self.build_electron_nuclei_strain()
        yield from self.build_electron_electron_strain()

    def build_electron_nuclei_strain(self) -> torch.Tensor:
        """Constructs the nuclei strained part."""
        operator_cache = self.modified_spin_system.operator_cache
        for e_idx, n_idx, electron_nuclei_interaction in self.modified_spin_system.electron_nuclei:
            electron_nuclei = electron_nuclei_interaction.strained_tensor
            if electron_nuclei is None:
                pass
            else:
                electron_nuclei = electron_nuclei.to(self.complex_dtype)
                yield (
                    electron_nuclei_interaction.strain_correlation.to(self.complex_dtype),
                    operator_cache[e_idx],
                    operator_cache[len(self.modified_spin_system.electrons) + n_idx],
                    electron_nuclei)

    def build_electron_electron_strain(self) -> torch.Tensor:
        """Constructs the electron-electron strained part."""
        operator_cache = self.modified_spin_system.operator_cache
        for e_idx_1, e_idx_2, electron_electron_interaction in self.modified_spin_system.electron_electron:
            electron_electron = electron_electron_interaction.strained_tensor
            if electron_electron is None:
                pass
            else:
                electron_electron = electron_electron.to(self.complex_dtype)
                yield (
                    electron_electron_interaction.strain_correlation.to(self.complex_dtype),
                    operator_cache[e_idx_1],
                    operator_cache[e_idx_2],
                    electron_electron)

    def __repr__(self):
        spin_system_summary = str(self.base_spin_system)

        lines = []
        lines.append(spin_system_summary)
        lines.append("\n" + "=" * 60)
        lines.append("GENERAL INFO: ")
        lines.append("=" * 60)

        is_batched = self.base_spin_system.config_shape

        if is_batched:
            lorentz = self.lorentz.flatten(0, -1)
            gauss = self.gauss.flatten(0, -1)
            batch_size = lorentz.shape[0] if hasattr(self.lorentz, 'shape') else len(self.lorentz)
            lines.append(f"BATCHED (batch_size={batch_size}) - showing first instance:")

            lines.append(f"lorentz: {lorentz[0].item():.5f} T")
            lines.append(f"gauss: {gauss[0].item():.5f} T")

            ham_str = self.base_ham_strain.flatten(0, -2)[0]
            ham_components = [f"{val:.4e}" if abs(val) >= 1e4 else f"{val:.4f}"
                              for val in ham_str.tolist()]
            ham_dim = self.base_ham_strain.shape[1:]
            lines.append(f"ham_str (dim={ham_dim}): {ham_components} Hz")
        else:
            lines.append(f"lorentz: {self.lorentz.item():.5f} T")
            lines.append(f"gauss: {self.gauss.item():.5f} T")

            ham_str = self.base_ham_strain
            ham_components = [f"{val:.4e}" if abs(val) >= 1e4 else f"{val:.4f}"
                              for val in ham_str.tolist()]
            lines.append(f"ham_str: {ham_components} Hz")
        return '\n'.join(lines)


class SpinSystemOrientator:
    """
    The helper class that allow toa
    transform spin system to spin system at different rotation angles.
    Effectively rotate all Hamiltonian parts
    """
    def __call__(self, spin_system: SpinSystem, rotation_matrices: torch.Tensor) -> SpinSystem:
        """
        :param spin_system: spin_system with interactions
        :param rotation_matrices: rotation_matrices that rotate spin system
        :return: modified spin system with all rotated interactions
        """
        spin_system = self.transform_spin_system_to_oriented(copy.deepcopy(spin_system), rotation_matrices)
        return spin_system

    def interactions_to_multioriented(self, interactions: list[nn.Module], rotation_matrices: torch.Tensor):
        interactions_tensors = torch.stack([interaction.tensor for interaction in interactions], dim=0)
        interactions_tensors = utils.apply_expanded_rotations(rotation_matrices, interactions_tensors)
        not_none_strained = [
            interaction.strained_tensor for interaction in interactions if interaction.strained_tensor is not None
        ]
        none_strained_flag = [
            True if interaction.strained_tensor is None else False for interaction in interactions
        ]
        if not_none_strained:
            strained_tensors = torch.stack(not_none_strained, dim=0)
            strained_tensors = utils.apply_expanded_rotations(rotation_matrices, strained_tensors)
            strained_tensors = strained_tensors.transpose(-3, -4)
            strained_iterator = iter(strained_tensors)
            strined_res = [None if x else next(strained_iterator) for x in none_strained_flag]
        else:
            strined_res = [None] * len(interactions)
        return interactions_tensors, strined_res

    def _apply_reverse_transform(self, spin_system: SpinSystem, new_interactions: list[MultiOrientedInteraction]):
        # Determine how many interactions belong to each original group.
        num_g = len(spin_system.g_tensors)
        num_nuc = len(spin_system.electron_nuclei)

        # Split the new interactions list into the three groups.
        new_g_tensors = new_interactions[:num_g]
        new_electron_nuclei = new_interactions[num_g:num_g + num_nuc]
        new_electron_electron = new_interactions[num_g + num_nuc:]

        spin_system.g_tensors = nn.ModuleList(new_g_tensors)
        spin_system.electron_nuclei_interactions = nn.ModuleList(new_electron_nuclei)
        spin_system.electron_electron_interactions = nn.ModuleList(new_electron_electron)
        return spin_system

    def transform_spin_system_to_oriented(self, spin_system: SpinSystem, rotation_matrices: torch.Tensor):
        #rotation_matrices = rotation_matrices.to(torch.complex64)
        config_shape = torch.Size([*spin_system.config_shape, rotation_matrices.shape[0]])
        interactions = [g_tensor for g_tensor in spin_system.g_tensors] + \
                       [el_nuc[-1] for el_nuc in spin_system.electron_nuclei] +\
                       [el_el[-1] for el_el in spin_system.electron_electron]
        strain_correlations = [g_tensor.strain_correlation for g_tensor in spin_system.g_tensors] + \
                       [el_nuc[-1].strain_correlation for el_nuc in spin_system.electron_nuclei] +\
                       [el_el[-1].strain_correlation for el_el in spin_system.electron_electron]

        interactions_tensors, strained_tensors = self.interactions_to_multioriented(interactions, rotation_matrices)
        interactions =\
            [
                MultiOrientedInteraction(interactions_tensor, strained_tensor, config_shape, strain_correlation) for
                interactions_tensor, strained_tensor, strain_correlation in
                zip(interactions_tensors, strained_tensors, strain_correlations)
            ]

        spin_system = self._apply_reverse_transform(spin_system, interactions)
        return spin_system


class MultiOrientedSample(BaseSample):
    def __init__(self, spin_system: SpinSystem,
                 ham_strain: tp.Optional[torch.Tensor] = None,
                 gauss: torch.Tensor = None,
                 lorentz: torch.Tensor = None,
                 mesh: tp.Optional[tp.Union[BaseMesh, tuple[int, int]]] = None,
                 spin_system_frame: tp.Optional[torch.Tensor] = None,
                 device: torch.device = torch.device("cpu"),
                 dtype: torch.dtype = torch.float32,
                 ):
        """
        :param spin_system:
        SpinSystem
            The spin system describing electrons, nuclei, and their interactions.

        :param spin_system_frame:
        torch.Tensor | Sequence[float] optional
            Orientation of the spin system. Can be provided as:
              - A 1D tensor of shape (3,) representing Euler angles in ZYZ' convention.
              - A 2D tensor of shape (3, 3) representing a rotation matrix.
            Default is `None`, meaning lab frame.

        This parameter has meaning for tasks when there are other samples in the system and you try to create some
        complex system. For example, you have two triplets connected by dipole-dipole interaction
        but initial populations of the triplet states you set at the basises of individual triplets.

        :param ham_strain:
        torch.Tensor, optional
            Anisotropic line width, due to the unresolved hyperfine interactions.
            The tensor components, provided in one of the following forms:
              - A scalar (for isotropic interaction).
              - A sequence of two values (axial and z components).
              - A sequence of three values (principal components).

        :param gauss:
        torch.Tensor, optional
            Gaussian broadening parameter(s). Defines inhomogeneous linewidth
            contributions (e.g., due to static disorder). Default is `None`.

        :param lorentz:
        torch.Tensor, optional
            Lorentzian broadening parameter(s). Defines homogeneous linewidth
            contributions (e.g., due to relaxation). Default is `None`

        :param mesh: The mesh to perform rotations for powder samples. It can be:
           -tuple[initial_grid_frequency, interpolation_grid_frequency],
           where initial_grid_frequency is the size of the initial mesh,
           interpolation_grid_frequency is the size of the interpolation mesh
           For this case mesh will be initialize as DelaunayMeshNeighbour with given sizes

           -Inheritor of Base Mesh

        If it is None it will be initialize as DelaunayMeshNeighbour with initial_grid_frequency = 20

        :param device: device to compute (cpu / gpu)
        """
        super().__init__(spin_system, ham_strain, gauss, lorentz, device=device, dtype=dtype)
        self.mesh = self._init_mesh(mesh, device=device, dtype=dtype)
        self._construct_spin_system_rot_matrix(frame=spin_system_frame, dtype=dtype, device=device)
        rotation_matrices = self.mesh.rotation_matrices

        self._ham_strain = self._expand_hamiltonian_strain(
            self.base_ham_strain,
            self.orientation_vector(rotation_matrices)
        )
        if spin_system_frame is None:
            self.modified_spin_system = SpinSystemOrientator()(spin_system, rotation_matrices)
        else:
            self.modified_spin_system = SpinSystemOrientator()(
                spin_system, torch.matmul(rotation_matrices, self._spin_system_rot_matrix)
            )

    def _construct_spin_system_rot_matrix(
            self, frame: tp.Optional[torch.tensor],
            device: torch.device,
            dtype: torch.dtype
    ):
        if frame is None:
            _frame = None
            _rot_matrix = None

        else:
            if not isinstance(frame, torch.Tensor):
                raise TypeError("frame must be a torch.Tensor or None.")
            if frame.shape[-2:] == (3, 3):
                _frame = utils.rotation_matrix_to_euler_angles(frame)
                _rot_matrix = frame.to(dtype).to(device)

            elif frame.shape == (3, ):
                _frame = frame.to(dtype)
                _rot_matrix = self.euler_to_rotmat(_frame).to(dtype).to(device)

            else:
                raise ValueError(
                    "frame must be either:\n"
                    "  • None (→ identity rotation),\n"
                    "  • a tensor of Euler angles with shape batch×3,\n"
                    "  • or a tensor of rotation matrices with shape batch×3×3."
                )
        self.register_buffer("_spin_system_frame", _frame)
        self.register_buffer("_spin_system_rot_matrix", _rot_matrix)

    def _init_mesh(
            self, mesh: tp.Optional[tp.Union[BaseMesh, tuple[int, int]]],
            device: torch.device, dtype: torch.dtype
    ):
        if mesh is None:
            mesh = mesher.DelaunayMeshNeighbour(interpolate=False,
                                                initial_grid_frequency=20,
                                                interpolation_grid_frequency=40, device=device, dtype=dtype)
        elif isinstance(mesh, tuple):
            initial_grid_frequency = mesh[0]
            interpolation_grid_frequency = mesh[1]
            if initial_grid_frequency >= interpolation_grid_frequency:
                interpolate = False
            else:
                interpolate = True
            mesh = mesher.DelaunayMeshNeighbour(interpolate=interpolate,
                                                initial_grid_frequency=initial_grid_frequency,
                                                interpolation_grid_frequency=interpolation_grid_frequency,
                                                device=device, dtype=dtype)
        return mesh

    def _expand_hamiltonian_strain(self, ham_strain: torch.Tensor, orientation_vector: torch.Tensor):
        ham_shape = ham_strain.shape[:-1]
        orient_shape = orientation_vector.shape[:-1]
        ham_expanded = ham_strain.view(*ham_shape, *([1] * len(orient_shape)), ham_strain.shape[-1])
        orient_expanded = orientation_vector.view(*([1] * len(ham_shape)), *orient_shape, orientation_vector.shape[-1])
        result = ((ham_expanded ** 2) * (orient_expanded ** 2)).sum(dim=-1).sqrt()
        return result

    def orientation_vector(self, rotation_matrices: torch.Tensor):
        return rotation_matrices[..., -1, :]

    def build_ham_strain(self) -> torch.Tensor:
        """Constructs the zero-field strained part of Hamiltonian"""
        return self._ham_strain

    def update(self,
               g_tensors: list[BaseInteraction] = None,
               electron_nuclei: list[tuple[int, int, BaseInteraction]] | None = None,
               electron_electron: list[tuple[int, int, BaseInteraction]] | None = None,
               nuclei_nuclei: list[tuple[int, int, BaseInteraction]] | None = None,
               ham_strain: tp.Optional[torch.Tensor] = None,
               gauss: tp.Union[torch.Tensor, float] = None,
               lorentz: tp.Union[torch.Tensor, float] = None
               ):
        """
        Update the parameters of a sample. No recomputation of spin vectors does not occur
        :param g_tensors:
        list[BaseInteraction]
            g-tensors corresponding to each electron in `electrons`.
            Each element must be an instance of `BaseInteraction` (e.g., `Interaction`).

        :param electron_nuclei:
        list[tuple[int, int, BaseInteraction]], optional
            Hyperfine interactions between electrons and nuclei.
            Each tuple is of the form (electron_index, nucleus_index, interaction_tensor).
            Default is `None`.

        :param electron_electron:
        list[tuple[int, int, BaseInteraction]], optional
            Dipolar or exchange interactions between pairs of electrons.
            Each tuple is of the form (electron_index, electron_index, interaction_tensor).
            Default is `None`.

        :param nuclei_nuclei:
        list[tuple[int, int, BaseInteraction]], optional
            Dipolar or J-coupling interactions between pairs of nuclei.
            Each tuple is of the form (nucleus_index, nucleus_index, interaction_tensor).
            Default is `None`

        :param ham_strain:
        torch.Tensor, optional
            Anisotropic line width, due to the unresolved hyperfine interactions.
            The tensor components, provided in one of the following forms:
              - A scalar (for isotropic interaction).
              - A sequence of two values (axial and z components).
              - A sequence of three values (principal components).

        :param gauss:
        torch.Tensor, optional
            Gaussian broadening parameter(s). Defines inhomogeneous linewidth
            contributions (e.g., due to static disorder). Default is `None`.

        :param lorentz:
        torch.Tensor, optional
            Lorentzian broadening parameter(s). Defines homogeneous linewidth
            contributions (e.g., due to relaxation). Default is `None`
        """

        rotation_matrices = self.mesh.rotation_matrices
        self.base_spin_system.update(g_tensors, electron_nuclei, electron_electron, nuclei_nuclei)

        if ham_strain is not None:
            self.base_ham_strain = self._init_ham_str(ham_strain, self.device, self.dtype)
            self._ham_strain = self._expand_hamiltonian_strain(
                self.base_ham_strain,
                self.orientation_vector(rotation_matrices)
            )
        self.gauss = self._init_gauss_lorentz(gauss, self.device, self.dtype)
        self.lorentz = self._init_gauss_lorentz(lorentz, self.device, self.dtype)
        if self._spin_system_frame is None:
            self.modified_spin_system = SpinSystemOrientator()(self.base_spin_system, rotation_matrices)
        else:
            self.modified_spin_system = SpinSystemOrientator()(
                self.base_spin_system, torch.matmul(rotation_matrices, self._spin_system_rot_matrix)
            )