#!/usr/bin/env python

from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import RWMol
#import random

# Filter for mols with more than 2 adjacent heteroatoms
def hetero2(mol):
    cond = False
    for atom in mol.GetAtoms():
        if atom.GetAtomicNum() != 6:
            count = 0
            for x in atom.GetNeighbors():
                if x.GetAtomicNum() != 6:
                    count += 1
            if count > 1:
                cond = True
    return cond 

# Filter for mols with more than 1 adjacent heteroatoms
def hetero1(mol):
    cond = False
    for atom in mol.GetAtoms():
        if atom.GetAtomicNum() != 6:
            for x in atom.GetNeighbors():
                if x.GetAtomicNum() != 6:
                    cond = True
    return cond


# Filter for mols with more than maxh number of heteroatoms
def heteromax(mol,maxh):
    cond = False
    count = 0
    for atom in mol.GetAtoms():
        if atom.GetAtomicNum() != 6:
            count = count + 1
#            for x in atom.GetNeighbors():
#                if x.GetAtomicNum() != 6:
    if count > maxh:
            cond = True
    return cond

# Filter for mols with cummulated double bonds or any triple bonds in rings
def cummul_ring(mol):
    cond = False
    for atom in mol.GetAtoms():
        count = 0
        if atom.IsInRing():
            for bond in atom.GetBonds():
                if bond.GetBondType() is Chem.BondType.DOUBLE:
                    count = count + 1
                if bond.GetBondType() is Chem.BondType.TRIPLE:
                    count = count + 2
        if count > 1:
            cond = True
    return cond


# Filter for identical mols
def duplicate(mol1,mol2):
    cond = False
    if Chem.MolToSmiles(mol1) == Chem.MolToSmiles(mol2):
        cond = True
    return cond

