"""Module for calculating magnetic properties."""

import numpy as np

from tbplas.diag.base import DiagSolver
from tbplas.diag.io import load_spin

__all__ = ["SpinTexture", "SpinData"]


class SpinData:
    """
    Data class of spin texture.

    NOTE: DO NOT use dataclass.

    Attributes
    ----------
    kpt_cart: (num_kpt, 3) float64 array
        Cartesian coordinates of k-points in 1/NM
    bands: (num_kpt, num_states) float64 array
        Energies on k-points in eV
    sigma_x, sigma_y, sigma_z: (num_kpt, num_states) float64 array
        x, y, z components of expectation value of Pauli operators
    """
    def __init__(self, prefix: str) -> None:
        """
        :param prefix: prefix of data files
        :return: None
        """
        self.kpt_cart = None
        self.bands = None
        self.sigma_x = None
        self.sigma_y = None
        self.sigma_z = None
        self.load(prefix)

    def load(self, prefix: str) -> None:
        """
        Load spin data from file.
        :param prefix: prefix of data files
        :return: None
        """
        self.kpt_cart, self.bands, self.sigma_x, self.sigma_y, self.sigma_z = \
            load_spin(prefix)


class SpinTexture(DiagSolver):
    """
    Spin texture calculator.
    """
    def calc_spin_texture(self) -> SpinData:
        """
        Calculate spin texture.
        :return: spin texture on given k-points.
        """
        if __debug__:
            self._check_k_points()
        self._pre_run("spin")
        self._cpp_solver.calc_spin_texture(self.config.k_points)
        spin_data = SpinData(self.config.prefix)
        return spin_data
