"""Module for evaluating Berry curvature, Chern number, etc."""

import numpy as np

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


__all__ = ["BerryData", "Berry"]


class BerryData:
    """
    Data class of Berry connection and curvature.

    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
    omega_xy: (num_kpt, num_states) float64 array
        xy components of Berry curvature in NM^2
    """
    def __init__(self, prefix: str) -> None:
        """
        :param prefix: prefix of data files
        :return: None
        """
        self.kpt_cart = None
        self.bands = None
        self.omega_xy = 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.omega_xy = load_berry(prefix)


class Berry(DiagSolver):
    """
    Berry curvature and Chern number calculator.
    """
    def _check_k_grid_dim(self) -> None:
        """Check k_grid_size and dimension."""
        k_grid_size = self.config.k_grid_size
        assert len(k_grid_size) == 3, "Length of k_grid_size should be 3"
        for i in k_grid_size:
            assert i >= 1, f"Dimension of k_grid_size should >= 1"
        assert k_grid_size[2] == 1, "kmesh_size[2] should == 1"

    def calc_berry_curvature_kubo(self) -> BerryData:
        """
        Calculate Berry curvature and Chern number for each band in FBZ.
        :return: Berry curvature in FBZ
        """
        if __debug__:
            self._check_k_grid_dim()
        self._pre_run("berry")
        self._cpp_solver.calc_berry_curvature_kubo()
        berry_data = BerryData(self.config.prefix)
        return berry_data

    def calc_berry_curvature_wilson(self) -> BerryData:
        """
        Calculate total Berry curvature and Chern number in FBZ.
        :return: total Berry curvature in FBZ
        """
        if __debug__:
            self._check_k_grid_dim()
            if self.config.diag_algo.find("feast") != -1:
                num_max = self.config.feast_num_states
            else:
                num_max = self._model.num_orb
            num_occ = self.config.num_occ
            assert num_occ in range(1, num_max+1), (f"num_occ {num_occ} should"
                                                    f" be in [1, {num_max}]")
        self._pre_run("berry")
        self._cpp_solver.calc_berry_curvature_wilson()
        berry_data = BerryData(self.config.prefix)
        return berry_data
