"""Atom properties descriptor calculation."""

from pathlib import Path
from subprocess import DEVNULL as null
from typing import List
import pandas as pd
import subprocess
import shutil
from tqdm import tqdm

from ..utils import (
    get_data,
    get_dataframe,
    get_newly_construct_dataframe,
    deal_with_dataframe,
    file_mode_deal,
    fchk_file_check,
    generate_atoms_desc_name,
    file_append,
    safe_get,
)
from ..constants import AtomPropertiesDescriptorName, GlobalAtomPropertiesDescrptorName
from ..operators import MultiwfnOperators


class AtomDescriptorCalculator:
    """Calculate atom property descriptors using Multiwfn."""

    def __init__(self, multiwfn_path: str, temp_file_dir: Path, file_mode: int = 1, version: str = "stem", env: dict = None):
        """
        Initialize atom descriptor calculator.

        Args:
            multiwfn_path: Path to Multiwfn executable
            temp_file_dir: Directory for temporary files
            file_mode: 0=don't read temp, 1=read if exist, 2=only read temp
        """
        self.multiwfn_path = multiwfn_path
        self.temp_file_dir = temp_file_dir
        self.file_mode = file_mode
        self.operators = MultiwfnOperators()
        self.version = version
        self.env = env

    def calculate(self, calc_data_df: pd.DataFrame, n_molecules: int) -> pd.DataFrame:
        """
        Calculate atom property descriptors for all molecules.

        Args:
            calc_data_df: DataFrame containing calculation data
            n_molecules: Number of molecules
            version: Descriptor version ('all' or 'stem')

        Returns:
            DataFrame with atom property descriptors
        """
        # Steam version: only charge_CM5
        if self.version == "stem":
            stem_atom = ["charge_CM5"]
        else:
            stem_atom = None

        atom_lists = calc_data_df["atom_lists"].tolist()
        atom_lists, all_atom_name = generate_atoms_desc_name(atom_lists, AtomPropertiesDescriptorName, n_molecules)

        # Filter descriptor names based on version
        if self.version == "stem":
            all_descriptor_names = [name for name in all_atom_name if any(pattern in name for pattern in stem_atom)]
        else:
            all_descriptor_names = GlobalAtomPropertiesDescrptorName + all_atom_name

        atom_descriptors = pd.DataFrame(columns=all_descriptor_names, dtype="float64")

        (self.temp_file_dir / "atom_temp").mkdir(parents=True, exist_ok=True)

        pbar = tqdm(
            zip(calc_data_df["name_list"], calc_data_df["atom_lists"], calc_data_df["sp_file"]),
            total=len(calc_data_df),
            desc="[Atom] Descriptors ",
            ncols=100,
            unit="mol",
            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]",
        )

        for tag, atoms, fchk_file in pbar:
            pbar.set_postfix_str(f"Processing {tag}")
            charge_output = self.temp_file_dir / f"atom_temp/charge-{tag}.txt"
            oxidation_output = self.temp_file_dir / f"atom_temp/oxidation-{tag}.txt"
            dipole_output = self.temp_file_dir / f"atom_temp/multipole-{tag}.txt"
            total_esp_output = self.temp_file_dir / f"atom_temp/AtomESP-{tag}.txt"
            ei_output = self.temp_file_dir / f"atom_temp/energy_index-{tag}.txt"

            use_temp_file = file_mode_deal(
                self.file_mode,
                [charge_output, oxidation_output, dipole_output, total_esp_output, ei_output],
            )
            fchk_file = fchk_file_check(fchk_file)

            if not use_temp_file:
                self._run_multiwfn_atom(fchk_file, atoms, charge_output, oxidation_output, dipole_output, total_esp_output, ei_output)

            self._extract_global_descriptors(atom_descriptors, tag, dipole_output)
            self._extract_atomic_descriptors(
                atom_descriptors, tag, atoms, charge_output, oxidation_output, dipole_output, total_esp_output, ei_output
            )

        self._cleanup_temp_files()

        return deal_with_dataframe(atom_descriptors)

    def _run_multiwfn_atom(
        self,
        fchk_file: Path,
        atoms: List[int],
        charge_output: Path,
        oxidation_output: Path,
        dipole_output: Path,
        total_esp_output: Path,
        ei_output: Path,
    ) -> None:
        """Run Multiwfn to generate atom property outputs."""
        charge_origin_output = Path.cwd() / f"{Path(fchk_file).stem}.chg"
        charge_output_working = Path.cwd() / "charge.txt"
        oxidation_output_working = Path.cwd() / "oxidation.txt"
        dipole_output_working = Path.cwd() / "multipole.txt"
        total_esp_output_working = Path.cwd() / "AtomESP.txt"
        ei_output_working = Path.cwd() / "energy_index.txt"

        subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
            input=self.operators.cm5_charge()
        )
        file_append(charge_output_working, charge_origin_output, "CM5:")

        if self.version == "all":
            # Run all charge analyses
            # TODO: do not need to run all analyses
            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
                input=self.operators.hirshfeld_charge()
            )
            file_append(charge_output_working, charge_origin_output, "Hirshfeld:")

            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
                input=self.operators.mulliken_charge()
            )
            file_append(charge_output_working, charge_origin_output, "Mulliken:")

            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
                input=self.operators.scp_charge()
            )
            file_append(charge_output_working, charge_origin_output, "SCPA:")

            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
                input=self.operators.modified_mulliken_charge()
            )
            file_append(charge_output_working, charge_origin_output, "Mod Mulliken:")

            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
                input=self.operators.becke_charge()
            )
            file_append(charge_output_working, charge_origin_output, "Becke:")

            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
                input=self.operators.adch_charge()
            )
            file_append(charge_output_working, charge_origin_output, "ADCH:")

            # Run oxidation state analysis
            with open(oxidation_output_working, "w") as oxidation_outstd:
                subprocess.Popen(
                    [self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=oxidation_outstd, env=self.env
                ).communicate(input=self.operators.loba_oxidation())

            # Run atomic dipole analysis
            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
                input=self.operators.atom_dipole()
            )

            # Run total ESP analysis
            with open(total_esp_output_working, "w") as total_esp_outstd:
                subprocess.Popen(
                    [self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=total_esp_outstd, env=self.env
                ).communicate(input=self.operators.total_esp(atoms))

            # Run energy index analysis
            with open(ei_output_working, "w") as ei_outstd:
                subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=ei_outstd, env=self.env).communicate(
                    input=self.operators.energy_index(atoms)
                )

        # Move working files to temp directory
        shutil.move(charge_output_working, charge_output)
        if self.version == "all":
            shutil.move(oxidation_output_working, oxidation_output)
            shutil.move(dipole_output_working, dipole_output)
            shutil.move(total_esp_output_working, total_esp_output)
            shutil.move(ei_output_working, ei_output)

    def _extract_global_descriptors(self, df: pd.DataFrame, tag: str, dipole_output: Path) -> None:
        """Extract global molecular dipole descriptors."""
        # Only extract if columns exist
        if "dipole" in df.columns:
            df.loc[tag, "dipole"] = get_data(dipole_output, "dipole moment (a.u.&Debye):", "\n")[0]
        if "<r^2>" in df.columns:
            df.loc[tag, "<r^2>"] = get_data(dipole_output, "Molecular electronic spatial extent <r^2>:", "\n")

    def _extract_atomic_descriptors(
        self,
        df: pd.DataFrame,
        tag: str,
        atoms: List[int],
        charge_output: Path,
        oxidation_output: Path,
        dipole_output: Path,
        total_esp_output: Path,
        ei_output: Path,
    ) -> None:
        """Extract atomic-level descriptors."""
        # Charge data
        hirshfeld_data = get_dataframe(charge_output, "Hirshfeld:\n", "\n-------------", remove=(38, None))
        mulliken_data = get_dataframe(charge_output, "Mulliken:\n", "\n-------------", remove=(38, None))
        scp_data = get_dataframe(charge_output, "SCPA:\n", "\n-------------", remove=(38, None))
        mod_mulliken_data = get_dataframe(charge_output, "Mod Mulliken:\n", "\n-------------", remove=(38, None))
        becke_data = get_dataframe(charge_output, "Becke:\n", "\n-------------", remove=(38, None))
        adch_data = get_dataframe(charge_output, "ADCH:\n", "\n-------------", remove=(38, None))
        chelg_data = get_dataframe(charge_output, "CHELPG:\n", "\n-------------", remove=(38, None))
        mk_data = get_dataframe(charge_output, "MK_ESP:\n", "\n-------------", remove=(38, None))
        cm5_data = get_dataframe(charge_output, "CM5:\n", "\n-------------", remove=(38, None))
        resp_data = get_dataframe(charge_output, "RESP:\n", "\n-------------", remove=(38, None))

        # Oxidation state data
        if self.version == "all":
            oxidation_data = get_dataframe(oxidation_output, " Input 0 can exit\n", " The sum of oxidation states", remove=(34, None))

            # Dipole and spatial extent data
            dipole_data = get_newly_construct_dataframe(
                dipole_output,
                ("*****  Atom", "\n \n"),
                ([r" Atomic dipole moments:\n X=.*Y=.*Z=.*Norm=", r"Atomic electronic spatial extent \<r\^2\>:"], ["\n", "\n"]),
                ["atom_dipole", "<r^2>"],
                re_escape_on_data=False,
            )

            # Total ESP data
            total_esp_data = get_newly_construct_dataframe(
                total_esp_output,
                ("------------ Calculate properties at a point ------------", " \n           "),
                ("Total ESP:", "a.u."),
                "total_ESP",
            )

            # Energy index data
            ei_data = get_newly_construct_dataframe(
                ei_output,
                (" Calculate EI index for which atom? e.g. 5", "\n\n"),
                ("The EI index:", "a.u."),
                "energy_index",
            )

        for i, atom in enumerate(atoms):
            # Only extract descriptors that are in DataFrame columns
            if f"charge_Hirshfeld-atom{i}" in df.columns:
                df.loc[tag, f"charge_Hirshfeld-atom{i}"] = safe_get(hirshfeld_data, atom, 0)
            if f"charge_Mulliken-atom{i}" in df.columns:
                df.loc[tag, f"charge_Mulliken-atom{i}"] = safe_get(mulliken_data, atom, 0)
            if f"charge_SCPA-atom{i}" in df.columns:
                df.loc[tag, f"charge_SCPA-atom{i}"] = safe_get(scp_data, atom, 0)
            if f"charge_Modified_Mulliken_by_Bickelhaupt-atom{i}" in df.columns:
                df.loc[tag, f"charge_Modified_Mulliken_by_Bickelhaupt-atom{i}"] = safe_get(mod_mulliken_data, atom, 0)
            if f"charge_Becke-atom{i}" in df.columns:
                df.loc[tag, f"charge_Becke-atom{i}"] = safe_get(becke_data, atom, 0)
            if f"charge_ADCH-atom{i}" in df.columns:
                df.loc[tag, f"charge_ADCH-atom{i}"] = safe_get(adch_data, atom, 0)
            if f"charge_CHELPG-ESP-atom{i}" in df.columns:
                df.loc[tag, f"charge_CHELPG-ESP-atom{i}"] = safe_get(chelg_data, atom, 0)
            if f"charge_MK-ESP-atom{i}" in df.columns:
                df.loc[tag, f"charge_MK-ESP-atom{i}"] = safe_get(mk_data, atom, 0)
            if f"charge_CM5-atom{i}" in df.columns:
                df.loc[tag, f"charge_CM5-atom{i}"] = safe_get(cm5_data, atom, 0)
            if f"charge_RESP-atom{i}" in df.columns:
                df.loc[tag, f"charge_RESP-atom{i}"] = safe_get(resp_data, atom, 0)
            if f"oxidation_state_LOBA-atom{i}" in df.columns:
                df.loc[tag, f"oxidation_state_LOBA-atom{i}"] = safe_get(oxidation_data, atom, 0)
            if f"total_ESP-atom{i}" in df.columns:
                df.loc[tag, f"total_ESP-atom{i}"] = total_esp_data.loc[i, "total_ESP"]
            if f"dipole-atom{i}" in df.columns:
                df.loc[tag, f"dipole-atom{i}"] = safe_get(dipole_data, atom, "atom_dipole")
            if f"<r^2>-atom{i}" in df.columns:
                df.loc[tag, f"<r^2>-atom{i}"] = safe_get(dipole_data, atom, "<r^2>")
            if f"energy_index-atom{i}" in df.columns:
                df.loc[tag, f"energy_index-atom{i}"] = safe_get(ei_data, i, "energy_index")

    def _cleanup_temp_files(self) -> None:
        """Clean up temporary files."""
        for f in self.temp_file_dir.glob("*.chg"):
            f.unlink() if f.exists() else None
        (self.temp_file_dir / "new.fch").unlink() if (self.temp_file_dir / "new.fch").exists() else None
