import numpy as np
import mdtraj as md
from unittest import skipIf
import logging
from openmoltools import utils, amber,packmol
from distutils.spawn import find_executable


logging.basicConfig(level=logging.DEBUG, format="LOG: %(message)s")

@skipIf(packmol.PACKMOL_PATH is None, "Skipping testing of packmol conversion because packmol not found.")
def test_amber_box():
    etoh_filename = utils.get_data_filename("chemicals/etoh/etoh.mol2")
    trj_list = [md.load(etoh_filename)]
    
    with utils.enter_temp_directory():  # Prevents creating tons of GAFF files everywhere.
        
        box_filename = "./box.pdb"
        box_trj = packmol.pack_box(trj_list, [50])
        box_trj.save(box_filename)
    
        gaff_mol2_filename1, frcmod_filename1 = utils.run_antechamber("etoh", etoh_filename, charge_method=None)
        
        mol2_filenames = [gaff_mol2_filename1]
        frcmod_filenames =  [frcmod_filename1]
        
        prmtop_filename = "./out.prmtop"
        inpcrd_filename = "./out.inpcrd"

        tleap_cmd = amber.build_mixture_prmtop(mol2_filenames, frcmod_filenames, box_filename, prmtop_filename, inpcrd_filename)
        print(tleap_cmd)


@skipIf(packmol.PACKMOL_PATH is None, "Skipping testing of packmol conversion because packmol not found.")
def test_amber_binary_mixture():
    sustiva_filename = utils.get_data_filename("chemicals/etoh/etoh.mol2")
    etoh_filename = utils.get_data_filename("chemicals/etoh/etoh_renamed.mol2")

    trj0, trj1 = md.load(sustiva_filename), md.load(etoh_filename)
    
    # Hack to assign unique residue names that are consistent with contents of mol2 files
    trj0.top.residue(0).name = "LIG"
    trj1.top.residue(0).name = "LI2"
    
    trj_list = [trj0, trj1]
    
    with utils.enter_temp_directory():  # Prevents creating tons of GAFF files everywhere.
        
        box_filename = "./box.pdb"
        box_trj = packmol.pack_box(trj_list, [25, 50])
        box_trj.save(box_filename)
        
        gaff_mol2_filename0, frcmod_filename0 = utils.run_antechamber("sustiva", sustiva_filename, charge_method=None)
        gaff_mol2_filename1, frcmod_filename1 = utils.run_antechamber("etoh", etoh_filename, charge_method=None)
        
        mol2_filenames = [gaff_mol2_filename0, gaff_mol2_filename1]
        frcmod_filenames = [frcmod_filename0, frcmod_filename1]
        
        
        prmtop_filename = "./out.prmtop"
        inpcrd_filename = "./out.inpcrd"

        tleap_cmd = amber.build_mixture_prmtop(mol2_filenames, frcmod_filenames, box_filename, prmtop_filename, inpcrd_filename)
        print(tleap_cmd)

@@ -2584,28 +2584,1184 @@
 e)%0A        print(tleap_cmd)%0A
+%0A%0Adef test_run_antechamber():%0A    molecule_name = %22sustiva%22%0A    input_filename = utils.get_data_filename(%22chemicals/sustiva/sustiva.mol2%22)%0A    with utils.enter_temp_directory():  # Prevents creating tons of GAFF files everywhere.%0A        gaff_mol2_filename, frcmod_filename = amber.run_antechamber(molecule_name, input_filename, charge_method=None)%0A%0Adef test_run_tleap():%0A    molecule_name = %22sustiva%22%0A    input_filename = utils.get_data_filename(%22chemicals/sustiva/sustiva.mol2%22)%0A    with utils.enter_temp_directory():  # Prevents creating tons of GAFF files everywhere.%0A        gaff_mol2_filename, frcmod_filename = amber.run_antechamber(molecule_name, input_filename, charge_method=None)%0A        prmtop, inpcrd = utils.run_tleap(molecule_name, gaff_mol2_filename, frcmod_filename)%0A%0Adef test_run_antechamber_charges():%0A    molecule_name = %22acetate%22%0A    input_filename = utils.get_data_filename(%22chemicals/acetate/acetate.mol2%22)%0A    with utils.enter_temp_directory():  # Prevents creating tons of GAFF files everywhere.%0A        gaff_mol2_filename, frcmod_filename = amber.run_antechamber(molecule_name, input_filename, charge_method=None, net_charge=-1)%0A%0A

Move amber functionality tests to test_amber