from __future__ import annotations

import re
from collections import defaultdict
from itertools import islice
from typing import TYPE_CHECKING, Any

import numpy as np

from pymatgen.core.periodic_table import Element
from pymatgen.electronic_structure.core import Spin
from pymatgen.electronic_structure.dos import Dos
from pymatgen.io.lobster.future.core import LobsterFile
from pymatgen.io.lobster.future.versioning import version_processor

if TYPE_CHECKING:
    from typing import ClassVar

    from numpy import floating
    from numpy.typing import NDArray

    from pymatgen.util.typing import PathLike


class DOSCAR(LobsterFile):
    """Represents LOBSTER's projected DOS and local projected DOS.

    This class parses and stores data from the DOSCAR file generated by LOBSTER,
    which contains information about the total and projected density of states
    (DOS) for a quantum-chemical calculation performed with VASP.

    Attributes:
        completedos (LobsterCompleteDos): Complete DOS data.
        pdos (list[dict[str, dict[Spin, np.ndarray]]]): Projected DOS data.
            Access as `pdos[atomindex]['orbitalstring'][Spin.up/Spin.down]`.
        tdos (Dos): Total density of states.
        energies (np.ndarray): Energies at which the DOS was calculated
            (in eV, relative to Efermi).
        tdensities (dict[Spin, np.ndarray]): Total density arrays for each spin channel.
        itdensities (dict[Spin, np.ndarray]): Integrated total density arrays for each spin channel.
        is_spin_polarized (bool): Whether the system is spin polarized.
    """

    is_lcfo: ClassVar[bool] = False

    def __init__(
        self,
        filename: PathLike | None = None,
        process_immediately: bool = True,
        lobster_version: str | None = None,
    ) -> None:
        """Initialize a DOSCAR object.

        Args:
            filename (PathLike | None): Path to the DOSCAR file, typically "DOSCAR.lobster".
            process_immediately (bool): Whether to process the file immediately upon initialization.
            lobster_version (str | None): LOBSTER version string to use for parsing. If None, attempts to detect

        Raises:
            ValueError: If neither `structure_file` nor `structure` is provided.
        """
        self.projected_dos: dict[str, dict[str, Dos]]
        self.total_dos: Dos
        self.integrated_total_dos: Dos

        super().__init__(
            filename=filename,
            process_immediately=process_immediately,
            lobster_version=lobster_version,
        )

    @version_processor()
    def process(self) -> None:
        """Process the DOSCAR file and extract DOS data.

        Parses the DOSCAR file to extract total DOS, projected DOS, energies,
        and integrated densities. Sets the appropriate attributes based on
        whether the calculation is spin-polarized.

        Raises:
            ValueError: If the DOSCAR file format is invalid or spin polarization
                cannot be determined.
        """
        total_dos, integrated_total_dos = {}, {}
        data: list[NDArray] = []

        centers: list[str] = []
        orbitals: list[list[str]] = []

        self.spins = [Spin.up]

        header_regex = r"\s*\S+\s+\S+\s+(\d+)\s+(\S+)\s+1\.0+(?:;(.*);(.*))?"

        efermi = None
        ndos = 0

        lines_iter = iter(self.iterate_lines())

        center_counts: defaultdict[str, int] = defaultdict(int)

        for line in islice(lines_iter, 5, None):
            if match := re.match(header_regex, line):
                ndos = int(match.group(1))

                if efermi is None:
                    efermi = float(match.group(2))

                if center_match := match.group(3):
                    center_match = center_match.strip()

                    if center_match.startswith("Z="):
                        center_match = Element.from_Z(int(center_match.split()[-1])).symbol

                    center_counts[center_match] += 1

                    separator = "_" if self.is_lcfo else ""

                    centers.append(f"{center_match}{separator}{center_counts[center_match]}")

                if orbital_match := match.group(4):
                    orbitals += [[orb.strip() for orb in orbital_match.split()]]

            tmp_dos = []
            if line.strip():
                for _ in range(ndos):
                    row_parts = next(lines_iter).split()
                    tmp_dos.append(row_parts)

                data.append(np.array(tmp_dos, dtype=float))

        if len(data[0][0, :]) == 5:
            self.spins.append(Spin.down)
        elif len(data[0][0, :]) != 3:
            raise ValueError("There is something wrong with the DOSCAR. Can't extract spin polarization.")

        if efermi is None:
            raise ValueError("There is something wrong with the DOSCAR. Can't find efermi.")

        energies = data[0][:, 0]
        projected_dos: dict[str, Any] = {}

        if self.is_spin_polarized:
            total_dos[Spin.up] = data[0][:, 1]
            total_dos[Spin.down] = data[0][:, 2]

            integrated_total_dos[Spin.up] = data[0][:, 3]
            integrated_total_dos[Spin.down] = data[0][:, 4]
        else:
            total_dos[Spin.up] = data[0][:, 1]
            integrated_total_dos[Spin.up] = data[0][:, 2]

        for atom_counter in range(len(data) - 1):
            block_data = data[atom_counter + 1]
            center = centers[atom_counter]

            if center not in projected_dos:
                projected_dos[center] = {}

            for spin_index, spin in enumerate(self.spins):
                for orbital_index, row in enumerate(range(spin_index + 1, block_data.shape[1], len(self.spins))):
                    orbital = orbitals[atom_counter][orbital_index]

                    if orbital not in projected_dos[center]:
                        projected_dos[center][orbital] = {}

                    projected_dos[center][orbital][spin] = block_data[:, row]

        for center, orbitals in projected_dos.items():
            for orbital, dos in orbitals.items():
                projected_dos[center][orbital] = Dos(efermi, energies, dos)

        self.projected_dos = dict(projected_dos)
        self.total_dos = Dos(efermi, energies, total_dos)
        self.integrated_total_dos = Dos(efermi, energies, integrated_total_dos)

    @property
    def efermi(self) -> float:
        """Fermi energy in eV."""
        return self.total_dos.efermi

    @property
    def energies(self) -> NDArray[floating]:
        """Energies at which the DOS was calculated (in eV, relative to efermi)."""
        return self.total_dos.energies

    @property
    def is_spin_polarized(self) -> bool:
        """Whether the system is spin polarized.

        Returns:
            bool: True if the system is spin polarized, False otherwise.
        """
        spins = self.spins
        return spins is not None and len(spins) == 2

    @classmethod
    def get_default_filename(cls) -> str:
        """Get the default filename for the DOSCAR.

        Returns:
            str: Default filename string. "DOSCAR.lobster" for regular DOSCAR,
            "DOSCAR.LCFO.lobster" for LCFO analysis.
        """
        return "DOSCAR.lobster" if not cls.is_lcfo else "DOSCAR.LCFO.lobster"


class DOSCAR_LCFO(DOSCAR):
    """Represents LOBSTER's projected DOS and local projected DOS for LCFO analysis.

    This class handles DOSCAR files generated from LCFO analysis, which have
    a different format than regular DOSCAR files.

    Attributes:
        is_lcfo (bool): Class variable indicating this is for LCFO analysis. Always True.
    """

    is_lcfo: ClassVar[bool] = True
