"""CDFT (Conceptual Density Functional Theory) descriptor calculation."""

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

from qspoc.qm.xtb import XTBComputer

from ..utils import (
    get_data,
    get_dataframe,
    deal_with_dataframe,
    file_mode_deal,
    fchk_file_check,
    generate_atoms_desc_name,
    update_wfn_energy,
    safe_get,
    to_eV,
)
from ..constants import GlobalCDFTDescriptorName, AtomCDFTDescriptorName
from ..operators import MultiwfnOperators


class CDFTDescriptorCalculator:
    """Calculate CDFT descriptors using Multiwfn."""

    def __init__(
        self,
        multiwfn_path: str,
        temp_file_dir: Path,
        kill: bool = False,
        file_mode: int = 1,
        acc: str = "normal",
        version: str = "stem",
        env: dict = None,
    ):
        """
        Initialize CDFT descriptor calculator.

        Args:
            multiwfn_path: Path to Multiwfn executable
            temp_file_dir: Directory for temporary files
            kill: Continue calculation after Gaussian not converged
            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.kill = kill
        self.file_mode = file_mode
        self.operators = MultiwfnOperators()
        self.version = version
        self.env = env
        self._set_cdft_acc(acc)

    def _set_cdft_acc(self, acc):
        settings_file = (Path(self.multiwfn_path) / Path("../settings.ini")).resolve()

        params = {"normal": (25, 110), "high": (75, 434), "ultra": (100, 974)}
        rad, sph = params.get(acc, (25, 110))

        lines = settings_file.read_text().splitlines()
        found_rad, found_sph = False, False
        for i, line in enumerate(lines):
            if line.strip().startswith("radpot"):
                lines[i] = f"radpot= {rad}"
                found_rad = True
            elif line.strip().startswith("sphpot"):
                lines[i] = f"sphpot= {sph}"
                found_sph = True
        # 如果没找到则追加
        if not found_rad:
            lines.append(f"radpot= {rad}")
        if not found_sph:
            lines.append(f"sphpot= {sph}")
        settings_file.write_text("\n".join(lines))

    def calculate(
        self,
        calc_data_df: pd.DataFrame,
        n_molecules: int,
        additional_prop_lists: List[List[Tuple]] = None,
        calc_command: str = "b3lyp/def2tzvp",
    ) -> pd.DataFrame:
        """
        Calculate CDFT descriptors for all molecules.

        Args:
            calc_data_df: DataFrame containing calculation data
            n_molecules: Number of molecules
            additional_prop_lists: Additional property lists for CDFT
            calc_command: Calculation command for Multiwfn
            version: Descriptor version ('all' or 'stem')

        Returns:
            DataFrame with CDFT descriptors
        """
        # Steam version: only VIP, VEA, Cond_Fukui+, Cond_Fukui-, Cond_dual_desc
        if self.version == "stem":
            stem_global = ["VIP", "VEA"]
            stem_atom = ["Cond_Fukui+_atom", "Cond_Fukui-_atom", "Cond_dual_desc_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, AtomCDFTDescriptorName, 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 = GlobalCDFTDescriptorName + all_atom_name

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

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

        prop_lists = self._generate_prop_lists(additional_prop_lists, calc_data_df)

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

        for tag, atoms, p, fchk_file in pbar:
            pbar.set_postfix_str(f"Processing {tag}")
            cdft_output = self.temp_file_dir / f"CDFT_temp/CDFT-{tag}.txt"
            use_temp_file = file_mode_deal(self.file_mode, cdft_output)
            fchk_file = fchk_file_check(fchk_file)

            spin_state = f"({p[0][0]} {p[0][1]})\n({p[1][0]} {p[1][1]})\n({p[2][0]} {p[2][1]})"
            dft_calc = self._get_dft_recalc(calc_command=calc_command)

            if not use_temp_file:
                self._run_multiwfn_cdft(fchk_file, calc_command, spin_state, dft_calc, cdft_output)

            self._extract_molecular_descriptors(cdft_descriptors, tag, cdft_output)
            self._extract_atomic_descriptors(cdft_descriptors, tag, atoms, cdft_output)

            if not use_temp_file:
                self._cleanup_temp_files()

        # Clean up fort.7
        (Path.cwd() / "fort.7").unlink() if (Path.cwd() / "fort.7").exists() else None

        return deal_with_dataframe(cdft_descriptors)

    def _generate_prop_lists(self, additional_prop_lists: List[List[Tuple]], calc_data_df: pd.DataFrame) -> List:
        """Generate property lists for CDFT calculations."""
        if additional_prop_lists is None:
            mp_calc = lambda a: 2 if a % 2 == 1 else 1
            prop_calc = lambda x: [(x[0], x[1]), (x[0] - 1, mp_calc(x[1])), (x[0] + 1, mp_calc(x[1]))]
            return [prop_calc(p) for p in calc_data_df["prop_lists"].tolist()]
        else:
            return additional_prop_lists

    def _run_multiwfn_cdft(self, fchk_file: Path, cdft_command: str, spin_state: str, dft_calc: bool, cdft_output: Path) -> None:
        """Run Multiwfn to generate CDFT output."""
        cdft_operator = self.operators.cdft(cdft_command, spin_state, dft_calc)

        if cdft_command == "xtb":
            xtb_computer = XTBComputer(xtb_path=cdft_command)

            name_list = ["N", "N+1", "N-1"]
            for calc_state, tag in zip(spin_state.strip().split("\n"), name_list):
                char = int(calc_state.strip("()").split()[0])
                mult = int(calc_state.strip("()").split()[1])
                new_name = fchk_file.stem.replace("sp_with_", "opt_with_") + ".xyz"
                coord_file = fchk_file.with_name(new_name)

                # TODO: deal with FileNotFoundError: There is no file: tests/test_file/speed_test/2/coords/mol-0-opt_with_g16.xyz
                xtb_output = xtb_computer.singlepoint_energy(tag, coord_file, charge=char, multiplicity=mult)
                xtb_sp_out = xtb_computer.get_sp_output()
                wfn_convert = self.operators.wfn_convert(tag)
                prog = subprocess.Popen([self.multiwfn_path, xtb_output], stdin=subprocess.PIPE, stdout=null)

                prog.communicate(input=wfn_convert)

                update_wfn_energy(xtb_sp_out, f"{tag}.wfn")

                Path(xtb_output).unlink() if Path(xtb_output).exists() else None

        prog = subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env)
        prog.communicate(input=cdft_operator)
        calc_res = prog.returncode

        try:
            assert calc_res == 0
        except:
            print(f"ERROR: Molecule calculation failed with error code {calc_res}")
            if self.kill:
                raise Exception(
                    "Probably gaussian calculation is not converged. Try `SCF(novaracc,noincfock)` or see http://sobereva.com/61."
                )
            else:
                return

        if dft_calc:
            cdft_output_temp = Path.cwd() / "CDFT.txt"
            shutil.move(cdft_output_temp, cdft_output)

    def _extract_molecular_descriptors(self, df: pd.DataFrame, tag: str, cdft_output: Path) -> None:
        """Extract molecular-level CDFT descriptors."""
        # Only extract VIP and VEA for stem version
        if "VIP" in df.columns:
            df.loc[tag, "VIP"] = get_data(cdft_output, "First vertical IP:", "Hartree,")
        if "VEA" in df.columns:
            df.loc[tag, "VEA"] = get_data(cdft_output, "First vertical EA:", "Hartree,")
        if "E_ox_potential" in df.columns:
            df.loc[tag, "E_ox_potential"] = to_eV(get_data(cdft_output, "E(N-1):", "Hartree"), get_data(cdft_output, "E(N):", "Hartree"))
        if "E_red_potential" in df.columns:
            df.loc[tag, "E_red_potential"] = to_eV(get_data(cdft_output, "E(N):", "Hartree"), get_data(cdft_output, "E(N+1):", "Hartree"))
        if "Mulliken_Chi" in df.columns:
            df.loc[tag, "Mulliken_Chi"] = get_data(cdft_output, "Mulliken electronegativity:", "Hartree,")
        if "Chem_potential" in df.columns:
            df.loc[tag, "Chem_potential"] = get_data(cdft_output, "Chemical potential:", "Hartree,")
        if "Elec_hardness" in df.columns:
            df.loc[tag, "Elec_hardness"] = get_data(cdft_output, "Hardness (=fundamental gap):", "Hartree,")
        if "Elec_softness" in df.columns:
            df.loc[tag, "Elec_softness"] = get_data(cdft_output, "Softness:", "Hartree^-1")
        if "Electro_idx" in df.columns:
            df.loc[tag, "Electro_idx"] = get_data(cdft_output, "Electrophilicity index:", "Hartree,")
        if "Nucleo_idx" in df.columns:
            df.loc[tag, "Nucleo_idx"] = get_data(cdft_output, "Nucleophilicity index:", "Hartree,")

    def _extract_atomic_descriptors(self, df: pd.DataFrame, tag: str, atoms: List[int], cdft_output: Path) -> None:
        """Extract atomic-level CDFT descriptors."""
        hirshfeld_fukui_data = get_dataframe(
            cdft_output, "Atom     q(N)    q(N+1)   q(N-1)     f-       f+       f0      CDD\n", "\n \n", remove=(10, None)
        )
        cond_e_n_data = get_dataframe(
            cdft_output, "Atom              Electrophilicity          Nucleophilicity\n", "\n \n", remove=(10, None)
        )
        softness_data = get_dataframe(
            cdft_output, "Atom         s-          s+          s0        s+/s-       s-/s+       s(2)\n", "\n \n", remove=(10, None)
        )

        for i, atom in enumerate(atoms):
            # Only extract stem descriptors if present in columns
            if f"Hirshfeld_atom{i}" in df.columns:
                df.loc[tag, f"Hirshfeld_atom{i}"] = safe_get(hirshfeld_fukui_data, atom, 2)
            if f"Cond_Fukui-_atom{i}" in df.columns:
                df.loc[tag, f"Cond_Fukui-_atom{i}"] = safe_get(hirshfeld_fukui_data, atom, 3)
            if f"Cond_Fukui+_atom{i}" in df.columns:
                df.loc[tag, f"Cond_Fukui+_atom{i}"] = safe_get(hirshfeld_fukui_data, atom, 4)
            if f"Cond_Fukui0_atom{i}" in df.columns:
                df.loc[tag, f"Cond_Fukui0_atom{i}"] = safe_get(hirshfeld_fukui_data, atom, 5)
            if f"Cond_dual_desc_atom{i}" in df.columns:
                df.loc[tag, f"Cond_dual_desc_atom{i}"] = safe_get(hirshfeld_fukui_data, atom, 6)
            if f"Cond_softness-_atom{i}" in df.columns:
                df.loc[tag, f"Cond_softness-_atom{i}"] = safe_get(softness_data, atom, 0)
            if f"Cond_softness+_atom{i}" in df.columns:
                df.loc[tag, f"Cond_softness+_atom{i}"] = safe_get(softness_data, atom, 1)
            if f"Cond_softness0_atom{i}" in df.columns:
                df.loc[tag, f"Cond_softness0_atom{i}"] = safe_get(softness_data, atom, 2)
            if f"Relative_electro_idx_atom{i}" in df.columns:
                df.loc[tag, f"Relative_electro_idx_atom{i}"] = safe_get(softness_data, atom, 3)
            if f"Relative_nucleo_idx_atom{i}" in df.columns:
                df.loc[tag, f"Relative_nucleo_idx_atom{i}"] = safe_get(softness_data, atom, 4)
            if f"Cond_electro_idx_atom{i}" in df.columns:
                df.loc[tag, f"Cond_electro_idx_atom{i}"] = safe_get(cond_e_n_data, atom, 0)
            if f"Cond_nucleo_idx_atom{i}" in df.columns:
                df.loc[tag, f"Cond_nucleo_idx_atom{i}"] = safe_get(cond_e_n_data, atom, 1)

    def _cleanup_temp_files(self) -> None:
        """Clean up temporary .wfn files."""
        for wfn_file in list(Path.cwd().glob("*.wfn")):
            wfn_file.unlink() if wfn_file.exists() else None

    def _get_dft_recalc(self, calc_command: str):

        if all((self.temp_file_dir / f"{f}.wfn").exists() for f in ["N", "N-1", "N+1"]):
            return False
        else:
            return True
