import time
import warnings
warnings.filterwarnings("ignore", message="cuaev not installed")
import os
import sys
import torchani
import queue
from ase.optimize import BFGS
from ase.io import read, write
from ase import Atoms
from multiprocessing import Process, Queue, current_process
import contextlib
from aegon.libutils import readxyzs, writexyzs
#-------------------------------------------------------------------------------
eVtokcalpermol = 23.060548012069496
#-------------------------------------------------------------------------------
@contextlib.contextmanager
def suppress_stdout():
    with open(os.devnull, "w") as devnull:
        old_stdout = sys.stdout
        sys.stdout = devnull
        try:
            yield
        finally:
            sys.stdout = old_stdout
#-------------------------------------------------------------------------------
def ani_single_to_file(singlemolase, outname, opt='ANI1ccx', preclist=[1E-03, 1E-04, 1E-05]):
    for prec in preclist:
        with suppress_stdout():
            calculator = {
                'ANI1x': torchani.models.ANI1x().ase(),
                'ANI1ccx': torchani.models.ANI1ccx().ase(),
                'ANI2x': torchani.models.ANI2x().ase()
            }[opt]
        singlemolase.calc = calculator
        dyn = BFGS(singlemolase, logfile=None)
        dyn.run(fmax=prec, steps=200)
    energy = singlemolase.get_potential_energy()
    singlemolase.info['e'] = energy * eVtokcalpermol
    writexyzs(singlemolase, outname)
#-------------------------------------------------------------------------------
def do_job(tasks_to_accomplish, tasks_that_are_done, opt='ANI1ccx', preclist=[1E-03, 1E-04, 1E-05]):
    while True:
        try:
            task = tasks_to_accomplish.get_nowait()
        except queue.Empty:
            break
        else:
            mol = task
            timein=time.strftime("%c")
            print('%s at %s' %(mol.info['i'], timein))
            outname = f"{mol.info['i']}_opt.xyz"
            ani_single_to_file(mol, outname, opt, preclist)
            tasks_that_are_done.put(outname)
#-------------------------------------------------------------------------------
def ANI_parallel(moleculelist, nproc=1, opt='ANI1ccx', preclist=[1E-03, 1E-04, 1E-05]):
    tasks_to_accomplish = Queue()
    tasks_that_are_done = Queue()
    processes = []
    for mol in moleculelist:
        tasks_to_accomplish.put(mol)
    for _ in range(nproc):
        p = Process(target=do_job, args=(tasks_to_accomplish, tasks_that_are_done, opt, preclist))
        processes.append(p)
        p.start()
    for p in processes:
        p.join()
    optimized_molecules = []
    while not tasks_that_are_done.empty():
        xyzfile = tasks_that_are_done.get()
        mol = readxyzs(xyzfile)[0]
        mol.info['c']=1
        optimized_molecules.append(mol)
        os.remove(xyzfile)
    return optimized_molecules
#-------------------------------------------------------------------------------
