#!/usr/bin/python
import os, sys, string
from openeye.oechem import *

#----------------------------------------------
def a3_f2f(inf, outf, metals_file1, metals_file2):
    metal_list = []
    ifs = oemolistream(inf)
    ofs = oemolostream(outf)
    ofs.SetFormat(OEFormat_ISM)
    for mol in ifs.GetOEMols():
        metal = 0
        carbon = 0
        heavy_atom = 0
        for atom in mol.GetAtoms():
            if not atom.IsHydrogen():
                heavy_atom+=1
                if atom.IsMetal():
                   metal = 1
        #    elif not atom.IsHydrogen() and not atom.IsNitrogen() and not atom.IsCarbon and not atom.IsHalogen() and not atom.IsOxygen() and not atom.IsPhosphorus() and not atom.IsSulfur():
                if atom.IsCarbon():
                    carbon = 1
        # DEFINE here the cutoffs for heavy atoms
        if heavy_atom <= 48 and heavy_atom > 4 and carbon == 1 and string.find(OECreateAbsSmiString(mol), '*') == -1:
            OEWriteMolecule(ofs, mol)
            if metal == 1:
                print OECreateAbsSmiString(mol), mol.GetTitle()
                metal_list.append(mol.GetTitle())
        else:
            pass
    metalf1 = open(metals_file1, 'w')
    metalf2 = open(metals_file2, 'w')

    for cpd in metal_list:
        metalf1.write(cpd+'\n')
        metalf2.write(cpd+'\n')
    
    metalf1.close()
    metalf2.close()
    ifs.close()
    ofs.close()
#----------------------------------------------


try:
    inputfile = sys.argv[1]
    outputfile = 'a3.2'+sys.argv[1][string.find(sys.argv[1], '_'):string.find(sys.argv[1], '_out')]+'_out.ism'
    metal_file1 = 'a3.2'+sys.argv[1][string.find(sys.argv[1], '_'):string.find(sys.argv[1], '_out')]+'_metals_keep.txt'
    metal_file2 = 'a3.2'+sys.argv[1][string.find(sys.argv[1], '_'):string.find(sys.argv[1], '_out')]+'_metals.txt'

except IndexError:
    print '\nUsage: a3.2_size_cutoff_filter.py Molfile\n'
    raise SystemExit()

a3_f2f(inputfile, outputfile, metal_file1, metal_file2)