"""Bond properties descriptor calculation."""

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

from ..utils import (
    get_newly_construct_dataframe,
    deal_with_dataframe,
    file_mode_deal,
    fchk_file_check,
    generate_desc_names,
    deal_with_seperated_table,
    file_append,
    safe_get,
)
from ..constants import BondPropertiesDescriptorName
from ..operators import MultiwfnOperators


class BondDescriptorCalculator:
    """Calculate bond 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 bond descriptor calculator.

        Args:
            multiwfn_path: Path to Multiwfn executable
            temp_file_dir: Directory for temporary files
            bond_list: List of bond definitions [(i, j), ...]
            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

        # Bond order data strings
        self._bond_order_strings = {
            "mayer": "mayer bond order:\n Note: The diagonal elements are the sum of corresponding row elements\n ***************************** Bond order matrix *****************************",
            "wiberg": "wiberg bond order:\n Note: The diagonal elements are the sum of corresponding row elements\n ***************************** Bond order matrix *****************************",
            "mulliken": "mulliken bond order:\n Note:The diagonal elements are the sum of corresponding row elements\n ************************ Mulliken bond order matrix ************************",
            "fuzzy": "fuzzy bond order:\n ********************* Total delocalization index matrix *********************",
            "laplacian": "laplacian bond order:\n ************************ Laplacian bond order matrix ************************",
        }

    def calculate(self, calc_data_df: pd.DataFrame, n_bonds: int) -> pd.DataFrame:
        """
        Calculate bond 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 bond property descriptors
        """
        # Steam version: only Wiberg_bond_order, Bond_dipole
        if self.version == "stem":
            stem_descriptors = ["Wiberg_bond_order_bond", "Bond_dipole_bond"]
        else:
            stem_descriptors = None

        self.n_bonds = len(calc_data_df.loc[0, "bond_lists"])

        # Filter descriptor names based on version
        if self.version == "stem":
            all_bond_name = generate_desc_names(stem_descriptors, n_bonds)
        else:
            all_bond_name = generate_desc_names(BondPropertiesDescriptorName, n_bonds)

        bond_descriptors = pd.DataFrame(columns=all_bond_name, dtype="float64")

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

        pbar = tqdm(
            zip(calc_data_df["name_list"], calc_data_df["bond_lists"], calc_data_df["sp_file"]),
            total=len(calc_data_df),
            desc="[Bond] Descriptors ",
            ncols=100,
            unit="mol",
            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]",
        )
        for tag, bonds, fchk_file in pbar:
            pbar.set_postfix_str(f"Processing {tag}")
            order_output = self.temp_file_dir / f"bond_temp/bond_order-{tag}.txt"
            dipole_output = self.temp_file_dir / f"bond_temp/bond_dipole-{tag}.txt"

            use_temp_file = file_mode_deal(self.file_mode, [order_output, dipole_output])
            fchk_file = fchk_file_check(fchk_file)

            dipole_operator = self.operators.bond_dipole(bonds)

            if not use_temp_file:
                self._run_multiwfn_bond_orders(fchk_file, order_output)
                self._run_multiwfn_bond_dipole(fchk_file, dipole_output, dipole_operator)

            self._extract_bond_order_descriptors(bond_descriptors, tag, bonds, order_output)
            self._extract_bond_dipole_descriptors(bond_descriptors, tag, dipole_output)

        return deal_with_dataframe(bond_descriptors)

    def _run_multiwfn_bond_orders(self, fchk_file: Path, order_output: Path) -> None:
        """Run Multiwfn to generate bond order outputs."""
        order_origin_output = Path.cwd() / "bndmat.txt"
        order_output_working = Path.cwd() / "bond_order.txt"
        order_output_working.touch()

        # Run all bond order analyses
        subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
            input=self.operators.wiberg_bond_order()
        )
        file_append(order_output_working, order_origin_output, "wiberg bond order:")

        if self.version == "all":
            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
                input=self.operators.mayer_bond_order()
            )
            file_append(order_output_working, order_origin_output, "mayer bond order:")

            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
                input=self.operators.mulliken_bond_order()
            )
            file_append(order_output_working, order_origin_output, "mulliken bond order:")

            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
                input=self.operators.fuzzy_bond_order()
            )
            file_append(order_output_working, order_origin_output, "fuzzy bond order:")

            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=null, env=self.env).communicate(
                input=self.operators.laplacian_bond_order()
            )
            file_append(order_output_working, order_origin_output, "laplacian bond order:")

        import shutil

        shutil.move(order_output_working, order_output)

    def _run_multiwfn_bond_dipole(self, fchk_file: Path, dipole_output: Path, dipole_operator: bytes) -> None:
        """Run Multiwfn to generate bond dipole outputs."""
        with open(dipole_output, "w") as dipole_outstd:
            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=dipole_outstd, env=self.env).communicate(
                input=dipole_operator
            )

    def _extract_bond_order_descriptors(self, df: pd.DataFrame, tag: str, bonds: List[tuple], order_output: Path) -> None:
        """Extract bond order descriptors from Multiwfn output."""
        wiberg_order_data = deal_with_seperated_table(order_output, (self._bond_order_strings["wiberg"], "--------------"), "\n        ")
        for j, b in enumerate(bonds):
            if f"Wiberg_bond_order_bond{j}" in df.columns:
                df.loc[tag, f"Wiberg_bond_order_bond{j}"] = safe_get(wiberg_order_data, b[0] + 1, b[1] + 1)

        if self.version == "all":
            mayer_order_data = deal_with_seperated_table(order_output, (self._bond_order_strings["mayer"], "--------------"), "\n        ")

            mulliken_order_data = deal_with_seperated_table(
                order_output, (self._bond_order_strings["mulliken"], "--------------"), "\n        "
            )
            fuzzy_order_data = deal_with_seperated_table(order_output, (self._bond_order_strings["fuzzy"], "--------------"), "\n        ")
            laplacian_order_data = deal_with_seperated_table(
                order_output, (self._bond_order_strings["laplacian"], "--------------"), "\n        "
            )

            for j, b in enumerate(bonds):
                # Only extract bond orders that are in DataFrame columns
                if f"Mayer_bond_order_bond{j}" in df.columns:
                    df.loc[tag, f"Mayer_bond_order_bond{j}"] = safe_get(mayer_order_data, b[0] + 1, b[1] + 1)
                if f"Mulliken_bond_order_bond{j}" in df.columns:
                    df.loc[tag, f"Mulliken_bond_order_bond{j}"] = safe_get(mulliken_order_data, b[0] + 1, b[1] + 1)
                if f"Fuzzy_bond_order_bond{j}" in df.columns:
                    df.loc[tag, f"Fuzzy_bond_order_bond{j}"] = safe_get(fuzzy_order_data, b[0] + 1, b[1] + 1)
                if f"Laplacian_bond_order_bond{j}" in df.columns:
                    df.loc[tag, f"Laplacian_bond_order_bond{j}"] = safe_get(laplacian_order_data, b[0] + 1, b[1] + 1)

    def _extract_bond_dipole_descriptors(self, df: pd.DataFrame, tag: str, dipole_output: Path) -> None:
        """Extract bond dipole descriptors from Multiwfn output."""
        dipole_data = get_newly_construct_dataframe(
            dipole_output,
            (" Input index of two atoms, e.g.", "Contribution to system dipole moment"),
            (r" Bond dipole moment \(a\.u\.\):\n  X=.*Y=.*Z=.*Norm=", "\n"),
            "bond_dipole",
            re_escape_on_data=False,
        )

        for j in range(self.n_bonds):
            # Only extract bond dipole if column exists
            if f"Bond_dipole_bond{j}" in df.columns:
                df.loc[tag, f"Bond_dipole_bond{j}"] = dipole_data.loc[j, "bond_dipole"]
