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

import mwrun

class test_dump_per_species(unittest.TestCase):

  @classmethod
  def setUpClass(self):
    self.test_path = "dump_per_species"
    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)
	
    lammps_ref = os.path.join(self.workdir, "lammpstrj.ref")
    lammps_trj = os.path.join(self.workdir, "trajectories.lammpstrj")
    ok, msg = n.diff_files(lammps_ref, lammps_trj)
    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.lammpstrj"))
    
  def test_conf_dump_per_species_nacl(self):
    self.run_conf("nacl", 1)

  def test_conf_dump_per_species_tip4p_water(self):
    self.run_conf("tip4p_water", 1)

  def test_conf_dump_per_species_graphite(self):
    self.run_conf("graphite", 1)


