import numpy as np
from rdkit import Chem
from rdkit.Chem import AllChem

from .converter import coord2df
from .utils import smiles2mol


def get_RDKit_3D_struct(SMILES: str, sanitize=True, num_confs=10):
    try:
        assert "." not in SMILES
    except AssertionError:
        raise ValueError("Now we can only process single molecule SMILES strings. Please separate them first.")
    mol = smiles2mol(SMILES, sanitize=sanitize)
    mol_h = Chem.AddHs(mol)
    conf_ids = AllChem.EmbedMultipleConfs(mol_h, numConfs=num_confs, randomSeed=42)
    results = AllChem.MMFFOptimizeMoleculeConfs(mol_h, maxIters=200, numThreads=0)
    energies = [(res[1], cid) for cid, res in enumerate(results)]

    min_conf_id = min(energies, key=lambda x: x[0])[1]
    conf = mol_h.GetConformer(min_conf_id)
    coords_list = [(atom.GetSymbol(), *np.array(conf.GetAtomPosition(atom.GetIdx()))) for atom in mol_h.GetAtoms()]

    coord = coord2df(coords_list)
    coord["element"] = coord["element"].replace(to_replace="Al", value="La")
    return coord, mol_h
