import unittest
import os.path
import glob
import numpy as np

import mwrun

class test_aim(unittest.TestCase):

  @classmethod
  def setUpClass(self):
    self.test_path = "aim"
    self.mw_exec = os.path.abspath(mwrun.glob_mw_exec)

  def run_conf(self, confID, nranks):
    path_to_config = os.path.join(self.test_path, confID)
    n = mwrun.mwrun(self.mw_exec, path_to_config)
    self.workdir = n.workdir
    n.run_mw(nranks)

    ok, msg = n.compare_datafiles("forces.out", "forces.ref",1.0e-4,1.0e-6)
    self.assertTrue(ok, msg)

    ok, msg = n.compare_datafiles("dipoles.out", "dipoles.ref")
    self.assertTrue(ok, msg)

    ok, msg = n.compare_datafiles("radius.out", "radius.ref",1.0e-4,1.0e-6)
    self.assertTrue(ok, msg)

    ok, msg = n.compare_datafiles("energies.out", "energies.ref")
    self.assertTrue(ok, msg)

    ok, msg = n.compare_datafiles("energies_breakdown.out", "energies_breakdown.ref")
    self.assertTrue(ok, msg)

    ok, msg = n.compare_datafiles("pressure.out", "pressure.ref")
    self.assertTrue(ok, msg)

    ok, msg = n.compare_datafiles("stress_tensor.out", "stress_tensor.ref")
    self.assertTrue(ok, msg)    

  def tearDown(self):
    for f in glob.glob(os.path.join(self.workdir, "*.out")):
      os.remove(f)
    os.remove(os.path.join(self.workdir, "trajectories.xyz"))
    os.remove(os.path.join(self.workdir, "trajectories.pdb"))
    os.remove(os.path.join(self.workdir, "trajectories.lammpstrj"))

  def test_conf_3D_ishii_nodisp(self):
    self.run_conf("3D-ishii-nodisp", 1)

  def test_conf_3D_ishii_nodisp_8MPI(self):
    self.run_conf("3D-ishii-nodisp", 8)
