"""Module for evaluating Z2 topological invariant."""

from typing import Tuple, Union
from math import pi

import numpy as np

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


__all__ = ["Z2"]


class Z2(DiagSolver):
    """
    Class for evaluating the Z2 topological invariant.

    Reference: https://journals.aps.org/prb/abstract/10.1103/PhysRevB.84.075119
    """
    @staticmethod
    def _eval_phase(z: np.ndarray) -> np.ndarray:
        """
        Evaluate the phases of all the elements in an array.
        :param z: complex128 array for which the phases will be evaluated
        :return: phases of all the elements in the array
        """
        phase = np.angle(z)
        norm = np.absolute(z)
        for i in range(phase.shape[0]):
            if norm.item(i) < 1e-14:
                phase[i] = 0.0
        return phase

    @staticmethod
    def _get_min_key(data_dict: dict) -> Union[int, tuple, None]:
        """
        Return the key with the minimum value.
        :param data_dict: dictionary to search
            with keys being integers or tuples and values being floats
        :return: key with the minimum value
        """
        min_value = min(data_dict.values())
        result = None
        for key, value in data_dict.items():
            if abs(value - min_value) < 1.0e-15:
                result = key
                break
        return result

    def calc_phases(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate the phases of WF centers as function of kb (eqn. 12 of ref).
        :return: (kb_array, phases)
            kb_array: (num_kb,) float64 array
            FRACTIONAL coordinates of the loop along b-axis
            phases: (num_kpt, num_occ) float64 array
            phases of WF centers
        """
        # Check and set num_occ
        if __debug__:
            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 len(self.config.ka_array.shape) == 1
            assert len(self.config.kb_array.shape) == 1
            assert self.config.kb_array.shape[0] >= 1
            assert self.config.kb_array.shape[0] >= 1
            assert num_occ in range(1, num_max+1), (f"num_occ {num_occ} should"
                                                    f" be in [1, {num_max}]")
        self._pre_run("z2")
        self._cpp_solver.calc_phases(
            self.config.ka_array,
            self.config.kb_array,
            self.config.kc)
        return load_z2(self.config.prefix)

    def reorder_phases(self,
                       phases: np.ndarray,
                       threshold: float = 0.1,
                       smooth: bool = True) -> np.ndarray:
        """
        Reorder the phases to improve continuity and smoothness.

        NOTE: sometimes this method may not work properly. In that case,
        try to increase threshold and the density of kb_array.

        :param phases: (num_kpt, num_occ) float64 array
            phases of WF centers
        :param threshold: threshold for detecting discontinuity
        :param smooth: whether to smooth the phases by ensuring the continuity
            of 1st order derivative with respect to kb
        :return: (num_kpt, num_occ) float64 array
            reordered phases
        """
        phases = phases.copy()

        # Fix phase discontinuity
        for ib in range(phases.shape[0]-1):
            phase_i0 = phases[ib]
            phase_i1 = phases[ib + 1]
            phase_i1_cp = phase_i1.copy()
            for j in range(self.config.num_occ):
                phase_j0 = phase_i0.item(j)
                phase_j1 = phase_i1.item(j)

                # Correction should be done for data points will large
                # difference with respect to previous point
                if abs(phase_j0 - phase_j1) > threshold:
                    # Find candidates at i+1 that conserve phase continuity
                    # considering periodicity
                    # Key: (m,n) of |phase_m1 - phase_j1 - n * pi| < threshold
                    # value: |phase_m1 - phase_j1 - n * pi|
                    diff_dict = dict()
                    for m, phase_m1 in enumerate(phase_i1_cp):
                        for n in (-2, 0, 2):
                            phase_diff = abs(phase_j0 - (phase_m1 + n * pi))
                            if phase_diff < threshold:
                                diff_dict[(m, n)] = phase_diff

                    if len(diff_dict) > 0:
                        min_key = self._get_min_key(diff_dict)
                        min_pos = min_key[0]
                        min_shift = min_key[1] * pi
                        phases[ib + 1, j] = phase_i1_cp[min_pos] + min_shift
                        phase_i1_cp = np.delete(phase_i1_cp, min_pos)

        # Fix derivative discontinuity (smoothness)
        if smooth:
            for ib in range(phases.shape[0]-2):
                phase_i0 = phases[ib]
                phase_i1 = phases[ib + 1]
                phase_i2_cp = phases[ib + 2].copy()
                for j in range(self.config.num_occ):
                    phase_j0 = phase_i0.item(j)
                    phase_j1 = phase_i1.item(j)

                    # Find candidates at i+2 that conserve phase continuity
                    # phase_m2 := phase_i2[m]
                    # Key: index(m) of |phase_m2 - phase_j1| < threshold
                    # Value: phase_m2
                    diff_dict = dict()
                    for m, phase_m2 in enumerate(phase_i2_cp):
                        if abs(phase_j1 - phase_m2) < threshold:
                            diff_dict[m] = phase_m2

                    # Find the point that ensures the best smoothness
                    # Key: index(m) of |phase_m2 - phase_j1| < threshold
                    # Value: |(phase_m2 - phase_j1) - (phase_j1 - phase_j0)|
                    if len(diff_dict) > 0:
                        for m, phase_m2 in diff_dict.items():
                            phase_diff = abs(phase_m2 - 2 * phase_j1 + phase_j0)
                            diff_dict[m] = phase_diff
                        min_pos = self._get_min_key(diff_dict)

                        # Update phases at ib+2
                        phases[ib + 2, j] = phase_i2_cp[min_pos]
                        phase_i2_cp = np.delete(phase_i2_cp, min_pos)
        return phases

    @staticmethod
    def count_crossing(phases: np.ndarray, phase_ref: float = 0.5) -> int:
        """
        Count the number that the phases go across a reference value.

        The Z2 topological invariant is determined as num_crossing % 2.

        NOTE: the phases must be CORRECTLY reordered before passing to this
        function. Otherwise, the crossing number will be wrong!

        If the reordering algorithm fails to work by all means, then you have to
        plot the phases using scatter plot and count the crossing number manually.

        :param phases: (num_kpt, num_occ) float64 array
            phases of WF centers
        :param phase_ref: reference value
        :return: number of crossing
        """
        num_kpt = phases.shape[0]
        num_occ = phases.shape[1]
        num_crossing = 0
        for ik in range(0, num_kpt - 1):
            for ib in range(num_occ):
                d1 = phases.item(ik, ib) - phase_ref
                d2 = phases.item(ik+1, ib) - phase_ref
                if d1 * d2 < 0.0:
                    num_crossing += 1
        return num_crossing
