"""ESP (Electrostatic Potential) descriptor calculation."""

from pathlib import Path
from typing import List
import pandas as pd
from tqdm import tqdm

from ..utils import (
    get_data,
    get_dataframe,
    deal_with_dataframe,
    file_mode_deal,
    fchk_file_check,
    generate_atoms_desc_name,
    safe_get,
)
from ..constants import GlobalESPDescriptorName, AtomESPDescriptorName
from ..operators import MultiwfnOperators


class ESPDescriptorCalculator:
    """Calculate ESP descriptors using Multiwfn."""

    def __init__(
        self, multiwfn_path: str, temp_file_dir: Path, file_mode: int = 1, acc: str = "normal", version: str = "stem", env: dict = None
    ):
        """
        Initialize ESP 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.acc = acc
        self.version = version
        self.env = env

    def calculate(self, calc_data_df: pd.DataFrame, n_molecules: int) -> pd.DataFrame:
        """
        Calculate ESP 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 ESP descriptors
        """
        # Steam version: specific descriptors
        if self.version == "stem":
            stem_global = ["min_ESP", "max_ESP", "MPI", "char_balance", "area_pSurface", "area_nSurface", "aver_ALIE", "aver_LEAE"]
            stem_atom = ["max_ESP_atom", "min_ESP_atom", "icp_atom"]
        else:
            stem_global = None
            stem_atom = None

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

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

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

        (self.temp_file_dir / "ESP_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="[ESP] 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}")
            esp_output = self.temp_file_dir / f"ESP_temp/{tag}-ESP_output.txt"
            alie_output = self.temp_file_dir / f"ESP_temp/{tag}-ALIE_output.txt"
            leae_output = self.temp_file_dir / f"ESP_temp/{tag}-LEAE_output.txt"
            lea_output = self.temp_file_dir / f"ESP_temp/{tag}-LEA_output.txt"

            use_temp_file = file_mode_deal(self.file_mode, [esp_output, alie_output, leae_output, lea_output])
            fchk_file = fchk_file_check(fchk_file)

            if not use_temp_file:
                self._run_multiwfn_esp(fchk_file, esp_output, alie_output, leae_output, lea_output)

            self._extract_molecular_descriptors(esp_descriptors, tag, esp_output, alie_output, leae_output, lea_output)
            self._extract_atomic_descriptors(esp_descriptors, tag, atoms, esp_output)

        return deal_with_dataframe(esp_descriptors)

    def _run_multiwfn_esp(self, fchk_file: Path, esp_output: Path, alie_output: Path, leae_output: Path, lea_output: Path) -> None:
        """Run Multiwfn to generate ESP outputs."""
        import subprocess

        # Run ESP and LEA together
        with open(esp_output, "w") as esp_stdout, open(lea_output, "w") as lea_stdout:
            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=esp_stdout, env=self.env).communicate(
                input=self.operators.esp(self.acc)
            )
            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=lea_stdout, env=self.env).communicate(
                input=self.operators.lea(self.acc)
            )

        # Run ALIE and LEAE together
        with open(alie_output, "w") as alie_stdout, open(leae_output, "w") as leae_stdout:
            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=alie_stdout, env=self.env).communicate(
                input=self.operators.alie(self.acc)
            )
            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=leae_stdout, env=self.env).communicate(
                input=self.operators.leae(self.acc)
            )

    def _extract_molecular_descriptors(
        self, df: pd.DataFrame, tag: str, esp_output: Path, alie_output: Path, leae_output: Path, lea_output: Path
    ) -> None:
        """Extract molecular-level ESP descriptors."""
        # Surface ESP properties
        if "min_ESP" in df.columns:
            df.loc[tag, "min_ESP"] = get_data(esp_output, "Global surface minimum:", "a.u.")
        if "max_ESP" in df.columns:
            df.loc[tag, "max_ESP"] = get_data(esp_output, "Global surface maximum:", "a.u.")
        if "aver_pESP" in df.columns:
            df.loc[tag, "aver_pESP"] = get_data(esp_output, "Positive average value:", "a.u.")
        if "aver_nESP" in df.columns:
            df.loc[tag, "aver_nESP"] = get_data(esp_output, "Negative average value:", "a.u.")
        if "aver_oESP" in df.columns:
            df.loc[tag, "aver_oESP"] = get_data(esp_output, "Overall average value:", "a.u.")
        if "vari_p_ESP" in df.columns:
            df.loc[tag, "vari_p_ESP"] = get_data(esp_output, "Positive variance:", "a.u.^2")
        if "vari_n_ESP" in df.columns:
            df.loc[tag, "vari_n_ESP"] = get_data(esp_output, "Negative variance:", "a.u.^2")
        if "vari_o_ESP" in df.columns:
            df.loc[tag, "vari_o_ESP"] = get_data(esp_output, "Overall variance (sigma^2_tot):", "a.u.^2")
        if "char_balance" in df.columns:
            df.loc[tag, "char_balance"] = get_data(esp_output, "Balance of charges (nu):", "\n")
        if "MPI" in df.columns:
            df.loc[tag, "MPI"] = get_data(esp_output, "Molecular polarity index (MPI):", "eV")
        if "area_pSurface" in df.columns:
            df.loc[tag, "area_pSurface"] = get_data(esp_output, "Positive surface area:", "Bohr^2")
        if "area_nSurface" in df.columns:
            df.loc[tag, "area_nSurface"] = get_data(esp_output, "Negative surface area:", "Bohr^2")

        # ALIE properties
        if "aver_ALIE" in df.columns:
            df.loc[tag, "aver_ALIE"] = get_data(alie_output, "Average value:", "a.u.")
        if "vari_ALIE" in df.columns:
            df.loc[tag, "vari_ALIE"] = get_data(alie_output, "Variance:", "a.u.^2")

        # LEAE properties
        if "aver_LEAE" in df.columns:
            df.loc[tag, "aver_LEAE"] = get_data(leae_output, "Overall average value:", "a.u.")
        if "vari_LEAE" in df.columns:
            df.loc[tag, "vari_LEAE"] = get_data(leae_output, "Overall variance:", "a.u.")

        # LEA properties
        if "aver_LEA" in df.columns:
            df.loc[tag, "aver_LEA"] = get_data(lea_output, "Overall average value:", "a.u.")
        if "vari_LEA" in df.columns:
            df.loc[tag, "vari_LEA"] = get_data(lea_output, "Overall variance:", "a.u.")

    def _extract_atomic_descriptors(self, df: pd.DataFrame, tag: str, atoms: List[int], esp_output: Path) -> None:
        """Extract atomic-level ESP descriptors."""
        # Max/min/variance data
        max_min_data = get_dataframe(
            esp_output, "Atom#    All/Positive/Negative area (Ang^2)  Minimal value   Maximal value\n", "\n \n", reset_index=True
        )
        aver_vari_data = get_dataframe(
            esp_output, "Atom#    All/Positive/Negative average       All/Positive/Negative variance\n", "\n \n", reset_index=True
        )
        icp_nu_data = get_dataframe(esp_output, "  Atom#           Pi              nu         nu*sigma^2\n", "\n\n", reset_index=True)

        for i, atom in enumerate(atoms):
            # Only extract stem descriptors if present in columns
            if atom in max_min_data.index:
                if f"max_ESP_atom{i}" in df.columns:
                    df.loc[tag, f"max_ESP_atom{i}"] = safe_get(max_min_data, atom, 4)
                if f"min_ESP_atom{i}" in df.columns:
                    df.loc[tag, f"min_ESP_atom{i}"] = safe_get(max_min_data, atom, 5)
            if atom in aver_vari_data.index:
                if f"aver_ESP_atom{i}" in df.columns:
                    df.loc[tag, f"aver_ESP_atom{i}"] = safe_get(aver_vari_data, atom, 1)
                if f"vari_ESP_atom{i}" in df.columns:
                    df.loc[tag, f"vari_ESP_atom{i}"] = safe_get(aver_vari_data, atom, 4)
            if atom in icp_nu_data.index:
                if f"icp_atom{i}" in df.columns:
                    df.loc[tag, f"icp_atom{i}"] = safe_get(icp_nu_data, atom, 1)
                if f"nu_atom{i}" in df.columns:
                    df.loc[tag, f"nu_atom{i}"] = safe_get(icp_nu_data, atom, 2)
