from rdkit.Chem import AllChem as rdkit
import stk

file_1 = 'SBI.mol'
func_1 = ['diol']
file_2 = 'CNben.mol'
func_2 = ['difluorene']

# turn off cache
stk.OPTIONS['cache'] = False
    
name_1 = file_1.replace('.mol', '')
unit_1 = stk.StructUnit2(file_1, func_1)

name_2 = file_2.replace('.mol', '')
unit_2 = stk.StructUnit2(file_2, func_2)


# Should pass
units = 2
# make polymer
NAME = name_1+'_'+name_2+'_AB_poly'
print(f'Creating polymer: {NAME}')
polymer = stk.Polymer([unit_1, unit_2], stk.Linear('AB', [0, 0], n=units, ends='h'))
# write unoptimized structure
polymer.write(NAME+'.mol')
mol_polymer = rdkit.MolFromMolFile(NAME + '.mol')
# clean molecule with ETKDG
embedder = stk.ETKDG(use_cache=False)
embedder.optimize(polymer)
# write optimized polymer to json
polymer.dump(NAME+'_opt.json')
polymer.write(NAME+'_opt.mol')


# Should fail
units = 5
# make polymer
NAME = name_1+'_'+name_2+'_AB_poly'
print(f'Creating polymer: {NAME}')
polymer = stk.Polymer([unit_1, unit_2], stk.Linear('AB', [0, 0], n=units, ends='h'))
# write unoptimized structure
polymer.write(NAME+'.mol')
mol_polymer = rdkit.MolFromMolFile(NAME + '.mol')
# clean molecule with ETKDG
embedder = stk.ETKDG(use_cache=False)  # should fail here
embedder.optimize(polymer)
# write optimized polymer to json
polymer.dump(NAME+'_opt.json')
polymer.write(NAME+'_opt.mol')