from ast import literal_eval
from pathlib import Path
import subprocess
from typing import List, Optional, Union
import pandas as pd
from tqdm import tqdm

from qspoc.converter import csv2dataframe, df2xyz, xlsx2dataframe

from qspoc.multiwfn.interface import MultiwfnDescriptorGenerator
from qspoc.qm.interface import QMComputer
from qspoc.utils import canonicalize_SMILES, generate_prop_lists, legalizing_name

from qspoc.rdkit_processor import get_RDKit_3D_struct


class QSPOCDesc:

    def __init__(
        self,
        multiwfn_path,
        save_dir: Optional[str] = None,
        exe_path_dict: dict = {},
        drop_error: bool = True,
        acc: str = "normal",
        version: str = "stem",
    ):
        # version is stem or all
        self.mol_desc_instances = {}
        self.results = {}
        self.failed_molecules = {}
        self.computation_times = {}
        self.computed = False
        self.multiwfn_path = multiwfn_path
        self.version = version
        self.calc_data_df = pd.DataFrame()
        self.save_dir = Path.cwd() if save_dir is None else Path(save_dir)
        self.exe_path_dict = exe_path_dict
        self.drop_error = drop_error
        
        self.SUPPORT_METHOD = ["xtb", "g16", "orca"]
        if not self.save_dir.exists():
            try:
                self.save_dir.mkdir(parents=True)
            except Exception as e:
                raise ValueError(f"MultiMolDesc Init Fail: {e}")
        assert all(k in self.SUPPORT_METHOD for k in exe_path_dict), f"Unsupported QM method: {list(exe_path_dict.keys())}."

        self.qm_computer = QMComputer()
        if "xtb" in exe_path_dict.keys():
            self._check_command(exe_path_dict["xtb"], "xtb")
            self.qm_computer.load_xtb_computer(xtb_path=exe_path_dict["xtb"])
        if "g16" in exe_path_dict.keys():
            self._check_command(exe_path_dict["g16"], "g16")
            self.qm_computer.load_g16_computer(g16_path=exe_path_dict["g16"])
        if "orca" in exe_path_dict.keys():
            self._check_command(exe_path_dict["orca"], "orca")
            self.qm_computer.load_orca_computer(orca_path=exe_path_dict["orca"])

        self.mdg = MultiwfnDescriptorGenerator(
            multiwfn_path=self.multiwfn_path,
            temp_file_dir=Path(self.save_dir) / "multiwfn_desc",
            temp_coords_dir=Path(self.save_dir) / "coords",
            kill=False,
            file_mode=1,
            acc=acc,
            version=version,
        )

    def load_data(
        self,
        smiles_list: List[str],
        name_list: Optional[List[str]] = None,
        atom_lists: Optional[List[List[int]]] = None,
        bond_lists: Optional[List[List[int]]] = None,
        angle_lists: Optional[List[List[tuple]]] = None,
        dihedral_lists: Optional[List[List[tuple]]] = None,
        axis_lists: Optional[List[tuple]] = None,
        prop_lists: Optional[List[tuple]] = None,
        drop_duplicates: Optional[bool] = True,
    ):

        if not smiles_list:
            raise ValueError("MultiMolDesc Init Fail: smiles_list cannot be empty.")

        self.n_molecules = len(smiles_list)

        if not isinstance(smiles_list, list):
            raise TypeError("MultiMolDesc Init Fail: smiles_list must be a list.")

        if name_list is not None:
            if len(name_list) != len(set(name_list)):
                raise ValueError("Tags must be unique.")
            # convert name_list into valid name
            name_list = [legalizing_name(n) for n in name_list]

        if prop_lists is None:
            prop_lists = generate_prop_lists(smiles_list=smiles_list).tolist()

        # eval name_list length
        for l in [name_list, atom_lists, bond_lists, angle_lists, dihedral_lists, axis_lists, prop_lists]:
            if l is not None:
                if not isinstance(l, list):
                    raise TypeError("MultiMolDesc Init Fail: input must be a list or None.")

                if len(l) != self.n_molecules:
                    raise ValueError("MultiMolDesc Init Fail: The length of input list must be equal to the length of smiles_list.")

        self.calc_data_df["smiles_list"] = smiles_list
        self.calc_data_df["name_list"] = name_list if name_list is not None else [f"mol-{i}" for i in range(self.n_molecules)]
        self.calc_data_df["atom_lists"] = self._get_double_lists(atom_lists, "atom lists")
        self.calc_data_df["bond_lists"] = self._get_triple_lists(bond_lists, "bond lists")
        self.calc_data_df["angle_lists"] = self._get_triple_lists(angle_lists, "angle lists", assert_length=3)
        self.calc_data_df["dihedral_lists"] = self._get_triple_lists(dihedral_lists, "dihedral lists", assert_length=4)
        self.calc_data_df["axis_lists"] = self._get_triple_lists(axis_lists, "axis lists")
        self.calc_data_df["prop_lists"] = self._get_double_lists(prop_lists, "prop lists")

        self._drop_duplicates() if drop_duplicates else None

    def load_data_from_file(
        self,
        file_path: str,
        smiles_tag: str = "smiles_list",
        name_tag: Optional[str] = "name_list",
        atomlists_tag: Optional[str] = "atom_lists",
        bondlists_tag: Optional[str] = "bond_lists",
        anglelists_tag: Optional[str] = "angle_lists",
        dihedrallists_tag: Optional[str] = "dihedral_lists",
        axlists_tag: Optional[str] = "axis_lists",
        prop_tag: Optional[str] = "prop_lists",
        drop_duplicates: Optional[bool] = True,
    ):
        if Path(file_path).suffix == ".csv":
            input_df = csv2dataframe(file_path)
        elif Path(file_path).suffix == ".xlsx":
            input_df = xlsx2dataframe(file_path)
        else:
            raise ValueError("Unsupported file format. Please use .csv or .xlsx.")

        if input_df.empty:
            raise ValueError("MultiMolDesc Load Fail: The input file contains no data.")

        try:
            assert smiles_tag in input_df.columns, f"SMILES tag '{smiles_tag}' not found in the file."
            input_df.rename(columns={smiles_tag: "smiles_list"}, inplace=True)
            self.n_molecules = len(input_df)
        except AssertionError as e:
            raise Exception(f"MultiMolDesc Load Fail: {e}")

        if name_tag in input_df.columns:
            assert input_df[name_tag].is_unique, f"Name in column '{name_tag}' must be unique."
            input_df = input_df.rename(columns={name_tag: "name_list"})
        else:
            input_df["name_list"] = [f"mol-{i}" for i in range(len(input_df))]

        # convert name_list into valid name
        input_df["name_list"] = input_df["name_list"].apply(legalizing_name)

        rename_map = {
            atomlists_tag: "atom_lists",
            bondlists_tag: "bond_lists",
            anglelists_tag: "angle_lists",
            dihedrallists_tag: "dihedral_lists",
            axlists_tag: "axis_lists",
            prop_tag: "prop_lists",
        }
        input_df = input_df.rename(columns={k: v for k, v in rename_map.items() if k in input_df.columns})

        if "prop_lists" not in input_df.columns:
            input_df["prop_lists"] = generate_prop_lists(input_df["smiles_list"])

        column_checker = {
            "atom_lists": lambda df: self._get_double_lists_with_df(df, "atom_lists"),
            "bond_lists": lambda df: self._get_triple_lists_with_df(df, "bond_lists"),
            "angle_lists": lambda df: self._get_triple_lists_with_df(df, "angle_lists", assert_length=3),
            "dihedral_lists": lambda df: self._get_triple_lists_with_df(df, "dihedral_lists", assert_length=4),
            "axis_lists": lambda df: self._get_triple_lists_with_df(df, "axis_lists"),
            "prop_lists": lambda df: self._get_double_lists_with_df(df, "prop_lists"),
        }

        for col, fn in column_checker.items():
            input_df[col] = fn(input_df)

        self.calc_data_df = input_df

        self._drop_duplicates() if drop_duplicates else None

    def get_init_sturct(self, method: str = "default", recalc: bool = False):
        sanitize = method != "loose"
        indices_to_drop = []

        pbar = tqdm(
            enumerate(zip(self.calc_data_df["smiles_list"], self.calc_data_df["name_list"])),
            total=len(self.calc_data_df),
            desc=f"[RDKIT] Init-Struct",
            ncols=100,
            unit="mol",
            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]",
        )

        for idx, (smiles, tag) in pbar:
            try:
                # Check if output already exists and recalc is False
                save_coord_file = self.save_dir / f"coords/{tag}.xyz"

                if save_coord_file.exists() and not recalc:
                    self.calc_data_df.at[idx, "init_coord_file"] = save_coord_file
                    pbar.set_postfix_str(f"✓ {tag} (cached)")
                else:
                    coord, _ = get_RDKit_3D_struct(smiles, sanitize=sanitize)
                    self._save_coord(coord, tag)
                    self.calc_data_df.at[idx, "init_coord_file"] = save_coord_file
            except Exception as e:
                if self.drop_error:
                    print(f"Error processing molecule {tag} ({smiles}). Dropping molecule.")
                    indices_to_drop.append(idx)
                else:
                    raise e

        if self.drop_error and indices_to_drop:
            self.calc_data_df.drop(indices_to_drop, inplace=True)
            self.calc_data_df.reset_index(drop=True, inplace=True)

    def geometric_opt(self, method: str = "xtb", calc_command: str = None, recalc=False, **settings):
        assert method in self.SUPPORT_METHOD, f"Unsupported method: {method}. Supported methods: {self.SUPPORT_METHOD}"
        computer = self.qm_computer.get_computer(method=method)
        indices_to_drop = []

        self._check_df()

        pbar = tqdm(
            enumerate(zip(self.calc_data_df["name_list"], self.calc_data_df["prop_lists"])),
            total=len(self.calc_data_df),
            desc=f"[{method.upper()}] Geometric Opt",
            ncols=100,
            unit="mol",
            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]",
        )

        for idx, (tag, p) in pbar:
            try:
                if "method" in ["g16", "orca"] and (self.save_dir / f"coords/{tag}-opt_with_xtb.xyz").exists():
                    coord_file = self.save_dir / f"coords/{tag}-opt_with_xtb.xyz"
                    expected_output = self.save_dir / "coords" / f"{tag}-opt_with_{method}.xyz"
                else:
                    coord_file = self.save_dir / "coords" / f"{tag}.xyz"
                    expected_output = self.save_dir / "coords" / f"{tag}-opt_with_{method}.xyz"

                # Check if output already exists and recalc is False
                if expected_output.exists() and not recalc:
                    self.calc_data_df.at[idx, "opt_coord_file"] = expected_output
                    pbar.set_postfix_str(f"✓ {tag} (cached)")
                else:
                    optimized_coord_file = computer.struct_opt(
                        coord_file=coord_file, calc_command=calc_command, charge=p[0], multiplicity=p[1], **settings
                    )
                    self.calc_data_df.at[idx, "opt_coord_file"] = optimized_coord_file
                    pbar.set_postfix_str(f"✓ {tag}")
            except Exception as e:
                if self.drop_error:
                    smiles = self.calc_data_df.at[idx, "smiles_list"]
                    print(f"Error during geometric optimization for molecule {tag} ({smiles}). Dropping molecule.")
                    indices_to_drop.append(idx)
                    pbar.set_postfix_str(f"✗ {tag} (error)")
                else:
                    pbar.set_postfix_str(f"✗ {tag} (error)")
                    raise e

        pbar.close()

        if self.drop_error and indices_to_drop:
            self.calc_data_df.drop(indices_to_drop, inplace=True)
            self.calc_data_df.reset_index(drop=True, inplace=True)

    def _save_coord(self, coord, tag):
        coord_dir = self.save_dir / "coords"
        if coord_dir.exists() is False:
            coord_dir.mkdir(exist_ok=True)

        df2xyz(coord, filepath=(coord_dir / f"{tag}.xyz"), tag=tag)

    def singlepoint_energy(self, method: str = "xtb", calc_command: str = None, recalc: bool = False, **settings):
        assert method in self.SUPPORT_METHOD, f"Unsupported method: {method}. Supported methods: {self.SUPPORT_METHOD}"
        computer = self.qm_computer.get_computer(method=method)
        indices_to_drop = []

        self._check_df()

        pbar = tqdm(
            enumerate(zip(self.calc_data_df["name_list"], self.calc_data_df["prop_lists"])),
            total=len(self.calc_data_df),
            desc=f"[{method.upper()}] Single-Point ",
            ncols=100,
            unit="mol",
            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]",
        )

        for idx, (tag, p) in pbar:
            try:
                coord_file = self.calc_data_df.at[idx, "opt_coord_file"]

                # Determine expected output file based on method
                if method == "g16":
                    expected_output = self.save_dir / f"coords/{tag}-sp_with_{method}.fchk"
                else:  # xtb
                    expected_output = self.save_dir / f"coords/{tag}-sp_with_{method}.molden"

                # Check if output already exists and recalc is False
                if expected_output.exists() and not recalc:
                    self.calc_data_df.at[idx, "sp_file"] = expected_output
                    pbar.set_postfix_str(f"✓ {tag} (cached)")
                else:
                    energy_info = computer.singlepoint_energy(
                        tag=tag, coord_file=coord_file, calc_command=calc_command, charge=p[0], multiplicity=p[1], **settings
                    )
                    self.calc_data_df.at[idx, "sp_file"] = energy_info
                    pbar.set_postfix_str(f"✓ {tag}")
            except Exception as e:
                if self.drop_error:
                    smiles = self.calc_data_df.at[idx, "smiles_list"]
                    print(
                        f" Error during single point energy calculation for molecule {tag} ({smiles}). Dropping molecule.",
                    )
                    indices_to_drop.append(idx)
                    pbar.set_postfix_str(f"✗ {tag} (error)")
                else:
                    pbar.set_postfix_str(f"✗ {tag} (error)")
                    raise e

        pbar.close()

        if self.drop_error and indices_to_drop:
            self.calc_data_df.drop(indices_to_drop, inplace=True)
            self.calc_data_df.reset_index(drop=True, inplace=True)

    def descriptor_calc(
        self,
        save: bool = True,
        include: Optional[Union[List[str], str]] = "all",
        cdft_command: str = "xtb",
        recalc: bool = False,
        sh_filepath: str = None,
    ) -> pd.DataFrame:
        """
        计算 DFT 描述符。

        Args:
            save (bool): 是否将结果保存为 CSV 文件。
            include (List[str] or str): 指定要计算的描述符类别列表。
                可以传入 "all" 计算所有类别描述符。
                可选值包括: 'cdft', 'esp', 'shape', 'orbit', 'internal', 'atom', 'bond', 'steric'
            version (str): 描述符版本，可选 'all' 或 'stem'。
                'all': 计算所有描述符
                'stem': 只计算精简版描述符，包括：
                    - cdft: VIP, VEA, Cond_Fukui+, Cond_Fukui-, Cond_dual_desc
                    - esp: min_ESP, max_ESP, MPI, char_balance, area_pSurface, area_nSurface, aver_ALIE, aver_LEA, max_ESP_atom, min_ESP_atom, icp_atom
                    - shape: Molecular_volume, Molecular_sphericity, Molecular_all_surface_area
                    - orbit: HOMO_LUMO_energy_gap, HOMO_ODI, LUMO_ODI
                    - internal: 全部保留
                    - atom: charge_CM5
                    - bond: Wiberg_bond_order, Bond_dipole
                    - steric: 全部保留
        """
        self.mdg.set_file_mode(0 if recalc else 1)
        self.mdg.load_calc_data_df(self.calc_data_df)

        if sh_filepath is not None:
            self.mdg.get_multiwfn_env(sh_filepath)

        cdft_command = self.exe_path_dict["xtb"] if cdft_command == "xtb" else cdft_command

        descriptor_map = {
            "cdft": lambda: self.mdg.get_CDFT_descriptors(cdft_command=cdft_command),
            "esp": self.mdg.get_ESP_descriptors,
            "shape": self.mdg.get_shape_descriptors,
            "orbit": self.mdg.get_orbit_descriptors,
            "internal": self.mdg.get_internal_coordinate,
            "atom": self.mdg.get_atom_descriptors,
            "bond": self.mdg.get_bond_descriptors,
            "steric": self.mdg.get_steric_parameter,
        }

        if any(x == [] or x == [[]] for x in self.calc_data_df["atom_lists"]):
            descriptor_map.pop("atom")
        if any(x == [] or x == [[]] for x in self.calc_data_df["bond_lists"]):
            descriptor_map.pop("bond")
        if any(x == [] or x == [[]] for x in self.calc_data_df["axis_lists"]):
            descriptor_map.pop("steric")

        self._check_df()

        target_keys = []
        if include == "all":
            target_keys = list(descriptor_map.keys())
        elif isinstance(include, list):
            for key in include:
                if key in descriptor_map:
                    target_keys.append(key)
                else:
                    print(f"Warning: Descriptor '{key}' is not valid and will be skipped.")
                    print(f"Available options: {list(descriptor_map.keys())}")
        else:
            raise ValueError("Argument 'include' must be a list of strings or 'all'.")

        if not target_keys:
            print("No valid descriptors selected to calculate.")
            return pd.DataFrame()

        print(f"Calculating descriptors: {target_keys} (version: {self.version})")
        all_desc_dfs = []

        for key in target_keys:
            df = descriptor_map[key]()
            if df is not None and not df.empty:
                all_desc_dfs.append(df)

        if all_desc_dfs:
            final_df = pd.concat(all_desc_dfs, axis=1)
        else:
            final_df = pd.DataFrame()

        self.mdg._cleanup_temp_files()

        if save and not final_df.empty:
            save_path = self.save_dir / "multiwfn_descriptors.csv"
            final_df.index = self.calc_data_df["name_list"]
            final_df.index.name = "name"
            final_df.to_csv(save_path)
            print(f"Descriptors saved to {save_path}")

        return final_df

    def _drop_duplicates(self, keep_tag: str = "first"):
        if keep_tag == False:
            print("Warning: all duplicates will be dropped. This may lead to loss of data.")

        if self.calc_data_df["smiles_list"].duplicated().sum() > 0:
            print(
                f"{self.calc_data_df['smiles_list'].duplicated().sum()} duplicate SMILES found. Duplicates will be dropped.",
            )

        if self.calc_data_df["name_list"].duplicated().sum() > 0:
            raise Exception("🚨 MultiMolDesc Init Fail: Duplicate tags found. Tags must be unique.")

        self.calc_data_df["smiles_list"] = self.calc_data_df["smiles_list"].apply(canonicalize_SMILES)
        self.calc_data_df.drop_duplicates(subset=["smiles_list"], keep=keep_tag, inplace=True)

    def _get_double_lists_with_df(self, df, attr_name):
        if attr_name in df.columns:
            return self._get_double_lists(df[attr_name].tolist(), attr_name)
        else:
            return [[]] * self.n_molecules

    def _get_double_lists(self, a_lists, attr_name):

        if a_lists is not None:
            if all(type(a_list) == list for a_list in a_lists) or all(type(a_list) == tuple for a_list in a_lists):
                pass
            elif all(type(a_list) == int for a_list in a_lists):
                if attr_name == "prop lists":
                    raise Exception("🚨 MultiMolDesc Init Fail: prop lists must be a list of tuple.")
                a_lists = [[a_list] for a_list in a_lists]
            elif all(type(a_list) == str for a_list in a_lists):
                try:
                    a_lists = [literal_eval(a_list) for a_list in a_lists]
                except Exception as e:
                    raise Exception(f"MultiMolDesc Compute Fail: Failed to convert {attr_name} to list of lists: {e}")
            else:
                print([type(a_list) for a_list in a_lists])
                raise Exception(f"MultiMolDesc Init Fail: {attr_name} must be a list of list or None.")

        else:
            return [[]] * self.n_molecules

        return a_lists

    def _get_triple_lists_with_df(self, df, attr_name, assert_length: int = 2):
        if "attr_name" in df.columns:
            return self._get_triple_lists(df[attr_name].tolist(), attr_name, assert_length)
        else:
            return [[[]]] * self.n_molecules

    def _get_triple_lists(self, b_lists, attr_name, assert_length: int = 2):
        if b_lists is not None:

            if isinstance(b_lists[0], str):
                try:
                    b_lists = [literal_eval(l) for l in b_lists]
                except Exception as e:
                    print(f"MultiMolDesc Compute Fail: Failed to convert {attr_name} to list of lists: {e}")
                    raise e
        else:
            return [[[]]] * self.n_molecules

        def is_pair_seq(x):
            return isinstance(x, (list, tuple)) and len(x) == assert_length and (all(isinstance(i, int) for i in x))

        def is_2d_pairs(x):
            return isinstance(x, list) and len(x) > 0 and all(is_pair_seq(i) for i in x)

        def is_3d_2d_pairs(x):
            return isinstance(x, list) and len(x) > 0 and all(is_2d_pairs(i) for i in x)

        if is_2d_pairs(b_lists):
            return [[list(p) if isinstance(p, tuple) else p] for p in b_lists]

        if is_3d_2d_pairs(b_lists):
            return b_lists
        else:
            raise ValueError(f"MultiMolDesc Init Fail: {attr_name} must be a list of list of list or None.")

    def _check_df(self):
        assert len(self.calc_data_df) > 0, "There are no mols in calc_data_df."

    def _check_command(self, command: str, name: str):
        from shutil import which

        if which(command) is None:
            raise ValueError(f"Command '{command}' for {name} not found. Please check the path.")
