"""Molecular orbital descriptor calculation."""

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

from ..utils import (
    get_data,
    get_newly_construct_dataframe,
    deal_with_dataframe,
    file_mode_deal,
    fchk_file_check,
    generate_atoms_desc_name,
    safe_get,
    str2df,
)
from ..constants import GlobalOrbitDescriptorName, AtomOrbitDescriptorName
from ..operators import MultiwfnOperators


class OrbitDescriptorCalculator:
    """Calculate molecular orbital descriptors using Multiwfn."""

    def __init__(self, multiwfn_path: str, temp_file_dir: Path, file_mode: int = 1, version: str = "stem", env: dict = None):
        """
        Initialize orbital 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 orbital 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 orbital descriptors
        """
        # Steam version: only HOMO_LUMO_energy_gap, HOMO_ODI, LUMO_ODI
        if self.version == "stem":
            stem_global = ["HOMO_LUMO_energy_gap", "HOMO_ODI", "LUMO_ODI"]
        else:
            stem_global = None

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

        # Filter descriptor names based on version
        if self.version == "stem":
            all_descriptor_names = stem_global  # No atomic descriptors for stem version
        else:
            all_descriptor_names = GlobalOrbitDescriptorName + all_atom_name

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

        (self.temp_file_dir / "orbit_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="[Orbit] 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}")
            orbit_output = self.temp_file_dir / f"orbit_temp/{tag}-orbit.txt"

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

            if not use_temp_file:
                self._run_multiwfn_orbit(fchk_file, orbit_output)

            self._extract_descriptors(orbit_descriptors, tag, atoms, orbit_output)

        return deal_with_dataframe(orbit_descriptors)

    def _run_multiwfn_orbit(self, fchk_file: Path, orbit_output: Path) -> None:
        """Run Multiwfn to generate orbital output."""
        with open(orbit_output, "w") as orbit_stdout:
            subprocess.Popen([self.multiwfn_path, str(fchk_file)], stdin=subprocess.PIPE, stdout=orbit_stdout, env=self.env).communicate(
                input=self.operators.orbit()
            )

    def _extract_descriptors(self, df: pd.DataFrame, tag: str, atoms: List[int], orbit_output: Path) -> None:
        """Extract orbital descriptors from Multiwfn output."""
        # Extract ODI data
        odi_data = get_newly_construct_dataframe(
            orbit_output,
            ("Orbital delocalization index:", "\n"),
            ("Orbital delocalization index:", "\n"),
            desc_list="ODI",
            re_escape_on_paragraph=False,
        )

        # Extract orbital energies
        with orbit_output.open("r") as f:
            orbit_text = f.read()

        if "alpha" in orbit_text:
            # Unrestricted calculation
            homo_energy = (
                get_data(orbit_output, "alpha-HOMO, energy:", "a.u.") + get_data(orbit_output, "beta-HOMO, energy:", "a.u.")
            ) * 0.5
            lumo_energy = (
                get_data(orbit_output, "alpha-LUMO, energy:", "a.u.") + get_data(orbit_output, "beta-LUMO, energy:", "a.u.")
            ) * 0.5
            gap = (
                get_data(orbit_output, "HOMO-LUMO gap of alpha orbitals:", "a.u.")
                + get_data(orbit_output, "HOMO-LUMO gap of beta orbitals:", "a.u.")
            ) * 0.5
        else:
            # Restricted calculation
            homo_energy = get_data(orbit_output, "HOMO, energy:", "a.u.")
            lumo_energy = get_data(orbit_output, "LUMO, energy:", "a.u.")
            gap = get_data(orbit_output, "HOMO-LUMO gap:", "a.u.")

        # Set global orbital descriptors (only if columns exist)
        if "HOMO_energy" in df.columns:
            df.loc[tag, "HOMO_energy"] = homo_energy
        if "LUMO_energy" in df.columns:
            df.loc[tag, "LUMO_energy"] = lumo_energy
        if "HOMO_LUMO_energy_gap" in df.columns:
            df.loc[tag, "HOMO_LUMO_energy_gap"] = gap
        if "HOMO_ODI" in df.columns:
            df.loc[tag, "HOMO_ODI"] = safe_get(odi_data, 0, "ODI")
        if "LUMO_ODI" in df.columns:
            df.loc[tag, "LUMO_ODI"] = safe_get(odi_data, 1, "ODI")

        # Set atomic orbital descriptors (only if columns exist)
        for i, atom in enumerate(atoms):
            if f"HOMO_ODI-atom{i}" in df.columns:
                context = re.findall("Composition of each atom:\n(.*?)\n\n", orbit_text, re.DOTALL)
                homo_odi_atom_data = str2df(context[0], remove=(17, 30))
                df.loc[tag, f"HOMO_ODI-atom{i}"] = safe_get(homo_odi_atom_data, atom, 0)
            if f"LUMO_ODI-atom{i}" in df.columns:
                context = re.findall("Composition of each atom:\n(.*?)\n\n", orbit_text, re.DOTALL)
                lumo_odi_atom_data = str2df(context[-1], remove=(17, 30))
                df.loc[tag, f"LUMO_ODI-atom{i}"] = safe_get(lumo_odi_atom_data, atom, 0)
