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

import mwrun

class test_matrix_inversion(unittest.TestCase):

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

  def run_conf(self, confID, nranks, check_dipoles=False, check_electrodes=False, check_matrix=False, matrix_loaded=False):
    path_to_config = os.path.join(self.test_path, confID)
    n = mwrun.mwrun(self.mw_exec, path_to_config)
    self.workdir = n.workdir

    if matrix_loaded and mwrun.glob_compare_hessian:
      copyfile(os.path.join(n.workdir, "hessian_matrix.inpt.ref"), os.path.join(n.workdir, "hessian_matrix.inpt"))

    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_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_dipoles:
        ok, msg = n.compare_datafiles("dipoles.out", "dipoles.ref", rtol=1.0E-6, 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)

  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, "hessian_matrix.inpt")):
        os.remove(os.path.join(self.workdir, "hessian_matrix.inpt"))

#  def test_conf_2D_ishii_nodisp(self):
#    if mwrun.glob_skip_long_test:
#        self.skipTest("Test too long")
#    self.run_conf("dipoles/2D-ishii-nodisp", 1, check_dipoles=True)
#
#  def test_conf_2D_ishii_nodisp_4MPI(self):
#    if mwrun.glob_skip_long_test:
#        self.skipTest("Test too long")
#    self.run_conf("dipoles/2D-ishii-nodisp", 4, check_dipoles=True)

  def test_conf_LiClwAl(self):
     self.run_conf("electrodes/LiCl-Al", 1, check_matrix=True, check_electrodes=True)

  def test_conf_LiClwAl_4MPI(self):
    self.run_conf("electrodes/LiCl-Al", 4, check_matrix=True, check_electrodes=True)

  def test_conf_LiClwAl_invert_matrix_loaded(self):
     self.run_conf("electrodes/LiCl-Al", 1, matrix_loaded=True, check_electrodes=True)

  def test_conf_LiClwAl_charge_neutrality(self):
     self.run_conf("electrodes/LiCl-Al-charge-neutrality", 1, check_matrix=True, check_electrodes=True)

  def test_conf_LiClwAl_charge_neutrality_4MPI(self):
    self.run_conf("electrodes/LiCl-Al-charge-neutrality", 4, check_matrix=True, check_electrodes=True)

  def test_conf_LiClwAl_charge_neutrality_invert_matrix_loaded(self):
     self.run_conf("electrodes/LiCl-Al-charge-neutrality", 1, matrix_loaded=True, check_electrodes=True)

