"""Classes for parsing LOBSTER population analysis output files.

Provides classes to read and parse population analysis files generated by LOBSTER,
such as CHARGE.lobster, CHARGE.LCFO.lobster, GROSSPOP.lobster, and GROSSPOP.LCFO.lobster.
Extracts Mulliken and Loewdin charges and gross populations for further analysis.
"""

from __future__ import annotations

import re
from typing import TYPE_CHECKING, Literal

from pymatgen.electronic_structure.core import Spin
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 pymatgen.io.lobster.future.types import LobsterPopulations


class CHARGE(LobsterFile):
    """Parser for CHARGE.lobster and CHARGE.LCFO.lobster files.

    Attributes:
        is_lcfo (bool): Whether the CHARGE file is from LCFO analysis.
        centers (list[str]): Atom labels with indices.
        mulliken (list[float]): Mulliken charges for each atom.
        loewdin (list[float]): Loewdin charges for each atom.
    """

    charge_regex: ClassVar[str] = r"(\d+)\s+([a-zA-Z]+)\s+(\S+)\s+(\S+)"
    is_lcfo: ClassVar[bool] = False

    @version_processor()
    def parse_file(self) -> None:
        """Parse the CHARGE file and extract Mulliken and Loewdin charges.

        Populates the `centers`, `mulliken`, and `loewdin` attributes with data
        extracted from the file.

        Returns:
            None
        """
        self.centers = []
        self.mulliken = []
        self.loewdin = []

        for line in self.iterate_lines():
            if m := re.search(self.charge_regex, line):
                data = m.groups()
                self.centers.append(data[1] + data[0])

                if self.is_lcfo:
                    self.loewdin.append(float(data[2]))
                else:
                    self.mulliken.append(float(data[2]))
                    self.loewdin.append(float(data[3]))

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

        Returns:
            str: Default filename for the CHARGE file, depending on LCFO mode.
        """
        return "CHARGE.lobster" if not cls.is_lcfo else "CHARGE.LCFO.lobster"


class CHARGE_LCFO(CHARGE):
    """Parser for CHARGE.LCFO.lobster files.

    Attributes:
        is_lcfo (bool): Whether the CHARGE file is from LCFO analysis.
    """

    charge_regex: ClassVar[str] = r"(\d+)\s+([a-zA-Z]+)\s+(\S+)"
    is_lcfo: ClassVar[bool] = True


class GROSSPOP(LobsterFile):
    """Parser for GROSSPOP.lobster and GROSSPOP.LCFO.lobster files.

    Attributes:
        is_lcfo (bool): Whether the GROSSPOP file is from LCFO analysis.
        populations (LobsterPopulations): Population dictionaries for each atom.
        spins (list[Spin]): List of spins present in the file.
    """

    is_lcfo: ClassVar[bool] = False

    @version_processor()
    def parse_file(self) -> None:
        """Parse the GROSSPOP file and extract gross population data.

        Populates the `populations` attribute with a nested dictionary structure
        containing Mulliken and Loewdin populations for each atom and orbital.
        """
        self.populations: LobsterPopulations = {}
        self.spins = [Spin.up]

        keys = []

        current_atom = ""
        for line in self.iterate_lines():
            mulliken_key = re.findall(r"(?i)Mulliken GP", line)
            loewdin_key = re.findall(r"(?i)Loewdin GP", line)

            if mulliken_key:
                keys.append("mulliken")
            if loewdin_key:
                keys.append("loewdin")

            if len(loewdin_key) > 1 or len(mulliken_key) > 1:
                self.spins.append(Spin.down)

            if "total" in line.lower():
                pass
            elif data := re.search(
                r"(?:(\d+)\s+)?(?:([a-zA-Z]+)\s+)?(\S+)\s+([0-9\.]+)\s+([0-9\.]+)(?:\s+)?([0-9\.]+)?(?:\s+)?([0-9\.]+)?",
                line,
            ):
                groups = data.groups()
                populations: dict[Spin, dict[Literal["mulliken", "loewdin"], float]] = {}
                if groups[0] is not None and groups[1] is not None:
                    current_atom = groups[1] + groups[0]
                    self.populations[current_atom] = {}
                if groups[2] is not None:
                    self.populations[current_atom][groups[2]] = populations
                else:
                    continue

                for i, spin in enumerate(self.spins):
                    populations[spin] = {}
                    for j, key in enumerate(keys):
                        populations[spin][key] = float(groups[3 + i + j * len(self.spins)])

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

        Returns:
            str: Default filename for the GROSSPOP file, depending on LCFO mode.
        """
        return "GROSSPOP.lobster" if not cls.is_lcfo else "GROSSPOP.LCFO.lobster"


class GROSSPOP_LCFO(GROSSPOP):
    """Parser for GROSSPOP.LCFO.lobster files.

    Attributes:
        is_lcfo (bool): Whether the GROSSPOP file is in LCFO format.
        populations (LobsterPopulations): Population dictionaries for each atom.
    """

    is_lcfo: ClassVar[bool] = True
