import ezff
from ezff.interfaces import gulp, qchem

# Define ground truths
gt_gs = qchem.read_structure('ground_truths/optCHOSx.out')
gt_gs_energy = qchem.read_energy('ground_truths/optCHOSx.out')
gt_scan = qchem.read_structure('ground_truths/scanCHOSx.out')
gt_scan_energy = qchem.read_energy('ground_truths/scanCHOSx.out')

def my_error_function(rr):
    # Get a unique path for GULP jobs from the MPI rank. Set to '0' for serial jobs
    try:
        path = str(pool.rank)
    except:
        path = '0'

    # Calculate Ground State
    md_gs_job = gulp.job(path = path)
    md_gs_job.structure = gt_gs
    md_gs_job.forcefield = ezff.generate_forcefield(template, rr, FFtype = 'reaxff')
    md_gs_job.options['pbc'] = False
    md_gs_job.options['relax_atoms'] = False
    md_gs_job.options['relax_cell'] = False
    # Run GULP calculation
    md_gs_job.run(command='gulp')
    # Read output from completed GULP job and clean-up
    md_gs_energy = md_gs_job.read_energy()
    md_gs_job.cleanup()


    # Calculate PES Scan
    md_scan_job = gulp.job(path = path)
    md_scan_job.structure = gt_scan
    md_scan_job.forcefield = ezff.generate_forcefield(template, rr, FFtype = 'reaxff')
    md_scan_job.options['pbc'] = False
    md_scan_job.options['relax_atoms'] = False
    md_scan_job.options['relax_cell'] = False
    # Run GULP calculation
    md_scan_job.run(command='gulp')
    # Read output from completed GULP job and clean-up
    md_scan_energy = md_scan_job.read_energy()
    md_scan_job.cleanup()

    # Calculate error
    total_error = ezff.error_energy( md_scan_energy-md_gs_energy, gt_scan_energy-gt_gs_energy, weights = 'uniform')
    return [total_error]

# Read template and variable ranges
bounds = ezff.read_variable_bounds('variable_bounds', verbose=False)
template = ezff.read_forcefield_template('template')

problem = ezff.OptProblem(num_errors = 1, variable_bounds = bounds, error_function = my_error_function, template = template)
algorithm = ezff.Algorithm(problem, 'NSGAII', population = 16)
ezff.optimize(problem, algorithm, iterations = 5)

[((84, 134), 'ezff.interfaces.qchem.read_structure', 'qchem.read_structure', (['"""ground_truths/optCHOSx.out"""'], {}), "('ground_truths/optCHOSx.out')\n", (104, 134), False, 'from ezff.interfaces import gulp, qchem\n'), ((150, 197), 'ezff.interfaces.qchem.read_energy', 'qchem.read_energy', (['"""ground_truths/optCHOSx.out"""'], {}), "('ground_truths/optCHOSx.out')\n", (167, 197), False, 'from ezff.interfaces import gulp, qchem\n'), ((208, 259), 'ezff.interfaces.qchem.read_structure', 'qchem.read_structure', (['"""ground_truths/scanCHOSx.out"""'], {}), "('ground_truths/scanCHOSx.out')\n", (228, 259), False, 'from ezff.interfaces import gulp, qchem\n'), ((277, 325), 'ezff.interfaces.qchem.read_energy', 'qchem.read_energy', (['"""ground_truths/scanCHOSx.out"""'], {}), "('ground_truths/scanCHOSx.out')\n", (294, 325), False, 'from ezff.interfaces import gulp, qchem\n'), ((1725, 1784), 'ezff.read_variable_bounds', 'ezff.read_variable_bounds', (['"""variable_bounds"""'], {'verbose': '(False)'}), "('variable_bounds', verbose=False)\n", (1750, 1784), False, 'import ezff\n'), ((1796, 1837), 'ezff.read_forcefield_template', 'ezff.read_forcefield_template', (['"""template"""'], {}), "('template')\n", (1825, 1837), False, 'import ezff\n'), ((1849, 1960), 'ezff.OptProblem', 'ezff.OptProblem', ([], {'num_errors': '(1)', 'variable_bounds': 'bounds', 'error_function': 'my_error_function', 'template': 'template'}), '(num_errors=1, variable_bounds=bounds, error_function=\n    my_error_function, template=template)\n', (1864, 1960), False, 'import ezff\n'), ((1976, 2024), 'ezff.Algorithm', 'ezff.Algorithm', (['problem', '"""NSGAII"""'], {'population': '(16)'}), "(problem, 'NSGAII', population=16)\n", (1990, 2024), False, 'import ezff\n'), ((2027, 2074), 'ezff.optimize', 'ezff.optimize', (['problem', 'algorithm'], {'iterations': '(5)'}), '(problem, algorithm, iterations=5)\n', (2040, 2074), False, 'import ezff\n'), ((554, 573), 'ezff.interfaces.gulp.job', 'gulp.job', ([], {'path': 'path'}), '(path=path)\n', (562, 573), False, 'from ezff.interfaces import gulp, qchem\n'), ((635, 690), 'ezff.generate_forcefield', 'ezff.generate_forcefield', (['template', 'rr'], {'FFtype': '"""reaxff"""'}), "(template, rr, FFtype='reaxff')\n", (659, 690), False, 'import ezff\n'), ((1047, 1066), 'ezff.interfaces.gulp.job', 'gulp.job', ([], {'path': 'path'}), '(path=path)\n', (1055, 1066), False, 'from ezff.interfaces import gulp, qchem\n'), ((1134, 1189), 'ezff.generate_forcefield', 'ezff.generate_forcefield', (['template', 'rr'], {'FFtype': '"""reaxff"""'}), "(template, rr, FFtype='reaxff')\n", (1158, 1189), False, 'import ezff\n'), ((1556, 1658), 'ezff.error_energy', 'ezff.error_energy', (['(md_scan_energy - md_gs_energy)', '(gt_scan_energy - gt_gs_energy)'], {'weights': '"""uniform"""'}), "(md_scan_energy - md_gs_energy, gt_scan_energy -\n    gt_gs_energy, weights='uniform')\n", (1573, 1658), False, 'import ezff\n')]