// Compare forces acting on selected atom in system, obtained directly 
// from LAMMPS and indirectly, as a change of potential energy. 
//
// that's a serial program.
//
// written by Marcin Wojdyr wojdyr@gmail.com
// put into public domain

#include <iostream>
#include <iomanip>
#include <cstdlib>
#include <ctime>

// lammps header files
#include "lammps.h"         
#include "input.h"
#include "output.h"
#include "thermo.h"
#include "atom.h"

using namespace LAMMPS_NS;
using namespace std;

void print_usage_and_exit()
{
    cerr << "Usage:\n checkpot snippet [atom_nr [h]] >/dev/null\n" 
            "snippet is a LAMMPS script, that should set up system "
            "and potential\n";
    exit(1);
}

double calc_f_num(LAMMPS *lmp, int atom_nr, int coord, double h)
{
    double orig = lmp->atom->x[atom_nr][coord];
    lmp->atom->x[atom_nr][coord] = orig + h;
    lmp->input->one("run 1");
    double u1;
    lmp->output->thermo->evaluate_keyword("pe", &u1);
    lmp->atom->x[atom_nr][coord] = orig - h;
    lmp->input->one("run 1");
    double u2;
    lmp->output->thermo->evaluate_keyword("pe", &u2);
    lmp->atom->x[atom_nr][coord] = orig;
    return (u2 - u1) / (2*h);
}


int main(int argc, char **argv)
{
    if (argc != 2 && argc != 3 && argc != 4) 
        print_usage_and_exit();

    char *endptr;
    int atom_nr;
    if (argc >= 3) {
        atom_nr = strtol(argv[2], &endptr, 10);
        if (*endptr != '\0') 
            print_usage_and_exit();
    }

    double h = 0.01;
    if (argc == 4) {
        h = strtod(argv[3], &endptr);
        if (*endptr != '\0') 
            print_usage_and_exit();
    }

    LAMMPS *lmp = new LAMMPS(0, NULL, 0);
    lmp->input->file(argv[1]);

    int natoms = static_cast<int> (lmp->atom->natoms);

    if (argc == 2) {
        srand(time(0));
        atom_nr = rand() % natoms;
    }

    if (atom_nr < 0 || atom_nr >= natoms) {
        cerr << "Error: wrong atom_nr: " << atom_nr << endl;
        exit(1);
    }
    double *x = lmp->atom->x[atom_nr];
    cerr << "Selected atom #" << atom_nr 
        << ": type " << lmp->atom->type[atom_nr]
        << ", position (" << x[0] << ", " << x[1] << ", " << x[2] << ")" 
        << endl; 

    lmp->input->one("run 0");
    double epot;
    lmp->output->thermo->evaluate_keyword("pe", &epot);
    int p = 10;
    cerr << "Epot of system = " << setprecision(p) << epot 
        << "  |  Epot / N = " << setprecision(p) << epot / natoms << endl;
    double f[3];
    for (int i = 0; i < 3; ++i)
        f[i] = lmp->atom->f[atom_nr][i];
    cerr << "   force-in-program    ?=  force-calculated-numerically" << endl;
    cerr << "fx = " << setprecision(p) << f[0] << " ?= " 
        << setprecision(p) << calc_f_num(lmp, atom_nr, 0, h) << endl;
    cerr << "fy = " << setprecision(p) << f[1] << " ?= " 
        << setprecision(p) << calc_f_num(lmp, atom_nr, 1, h) << endl;
    cerr << "fz = " << setprecision(p) << f[2] << " ?= " 
        << setprecision(p) << calc_f_num(lmp, atom_nr, 2, h) << endl;
    delete lmp;
}