"""Functions to generate the conformers to be adsorbed and the most stable one.

functions:
remove_C_linked_Hs: Removes hydrogens bonded to a carbon atom from a molecule.
gen_confs: Generate a number of different conformers in random orientations.
run_isolated: directs the execution of functions to achieve the goal
"""
import logging

from formats import adapt_format

logger = logging.getLogger('DockOnSurf')


def remove_C_linked_Hs(mol):
    """Removes hydrogens bonded to a carbon atom from a rdkit mol object.

    @param mol: rdkit mol object of the molecule with hydrogens
    @return: rdkit mol object of the molecule without hydrogens linked to carbon
    atoms.
    The functions removes the hydrogens bonded to C atoms while keeping the
    ones bonded to heteroatoms.
    """
    from rdkit import Chem

    mol = Chem.RWMol(mol)
    rev_atm_idxs = [atom.GetIdx() for atom in reversed(mol.GetAtoms())]

    for atm_idx in rev_atm_idxs:
        atom = mol.GetAtomWithIdx(atm_idx)
        if atom.GetAtomicNum() != 1:
            continue
        for neigh in atom.GetNeighbors():
            if neigh.GetAtomicNum() == 6:
                mol.RemoveAtom(atom.GetIdx())
    return mol


def gen_confs(mol, num_confs, rms_thld):
    """Generate a number of different conformers in random orientations.

    @param mol: rdkit mol object of the molecule to be adsorbed.
    @param num_confs: number of raw conformers to randomly generate.
    @param rms_thld: rms threshold to consider two conformers as different.
    @return: mol: rdkit mol object containing the different conformers.
    """
    from rdkit.Chem import AllChem as Chem
    logger.info('Generating Conformers')
    conf_ids = Chem.EmbedMultipleConfs(mol, numConfs=num_confs, numThreads=0)
    for conf in conf_ids:
        Chem.UFFOptimizeMolecule(mol, confId=conf)  # TODO: Do we want this?

    Chem.AlignMolConformers(mol)
    mol_wo_Hs = remove_C_linked_Hs(mol)
    for conf1 in conf_ids:
        for conf2 in conf_ids[conf1 + 1:]:
            try:
                rms = Chem.GetBestRMS(mol_wo_Hs, mol_wo_Hs, prbId=conf2,
                                      refId=conf1)
            except ValueError:
                continue
            else:
                if rms <= rms_thld:
                    mol.RemoveConformer(conf2)

    for i, conf in enumerate(mol.GetConformers()):
        conf.SetId(i)

    logger.info(f'Generated {len(mol.GetConformers())} different conformers')
    return mol


def run_isolated(inp_vars):
    """Directs the execution of functions to obtain the conformers to adsorb

    @param inp_vars:
    @return:
    """
    logger.info('Carrying out procedures for the isolated molecule')
    rd_mol = adapt_format('rdkit', inp_vars['molec_file'])
    confs = gen_confs(rd_mol, inp_vars['num_conformers'], inp_vars['iso_rms'])