"""Internal coordinate (bond, angle, dihedral) descriptor calculation."""

from pathlib import Path
import numpy as np
import pandas as pd
from tqdm import tqdm

from qspoc.converter import xyz2coord

from ..utils import (
    bond_length_calc,
    angle_calc,
    dihedral_angle_calc,
    deal_with_dataframe,
    internal_name_generate,
)
from ..operators import MultiwfnOperators


class InternalCoordinateCalculator:
    """Calculate internal coordinate descriptors using Multiwfn."""

    def __init__(self):
        self.operators = MultiwfnOperators()

    def calculate(self, calc_data_df: pd.DataFrame, temp_coords_dir: str) -> pd.DataFrame:
        """
        Calculate internal coordinate descriptors for all molecules.
        Compatible with missing bond/angle/dihedral lists.
        """
        cols_to_check = ["bond_lists", "angle_lists", "dihedral_lists"]
        available_cols = [c for c in cols_to_check if c in calc_data_df.columns]

        def get_max_len(col_name):
            if col_name not in available_cols:
                return 0
            lengths = calc_data_df[col_name].dropna().apply(len)
            return lengths.max() if not lengths.empty else 0

        n_bonds = get_max_len("bond_lists")
        n_angles = get_max_len("angle_lists")
        n_dihedrals = get_max_len("dihedral_lists")
        tags = calc_data_df["name_list"]

        bond_data = calc_data_df.get("bond_lists", [None] * len(calc_data_df))
        angle_data = calc_data_df.get("angle_lists", [None] * len(calc_data_df))
        dihedral_data = calc_data_df.get("dihedral_lists", [None] * len(calc_data_df))
        sp_files = calc_data_df.get("sp_file", [None] * len(calc_data_df))
        results_list = []
        pbar = tqdm(
            zip(tags, bond_data, angle_data, dihedral_data, sp_files),
            total=len(calc_data_df),
            desc="[Inter] Coordinates",
            ncols=100,
            unit="mol",
            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]",
        )
        for tag, bonds, angles, dihedrals, fchk_file in pbar:
            pbar.set_postfix_str(f"Processing {tag}")

            # 初始化当前行的结果字典
            current_row = {"tag": tag}
            # 查找坐标文件
            try:
                opt_files = list(Path(temp_coords_dir).glob(f"{tag}-opt_with_*.xyz"))
                if not opt_files:
                    results_list.append(current_row)
                    continue

                # 确定优化方法
                opt_methods = [of.stem.split("_")[-1] for of in opt_files]
                opt_method = "g16" if "g16" in opt_methods else opt_methods[0]
                coord_path = Path(temp_coords_dir) / f"{tag}-opt_with_{opt_method}.xyz"
                coord = xyz2coord(str(coord_path))
            except Exception as e:
                print(f"Error loading coordinates for {tag}: {e}")
                results_list.append(current_row)
                continue
            # --- 计算键长 ---
            if isinstance(bonds, (list, np.ndarray)) and bonds != [[]]:
                for i, b in enumerate(bonds):
                    p1 = coord.loc[b[0], "coord"].astype(float).tolist()
                    p2 = coord.loc[b[1], "coord"].astype(float).tolist()
                    current_row[f"bond_{i}"] = bond_length_calc(p1, p2)
            # --- 计算角度 ---
            if isinstance(angles, (list, np.ndarray)) and angles != [[]]:
                for i, a in enumerate(angles):
                    p1 = coord.loc[a[0], "coord"].astype(float).tolist()
                    p2 = coord.loc[a[1], "coord"].astype(float).tolist()
                    p3 = coord.loc[a[2], "coord"].astype(float).tolist()
                    current_row[f"angle_{i}"] = angle_calc(p1, p2, p3)
            # --- 计算二面角 ---
            if isinstance(dihedrals, (list, np.ndarray)) and dihedrals != [[]]:
                for i, d in enumerate(dihedrals):
                    p1 = coord.loc[d[0], "coord"].astype(float).tolist()
                    p2 = coord.loc[d[1], "coord"].astype(float).tolist()
                    p3 = coord.loc[d[2], "coord"].astype(float).tolist()
                    p4 = coord.loc[d[3], "coord"].astype(float).tolist()
                    current_row[f"dihedral_{i}"] = dihedral_angle_calc(p1, p2, p3, p4)
            results_list.append(current_row)

        internal_descriptors = pd.DataFrame(results_list)
        if "tag" in internal_descriptors.columns:
            internal_descriptors.set_index("tag", inplace=True)

        all_desc_names = internal_name_generate(n_bonds, n_angles, n_dihedrals)
        for col in all_desc_names:
            if col not in internal_descriptors.columns:
                internal_descriptors[col] = np.nan

        internal_descriptors = internal_descriptors[all_desc_names]
        return deal_with_dataframe(internal_descriptors)
