import pytest
import os
import sys
import numpy as np
from flare.otf import OTF
from flare.gp import GaussianProcess
from flare.struc import Structure
import flare.kernels as en


def cleanup_espresso_run(target: str = None):
    os.system('rm pwscf.out')
    os.system('rm pwscf.wfc')
    os.system('rm -r pwscf.save')
    os.system('rm pwscf.in')
    os.system('rm pwscf.wfc1')
    os.system('rm pwscf.wfc2')
    if target:
        os.system('rm ' + target)


# ------------------------------------------------------
#                   test  otf runs
# ------------------------------------------------------
@pytest.mark.skipif(not os.environ.get('PWSCF_COMMAND',
                          False), reason='PWSCF_COMMAND not found '
                                  'in environment: Please install Quantum '
                                  'ESPRESSO and set the PWSCF_COMMAND env. '
                                  'variable to point to pw.x.')
def test_otf_h2():
    """
    Test that an otf run can survive going for more steps
    :return:
    """
    os.system('cp ./test_files/qe_input_1.in ./pwscf.in')

    qe_input = './pwscf.in'
    dt = 0.0001
    number_of_steps = 20
    cutoffs = np.array([5])
    dft_loc = os.environ.get('PWSCF_COMMAND')
    std_tolerance_factor = -0.1

    # make gp model
    kernel = en.two_body
    kernel_grad = en.two_body_grad
    hyps = np.array([1, 1, 1])
    hyp_labels = ['Signal Std', 'Length Scale', 'Noise Std']
    energy_force_kernel = en.two_body_force_en

    gp = \
        GaussianProcess(kernel=kernel,
                        kernel_grad=kernel_grad,
                        hyps=hyps,
                        cutoffs=cutoffs,
                        hyp_labels=hyp_labels,
                        par=True,
                        energy_force_kernel=energy_force_kernel,
                        maxiter=50)

    otf = OTF(qe_input, dt, number_of_steps, gp, dft_loc,
              std_tolerance_factor, init_atoms=[0],
              calculate_energy=True, max_atoms_added=1,
              no_cpus=2, par=True,
              output_name='h2_otf_qe_par')

    otf.run()
    os.system('mkdir test_outputs')
    os.system('mv h2_otf_qe_par* test_outputs')
    cleanup_espresso_run()

@@ -1180,10 +1180,1645 @@
 s = 
-20
+4%0A    cutoffs = np.array(%5B5%5D)%0A    dft_loc = os.environ.get('PWSCF_COMMAND')%0A    std_tolerance_factor = -0.1%0A%0A    # make gp model%0A    kernel = en.two_body%0A    kernel_grad = en.two_body_grad%0A    hyps = np.array(%5B1, 1, 1%5D)%0A    hyp_labels = %5B'Signal Std', 'Length Scale', 'Noise Std'%5D%0A    energy_force_kernel = en.two_body_force_en%0A%0A    gp = %5C%0A        GaussianProcess(kernel=kernel,%0A                        kernel_grad=kernel_grad,%0A                        hyps=hyps,%0A                        cutoffs=cutoffs,%0A                        hyp_labels=hyp_labels,%0A                        par=True,%0A                        energy_force_kernel=energy_force_kernel,%0A                        maxiter=50)%0A%0A    otf = OTF(qe_input, dt, number_of_steps, gp, dft_loc,%0A              std_tolerance_factor, init_atoms=%5B0%5D,%0A              calculate_energy=True, max_atoms_added=1,%0A              no_cpus=2, par=True,%0A              output_name='h2_otf_qe_par')%0A%0A    otf.run()%0A    os.system('mkdir test_outputs')%0A    os.system('mv h2_otf_qe_par* test_outputs')%0A    cleanup_espresso_run()%0A%0A@pytest.mark.skipif(not os.environ.get('PWSCF_COMMAND',%0A                          False), reason='PWSCF_COMMAND not found '%0A                                  'in environment: Please install Quantum '%0A                                  'ESPRESSO and set the PWSCF_COMMAND env. '%0A                                  'variable to point to pw.x.')%0Adef test_otf_Al_npool():%0A    %22%22%22%0A    Test that an otf run can survive going for more steps%0A    :return:%0A    %22%22%22%0A    os.system('cp ./test_files/qe_input_2.in ./pwscf.in')%0A%0A    qe_input = './pwscf.in'%0A    dt = 0.0001%0A    number_of_steps = 4
 %0A   
@@ -3657,32 +3657,32 @@
 _atoms_added=1,%0A
-
               no
@@ -3691,32 +3691,41 @@
 pus=2, par=True,
+ npool=2,
 %0A              o

add npool unit test