"""Molecular shape descriptor calculation."""

from pathlib import Path
import pandas as pd
import subprocess
from tqdm import tqdm

from ..utils import get_data, deal_with_dataframe, file_mode_deal, fchk_file_check
from ..constants import GlobalShapeDescriptorName
from ..operators import MultiwfnOperators


class ShapeDescriptorCalculator:
    """Calculate molecular shape 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 shape 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) -> pd.DataFrame:
        """
        Calculate shape descriptors for all molecules.

        Args:
            calc_data_df: DataFrame containing calculation data
            version: Descriptor version ('all' or 'stem')

        Returns:
            DataFrame with shape descriptors
        """
        # Steam version: only Molecular_volume, Molecular_sphericity, Molecular_all_surface_area
        if self.version == "stem":
            descriptor_names = ["Molecular_volume", "Molecular_sphericity", "Molecular_all_surface_area"]
        else:
            descriptor_names = GlobalShapeDescriptorName

        shape_descriptors = pd.DataFrame(columns=descriptor_names, dtype="float64")

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

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

        for tag, fchk_file in pbar:
            pbar.set_postfix_str(f"Processing {tag}")
            shape_output = self.temp_file_dir / f"shape_temp/shape-{tag}.txt"

            use_temp_file = file_mode_deal(self.file_mode, shape_output)
            fchk_file = fchk_file_check(fchk_file)

            if not use_temp_file:
                self._run_multiwfn_shape(fchk_file, shape_output)

            self._extract_descriptors(shape_descriptors, tag, shape_output)

        return deal_with_dataframe(shape_descriptors)

    def _run_multiwfn_shape(self, fchk_file: Path, shape_output: Path) -> None:
        """Run Multiwfn to generate shape output."""
        with open(shape_output, "w") as shape_stdout:
            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=shape_stdout, env=self.env).communicate(
                input=self.operators.shape(self.acc)
            )

    def _extract_descriptors(self, df: pd.DataFrame, tag: str, shape_output: Path) -> None:
        """Extract shape descriptors from Multiwfn output."""
        # Only extract descriptors that are in the DataFrame columns
        if "Molecular_ES_radius_min" in df.columns:
            df.loc[tag, "Molecular_ES_radius_min"] = get_data(shape_output, "The closest distance to the point:", "Bohr")
        if "Molecular_ES_radius_max" in df.columns:
            df.loc[tag, "Molecular_ES_radius_max"] = get_data(shape_output, "The farthest distance to the point:", "Bohr")
        if "Molecular_volume" in df.columns:
            df.loc[tag, "Molecular_volume"] = get_data(shape_output, "Volume enclosed by the isosurface:", "Bohr^3")
        if "Molecular_sphericity" in df.columns:
            df.loc[tag, "Molecular_sphericity"] = get_data(shape_output, "Sphericity:", "\n")
        if "Molecular_all_surface_area" in df.columns:
            df.loc[tag, "Molecular_all_surface_area"] = get_data(shape_output, "Isosurface area:", "Bohr^2")

        if "Molecular_length" in df.columns or "Molecular_widht" in df.columns or "Molecular_height" in df.columns:
            side_lengths = get_data(shape_output, "Length of the three sides:", "Angstrom")
            if "Molecular_length" in df.columns:
                df.loc[tag, "Molecular_length"] = side_lengths[0]
            if "Molecular_widht" in df.columns:
                df.loc[tag, "Molecular_widht"] = side_lengths[1]
            if "Molecular_height" in df.columns:
                df.loc[tag, "Molecular_height"] = side_lengths[2]
