import unittest
import os.path
import glob
from shutil import copyfile
import numpy as np

import mwrun

class test_charge_neutrality(unittest.TestCase):

  @classmethod
  def setUpClass(self):
    self.test_path = "charge_neutrality"
    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)

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

  def test_charge_neutrality(self):
    self.run_conf("Psi_average_0.5V", 12)
    self.run_conf("Psi_average_0V", 12)
    n = mwrun.mwrun(self.mw_exec, self.test_path)
    ok, msg = n.compare_datafiles("Psi_average_0.5V/charges.out", "Psi_average_0V/charges.out", rtol=1.0E-7, atol=1.0e-2)
    self.assertTrue(ok, msg)

