"""A class for interacting with KPath classes to
generate high-symmetry k-paths using different conventions.
"""

from __future__ import annotations

import itertools
import warnings
from typing import Any

import networkx as nx
import numpy as np

from pymatgen.electronic_structure.bandstructure import BandStructureSymmLine
from pymatgen.electronic_structure.core import Spin
from pymatgen.symmetry.analyzer import cite_conventional_cell_algo
from pymatgen.symmetry.kpath import KPathBase, KPathLatimerMunro, KPathSeek, KPathSetyawanCurtarolo

__author__ = "Jason Munro"
__copyright__ = "Copyright 2020, The Materials Project"
__version__ = "0.1"
__maintainer__ = "Jason Munro"
__email__ = "jmunro@lbl.gov"
__status__ = "Development"
__date__ = "March 2020"


@cite_conventional_cell_algo
class HighSymmKpath(KPathBase):
    """This class generates path along high symmetry lines in the
    Brillouin zone according to different conventions.
    The class is designed to be used with a specific primitive
    cell setting. The definitions for the primitive cell
    used can be found in: Computational Materials Science,
    49(2), 299-312. doi:10.1016/j.commatsci.2010.05.010.
    The space group analyzer can be used to produce the correct
    primitive structure
    (method get_primitive_standard_structure(international_monoclinic=False)).
    Ensure input structure is correct before 'get_kpoints()' method is used.
    See individual KPath classes for details on specific conventions.
    """

    def __init__(
        self,
        structure,
        has_magmoms=False,
        magmom_axis=None,
        path_type="setyawan_curtarolo",
        symprec=0.01,
        angle_tolerance=5,
        atol=1e-5,
    ):
        """Initialize a HighSymmKpath.

        Args:
            structure (Structure): Structure object
            has_magmoms (bool): Whether the input structure contains
                magnetic moments as site properties with the key 'magmom.'
                Values may be in the form of 3-component vectors given in
                the basis of the input lattice vectors, in
                which case the spin axis will default to a_3, the third
                real-space lattice vector (this triggers a warning).
            magmom_axis (list or numpy array): 3-component vector specifying
                direction along which magnetic moments given as scalars
                should point. If all magnetic moments are provided as
                vectors then this argument is not used.
            path_type (str): Chooses which convention to use to generate
                the high symmetry path. Options are: 'setyawan_curtarolo', 'hinuma',
                'latimer_munro' for the Setyawan & Curtarolo, Hinuma et al., and
                Latimer & Munro conventions. Choosing 'all' will generate one path
                with points from all three conventions. Equivalent labels between
                each will also be generated. Order will always be Latimer & Munro,
                Setyawan & Curtarolo, and Hinuma et al. Lengths for each of the paths
                will also be generated and output as a list. Note for 'all' the user
                will have to alter the labels on their own for plotting.
            symprec (float): Tolerance for symmetry finding
            angle_tolerance (float): Angle tolerance for symmetry finding.
            atol (float): Absolute tolerance used to determine symmetric
                equivalence of points and lines on the BZ.
        """
        super().__init__(structure, symprec=symprec, angle_tolerance=angle_tolerance, atol=atol)

        self._path_type = path_type

        self._equiv_labels = self._path_lengths = self._label_index = None

        if path_type != "all":
            if path_type == "latimer_munro":
                self._kpath = self._get_lm_kpath(has_magmoms, magmom_axis, symprec, angle_tolerance, atol).kpath
            elif path_type == "setyawan_curtarolo":
                self._kpath = self._get_sc_kpath(symprec, angle_tolerance, atol).kpath
            elif path_type == "hinuma":
                hin_dat = self._get_hin_kpath(symprec, angle_tolerance, atol, not has_magmoms)
                self._kpath = hin_dat.kpath
                self._hin_tmat = hin_dat._tmat

        else:
            if has_magmoms:
                raise ValueError("Cannot select 'all' with non-zero magmoms.")

            lm_bs = self._get_lm_kpath(has_magmoms, magmom_axis, symprec, angle_tolerance, atol)
            rpg = lm_bs._rpg

            sc_bs = self._get_sc_kpath(symprec, angle_tolerance, atol)
            hin_bs = self._get_hin_kpath(symprec, angle_tolerance, atol, not has_magmoms)

            index = 0
            cat_points = {}
            label_index = {}
            num_path = []
            self._path_lengths = []

            for bs in [lm_bs, sc_bs, hin_bs]:
                for value in bs.kpath["kpoints"]:
                    cat_points[index] = bs.kpath["kpoints"][value]
                    label_index[index] = value
                    index += 1

                total_points_path = 0
                for seg in bs.kpath["path"]:
                    total_points_path += len(seg)

                for block in bs.kpath["path"]:
                    new_block = [
                        ind
                        for label in block
                        for ind in range(len(label_index) - len(bs.kpath["kpoints"]), len(label_index))
                        if label_index[ind] == label
                    ]

                    num_path.append(new_block)

                self._path_lengths.append(total_points_path)

            self._label_index = label_index

            self._kpath = {"kpoints": cat_points, "path": num_path}

            self._equiv_labels = self._get_klabels(lm_bs, sc_bs, hin_bs, rpg)

    @property
    def path_type(self):
        """The type of kpath chosen."""
        return self._path_type

    @property
    def label_index(self):
        """The correspondence between numbers and kpoint symbols for the
        combined kpath generated when path_type = 'all'. None otherwise.
        """
        return self._label_index

    @property
    def equiv_labels(self):
        """The correspondence between the kpoint symbols in the Latimer and
        Munro convention, Setyawan and Curtarolo, and Hinuma
        conventions respectively. Only generated when path_type = 'all'.
        """
        return self._equiv_labels

    @property
    def path_lengths(self):
        """List of lengths of the Latimer and Munro, Setyawan and Curtarolo, and Hinuma
        conventions in the combined HighSymmKpath object when path_type = 'all' respectively.
        None otherwise.
        """
        return self._path_lengths

    def _get_lm_kpath(self, has_magmoms, magmom_axis, symprec, angle_tolerance, atol):
        """Get the Latimer and Munro k-path with labels."""
        return KPathLatimerMunro(self._structure, has_magmoms, magmom_axis, symprec, angle_tolerance, atol)

    def _get_sc_kpath(self, symprec, angle_tolerance, atol):
        """Get the Setyawan and Curtarolo k-path with labels."""
        kpath = KPathSetyawanCurtarolo(self._structure, symprec, angle_tolerance, atol)

        self.prim = kpath.prim
        self.conventional = kpath.conventional
        self.prim_rec = kpath.prim_rec
        self._rec_lattice = self.prim_rec

        return kpath

    def _get_hin_kpath(self, symprec, angle_tolerance, atol, tri):
        """Get the Hinuma et al. k-path with labels."""
        bs = KPathSeek(self._structure, symprec, angle_tolerance, atol, tri)

        kpoints = bs.kpath["kpoints"]
        tmat = bs._tmat
        for key in kpoints:
            kpoints[key] = np.dot(np.transpose(np.linalg.inv(tmat)), kpoints[key])

        bs.kpath["kpoints"] = kpoints
        self._rec_lattice = self._structure.lattice.reciprocal_lattice

        warnings.warn(
            "K-path from the Hinuma et al. convention has been transformed to the basis of the reciprocal lattice"
            "of the input structure. Use `KPathSeek` for the path in the original author-intended basis.",
            stacklevel=2,
        )

        return bs

    def _get_klabels(self, lm_bs, sc_bs, hin_bs, rpg):
        """Get equivalent k-point labels across the three k-path conventions.

        Returns:
            dict[str, dict[str, dict[str, str]]]: equivalent labels for paths if 'all' is chosen.
            If an exact kpoint match cannot be found, symmetric equivalency will be
            searched for and indicated with an asterisk in the equivalent label.
            If an equivalent label can still not be found, or the point is not in
            the explicit kpath, its equivalent label will be set to itself in the output.
        """
        lm_path = lm_bs.kpath
        sc_path = sc_bs.kpath
        hin_path = hin_bs.kpath

        n_op = len(rpg)

        pairs = itertools.permutations(
            [
                {"setyawan_curtarolo": sc_path},
                {"latimer_munro": lm_path},
                {"hinuma": hin_path},
            ],
            r=2,
        )
        labels = {"setyawan_curtarolo": {}, "latimer_munro": {}, "hinuma": {}}

        for a, b in pairs:
            [(a_type, a_path)] = list(a.items())
            [(b_type, b_path)] = list(b.items())

            sc_count = np.zeros(n_op)

            for o_num in range(n_op):
                a_tr_coord = [np.dot(rpg[o_num], coord_a) for coord_a in a_path["kpoints"].values()]

                for coord_a in a_tr_coord:
                    for value in b_path["kpoints"].values():
                        if np.allclose(value, coord_a, atol=self._atol):
                            sc_count[o_num] += 1
                            break

            a_to_b_labels = {}
            unlabeled = {}

            for label_a, coord_a in a_path["kpoints"].items():
                coord_a_t = np.dot(rpg[np.argmax(sc_count)], coord_a)
                assigned = False

                for label_b, coord_b in b_path["kpoints"].items():
                    if np.allclose(coord_b, coord_a_t, atol=self._atol):
                        a_to_b_labels[label_a] = label_b
                        assigned = True
                        break

                if not assigned:
                    unlabeled[label_a] = coord_a

            for label_a, coord_a in unlabeled.items():
                key = None

                for op in rpg:
                    coord_a_t = np.dot(op, coord_a)
                    key = [
                        key
                        for key, value in b_path["kpoints"].items()
                        if np.allclose(value, coord_a_t, atol=self._atol)
                    ]

                    if key != []:
                        a_to_b_labels[label_a] = key[0][0] + "^{*}"
                        break

                if key == []:
                    a_to_b_labels[label_a] = label_a

            labels[a_type][b_type] = a_to_b_labels

        return labels

    @staticmethod
    def get_continuous_path(bandstructure: BandStructureSymmLine) -> BandStructureSymmLine:
        """Obtain a continuous version of an inputted path using graph theory.
        This routine will attempt to add connections between nodes of
        odd-degree to ensure a Eulerian path can be formed. Initial
        k-path must be able to be converted to a connected graph. See
        npj Comput Mater 6, 112 (2020). 10.1038/s41524-020-00383-7
        for more details.

        Args:
            bandstructure (BandStructureSymmLine): BandStructureSymmLine object.

        Returns:
            bandstructure (BandStructureSymmLine): New BandStructureSymmLine object with continuous path.
        """
        G = nx.Graph()
        labels = [point.label for point in bandstructure.kpoints if point.label is not None]

        plot_axis = []
        for i in range(len(labels) // 2):
            G.add_edges_from([(labels[2 * i], labels[(2 * i) + 1])])
            plot_axis.append((labels[2 * i], labels[(2 * i) + 1]))

        G_euler = nx.algorithms.euler.eulerize(G)

        G_euler_circuit = nx.algorithms.euler.eulerian_circuit(G_euler)

        distances_map = []
        kpath_euler = []

        for edge_euler in G_euler_circuit:
            kpath_euler.append(edge_euler)
            for edge_reg in plot_axis:
                if edge_euler == edge_reg:
                    distances_map.append((plot_axis.index(edge_reg), False))
                elif edge_euler[::-1] == edge_reg:
                    distances_map.append((plot_axis.index(edge_reg), True))

        spins = [Spin.up, Spin.down] if bandstructure.is_spin_polarized else [Spin.up]

        new_kpoints = []
        new_bands = {spin: [np.array([]) for _ in range(bandstructure.nb_bands)] for spin in spins}

        new_projections: dict[Spin, Any] = {}
        if has_projections := isinstance(bandstructure.projections, dict) and all(
            spin in bandstructure.projections for spin in spins
        ):
            new_projections = {spin: [[] for _ in range(bandstructure.nb_bands)] for spin in spins}

        n_branches = len(bandstructure.branches)
        new_branches = []

        # This ensures proper format of bandstructure.branches
        processed = []
        for idx in range(n_branches):
            branch = bandstructure.branches[idx]

            if branch["name"] not in processed:
                if tuple(branch["name"].split("-")) in plot_axis:
                    new_branches.append(branch)
                    processed.append(branch["name"])
                else:
                    next_branch = bandstructure.branches[idx + 1]
                    combined = {
                        "start_index": branch["start_index"],
                        "end_index": next_branch["end_index"],
                        "name": f"{branch['name'].split('-')[0]}-{next_branch['name'].split('-')[1]}",
                    }
                    processed.extend((branch["name"], next_branch["name"]))

                    new_branches.append(combined)

        # Obtain new values
        for entry in distances_map:
            branch = new_branches[entry[0]]

            if not entry[1]:
                start = branch["start_index"]
                stop = branch["end_index"] + 1
                step = 1
            else:
                start = branch["end_index"]
                stop = branch["start_index"] - 1 if branch["start_index"] != 0 else None
                step = -1

            # kpoints
            new_kpoints += [point.frac_coords for point in bandstructure.kpoints[start:stop:step]]

            # eigenvals
            for spin in spins:
                for n, band in enumerate(bandstructure.bands[spin]):
                    new_bands[spin][n] = np.concatenate((new_bands[spin][n], band[start:stop:step]))

                # projections
                if has_projections:
                    for n, band in enumerate(bandstructure.projections[spin]):
                        new_projections[spin][n] += band[start:stop:step].tolist()

        if has_projections:
            for spin in spins:
                new_projections[spin] = np.array(new_projections[spin])

        new_labels_dict = {key: point.frac_coords for key, point in bandstructure.labels_dict.items()}

        return BandStructureSymmLine(
            kpoints=new_kpoints,
            eigenvals=new_bands,
            lattice=bandstructure.lattice_rec,
            efermi=bandstructure.efermi,
            labels_dict=new_labels_dict,  # type: ignore[arg-type]
            structure=bandstructure.structure,
            projections=new_projections or {},
        )
