from pathlib import Path
from rdkit import Chem
from typing import List, Optional
from ast import literal_eval
import pandas as pd
from tqdm import tqdm
from typing import Optional, List, Union
from indigo import Indigo
from indigo.renderer import IndigoRenderer


def smiles2mol(SMILES: str, sanitize: bool, removeHs: bool = False) -> Chem.Mol:
    params = Chem.SmilesParserParams()
    params.sanitize = sanitize
    params.removeHs = removeHs
    if SMILES == "None" or SMILES == "none":
        return None
    mol = Chem.MolFromSmiles(SMILES, params=params)
    if mol is None:
        print(f"Invalid SMILES '{SMILES}'")
        raise ValueError()
    return mol


def get_mol_charge(SMILES: str) -> int:
    mol = smiles2mol(SMILES, sanitize=True)
    return Chem.GetFormalCharge(mol)


def smiles2png(
    SMILES: str,
    file_path: str,
    highlight_atoms: Optional[Union[List[int], str, int]] = None,
    show_index: Optional[bool] = False,
    addHs: bool = False,
) -> None:
    if SMILES in ["None", "none", "blanck_cell", None]:
        return
    indigo = Indigo()
    renderer = IndigoRenderer(indigo)
    try:
        s = canonicalize_SMILES(SMILES)
        s = s.replace("->", "").replace("<-", "")
        mol = indigo.loadMolecule(s)
    except Exception as e:
        print(f"Invalid SMILES: {SMILES}")
        return

    mol.layout()

    if addHs:
        mol.unfoldHydrogens()

    if isinstance(highlight_atoms, int):
        highlight_atoms = [highlight_atoms]
    elif isinstance(highlight_atoms, list) and len(highlight_atoms) > 0 and isinstance(highlight_atoms[0], list):
        highlight_atoms = [item for sublist in highlight_atoms for item in sublist]
    elif isinstance(highlight_atoms, str):
        try:
            highlight_atoms = literal_eval(highlight_atoms)
        except:
            highlight_atoms = None
    if highlight_atoms is not None:
        if any(not isinstance(x, int) for x in highlight_atoms):
            print("🚨 Highlight atoms must be a list of integers")
            raise ValueError()

        highlight_set = set(highlight_atoms)
        for atom in mol.iterateAtoms():
            if atom.index() in highlight_set:
                atom.highlight()
        for bond in mol.iterateBonds():
            if bond.source().index() in highlight_set and bond.destination().index() in highlight_set:
                bond.highlight()

    indigo.setOption("render-image-width", 800)
    indigo.setOption("render-image-height", 800)

    indigo.setOption("render-background-color", "1.0, 1.0, 1.0")
    indigo.setOption("render-margins", "2, 2")

    indigo.setOption("render-bond-line-width", 2.0)
    indigo.setOption("render-stereo-style", "old")
    indigo.setOption("render-output-format", "png")
    if show_index:
        indigo.setOption("render-atom-ids-visible", True)
    renderer.renderToFile(mol, str(file_path))


def smileslist2png(
    SMILES_list: List[str],
    savepath: str = Path.cwd(),
    name_list: Optional[List[str]] = None,
    highlight_atoms_list: Optional[List[int]] = None,
    show_index: Optional[bool] = False,
    add_hydrogens: bool = False,
):
    if name_list is None:
        name_list = [f"mol-{i}.png" for i in range(len(SMILES_list))]
    if len(name_list) != len(SMILES_list):
        raise ValueError("The length of name_list must be equal to the length of SMILES_list")
    if highlight_atoms_list is None:
        highlight_atoms_list = [None for _ in range(len(SMILES_list))]
    else:
        if len(highlight_atoms_list) != len(SMILES_list):
            raise ValueError("The length of highlight_atoms_list must be equal to the length of SMILES_list")

    save_dir = Path(savepath)
    if not save_dir.exists():
        try:
            save_dir.mkdir(parents=True)
        except Exception as e:
            raise FileExistsError(f"Cannot create directory '{savepath}', because {e}.")

    for smiles, name, hl_atoms in tqdm(zip(SMILES_list, name_list, highlight_atoms_list), total=len(SMILES_list)):
        png_path = save_dir / name

        if not png_path.suffix:
            png_path = png_path.with_suffix(".png")

        try:
            smiles2png(smiles, str(png_path), highlight_atoms=hl_atoms, show_index=show_index, addHs=add_hydrogens)
        except Exception as e:
            if e.args == ():
                print(f"SMILES '{smiles}' cannot be draw, because ↑.")
            else:
                print(f"SMILES '{smiles}' cannot be draw, because {e}.")


def canonicalize_SMILES(s):
    try:
        if s in ["None", "none", "nan", "blanck_cell"] or s == "":
            return s
        else:
            return Chem.MolToSmiles(Chem.MolFromSmiles(s))
    except Exception as e:
        print(f"SMILES '{s}' cannot be canonicalized, because {e}.")
        return s


def canonicalize_input_SMILES_file(csv_path: str, SMILES_tag: str):
    df = pd.read_csv(csv_path)
    if type(SMILES_tag) == str:
        df[SMILES_tag] = df[SMILES_tag].apply(canonicalize_SMILES)
    elif type(SMILES_tag) == list:
        for tag in SMILES_tag:
            df[tag] = df[tag].apply(canonicalize_SMILES)
    else:
        raise ValueError("SMILES_tag must be str or list of str")

    df.to_csv(csv_path, index=False)


def combining_descriptors(
    ref_df: pd.DataFrame,
    to_combine_df_list: List[pd.DataFrame],
    ref_name_list: List[str],
    to_combine_name_list: List[str] = None,
    prefix_list: List[str] = None,
):
    to_combine_name_list = ["SMILES"] * len(to_combine_df_list) if to_combine_name_list is None else to_combine_name_list
    try:
        assert len(to_combine_df_list) == len(ref_name_list) == len(to_combine_name_list)
    except AssertionError:
        raise ValueError("Length mismatch: to_combine_df_list vs name_list")
    # 处理prefix_list（默认生成或检查长度）
    if prefix_list is None:
        prefix_list = [f"part{i}" for i in range(len(to_combine_df_list))]
    else:
        try:
            assert len(prefix_list) == len(to_combine_df_list)
        except AssertionError:
            raise ValueError("Length mismatch: prefix_list vs to_combine_df_list")
    # 检查所有tag都在ref_df的列中
    missing_tags = [tag for tag in ref_name_list if tag not in ref_df.columns]
    if missing_tags:
        raise ValueError(f"Missing tags in ref_df: {', '.join(missing_tags)}")
    # -------------------------- 2. 初始化结果表 & 合并逻辑 --------------------------
    # 复制参考表（避免修改原数据）
    result_df = ref_df.copy()
    result_df = result_df.loc[:, ref_name_list]
    # 遍历每个待合并的描述符表、标签、后缀
    for i in range(len(to_combine_df_list)):
        current_ref_tag = ref_name_list[i]
        current_df = to_combine_df_list[i]
        current_prefix = prefix_list[i]
        current_to_combine_tag = to_combine_name_list[i]
        try:
            # 重命名描述符表的列（添加后缀，避免列名冲突）
            current_df_renamed = current_df.rename(columns=lambda col: f"{current_prefix}_{col}")
            current_df_renamed.set_index(f"{current_prefix}_{current_to_combine_tag}", inplace=True)
        except Exception as e:
            raise ValueError(f"Column renaming failed for tag '{current_ref_tag}'") from e
        # try:
        # 左连接：参考表的tag列 → 描述符表的索引（保留参考表所有行）
        result_df = result_df.merge(
            current_df_renamed,
            left_on=current_ref_tag,  # 参考表中的tag列作为连接键
            right_index=True,  # 描述符表的索引作为连接键
            how="left",  # 左连接（保留参考表所有数据）
        )

    return result_df


def generate_prop_lists(smiles_list: pd.Series) -> pd.Series:
    from rdkit.Chem import rdmolops

    get_charge = lambda smi: rdmolops.GetFormalCharge(smiles2mol(smi, sanitize=True))
    return pd.Series([(get_charge(smi), 1) for smi in smiles_list])


if __name__ == "__main__":

    SMILES_list = [
        "O=S(=O)([O-])C(F)(F)F",
        "O=S(=O)([N-]S(=O)(=O)C(F)(F)F)C(F)(F)F",
        "O=C([O-])c1c(F)c(F)c(F)c(F)c1F",
        "O=C([O-])C(F)(F)F",
        "CC(=O)[O-]",
        "O=C([O-])c1ccccc1",
        "O=C([O-])C(F)(F)C(F)(F)F",
        "CC(F)(F)C(=O)[O-]",
        "hhs",
        "O=C([O-])C(F)(F)C(F)(F)C(F)(F)F",
        "O=C([O-])C(F)(F)C(F)(F)C(F)(F)C(F)(F)F",
        "O=C([O-])C(F)(F)C(F)(F)C(F)(F)C(F)(F)C(F)(F)F",
        "O=C([O-])C(F)(F)C(F)(F)C(F)(F)C(F)(F)C(F)(F)C(F)(F)F",
        "O=C([O-])C(F)(F)C(F)(F)C(F)(F)C(F)(F)C(F)(F)C(F)(F)C(F)(F)F",
    ]

    smileslist2png(SMILES_list, savepath="temp", show_index=True)


def legalizing_name(name: str) -> str:
    if name is None or name == "":
        return "None"
    else:
        name = name.replace(" ", "_")
        name = name.replace("/", "_")
        name = name.replace("\\", "_")
        name = name.replace(",", "_")
        return name
