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

import mwrun

class test_non_neutral(unittest.TestCase):

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

  def run_conf(self, confID, nranks, check_electrodes=False, check_matrix=False, check_maze=False):
    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)
    self.setup = True

    ok, msg = n.compare_datafiles("forces.out", "forces.ref", rtol=1.0E-6, atol=0.0)
    self.assertTrue(ok, msg)

    ok, msg = n.compare_datafiles("energies.out", "energies.ref", rtol=1.0E-6, atol=0.0)
    self.assertTrue(ok, msg)

    ok, msg = n.compare_datafiles("energies_breakdown.out", "energies_breakdown.ref", rtol=1.0E-6, atol=0.0)
    self.assertTrue(ok, msg)

    if check_maze and mwrun.glob_compare_hessian:
      ok, msg = n.compare_matrixfiles("maze_matrix.inpt", "maze_matrix.inpt.ref", rtol=1.0E-5, atol=0.0)
      self.assertTrue(ok, msg)

    if check_matrix and mwrun.glob_compare_hessian:
      ok, msg = n.compare_matrixfiles("hessian_matrix.inpt", "hessian_matrix.inpt.ref", rtol=1.0E-5, atol=0.0)
      self.assertTrue(ok, msg)

    if check_electrodes:
        ok, msg = n.compare_datafiles("charges.out", "charges.ref", rtol=1.0E-6, atol=0.0)
        self.assertTrue(ok, msg)
        ok, msg = n.compare_datafiles("total_charges.out", "total_charges.ref", rtol=1.0E-6, atol=0.0)
        self.assertTrue(ok, msg)

  def tearDown(self):
    if self.setup:
      for f in glob.glob(os.path.join(self.workdir, "*.out")):
        os.remove(f)
      if os.path.isfile(os.path.join(self.workdir, "maze_matrix.inpt")):
        os.remove(os.path.join(self.workdir, "maze_matrix.inpt"))
      if os.path.isfile(os.path.join(self.workdir, "hessian_matrix.inpt")):
        os.remove(os.path.join(self.workdir, "hessian_matrix.inpt"))

  def test_conf_LiCl_Al_Na_maze(self):
     self.run_conf("LiCl-Al-Na-maze", 1, check_maze=True, check_electrodes=True)

  def test_conf_LiCl_Al_Na_maze_4MPI(self):
     self.run_conf("LiCl-Al-Na-maze", 4, check_maze=True, check_electrodes=True)

  def test_conf_LiCl_Al_Na_cg(self):
     self.run_conf("LiCl-Al-Na-cg", 1, check_electrodes=True)

  def test_conf_LiCl_Al_Na_cg_4MPI(self):
     self.run_conf("LiCl-Al-Na-cg", 4, check_electrodes=True)

  def test_conf_LiCl_Al_Na_MI(self):
     self.run_conf("LiCl-Al-Na-MI", 1, check_matrix=True, check_electrodes=True)

  def test_conf_LiCl_Al_Na_MI_4MPI(self):
     self.run_conf("LiCl-Al-Na-MI", 4, check_matrix=True, check_electrodes=True)
