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

import mwrun

class python_interface(mwrun.mwrun):

  def __init__(self, mw_exec, path_to_config, path_to_ref):
    super(python_interface, self).__init__(mw_exec, path_to_config)
    # Constants
    self.avogadro = 6.022140857e+23
    self.boltzmann = 1.38064852e-23
    self.K2kJpermol = self.boltzmann * self.avogadro * 1.0e-3
    self.hartree2J = 4.359744650e-18
    self.workdir = os.path.join(path_to_config)

    self.refdir = os.path.join(path_to_ref)
    self.energy_file = os.path.join(self.workdir, "energies_breakdown.out")
    self.xyz_file = os.path.join(self.workdir, "trajectories.xyz")
    self.xyz_ref = os.path.join(self.refdir, "trajectories.ref.xyz")
    self.pdb_file = os.path.join(self.workdir, "trajectories.pdb")
    self.pdb_ref = os.path.join(self.refdir, "trajectories.ref.pdb")
    self.lammps_file = os.path.join(self.workdir, "trajectories.lammpstrj")
    self.lammps_ref = os.path.join(self.refdir, "trajectories.ref.lammpstrj")

    # Reference data by configuration
    # Values correspond to E/k_b in Kelvin
    # E_disp, E_LRC, E_real, E_fourier, E_self, E_intra, E_total
    # E_real does not include interaction from particules in the same molecules
    # E_intra is the correction for particules in the same molecule for fourier space only
    self.Eref = np.array([
        [9.95387e+04, -8.23715e+02, -5.58889e+05, 6.27009e+03, -2.84469e+06, 2.80999e+06, -4.88604e+05],
        [1.93712e+05, -3.29486e+03, -1.19295e+06, 6.03495e+03, -5.68938e+06, 5.61998e+06, -1.06590e+06],
        [3.54344e+05, -7.41343e+03, -1.96297e+06, 5.24461e+03, -8.53407e+06, 8.42998e+06, -1.71488e+06],
        [4.48593e+05, -1.37286e+04, -3.57226e+06, 7.58785e+03, -1.42235e+07, 1.41483e+07, -3.20501e+06]])

  def read_Ebreakdown(self, energy_file):
      """Read energy break down file

      col 1: Potential energy
      col 2: Coulomb LR energy
      col 3: Coulomb SR energy
      col 4: Coulomb self energy
      col 5: Coulomb molecule energy

      return E: energies by time step
      """
      E = np.loadtxt(energy_file, usecols=(1,2,3,4,5))
      return E

  def check_values(self, rtol=None, atol=None):
      """Compare ref values against actual values

      """

      if rtol is None:
        rtol = self.default_rtol
      if atol is None:
        atol = self.default_atol
    
      Emw = self.read_Ebreakdown(self.energy_file)
      Eref_hartree = self.Eref * (self.boltzmann / self.hartree2J)

      Eref_select = np.array([Eref_hartree[0,2]+Eref_hartree[0,5], # coulomb real + intra
                              Eref_hartree[0,3],                        # coulomb fourier
                              Eref_hartree[0,4]])                        # coulomb self


      Emw_select = np.array([Emw[2] + Emw[4],               # coulomb real + intra
                             Emw[1],                        # coulomb fourier
                             Emw[3]])                        # coulomb self

      allclose = np.allclose(Eref_select, Emw_select, rtol, atol)
      msg = "Ok"
      if (not allclose):
        msg = "\nEref_select:\n" + str(Eref_select) + "\n"
        msg += "Emw_select:\n" + str(Emw_select) + "\n"
        msg += "err:\n" + str((Eref_select-Emw_select) / Eref_select) + "\n"

      return allclose, msg

class test_python_interface(unittest.TestCase):

  @classmethod
  def setUpClass(self):
    self.nist_path = "nist/SPCE-water-reference-calculations/config1"
    self.python_interface_path = "python_interface/nist"
    self.mw_exec = os.path.abspath(mwrun.glob_mw_exec)

  def run_conf(self, nranks):
    n = python_interface(self.mw_exec, self.python_interface_path, self.nist_path)
    self.workdir = n.workdir
    self.refdir = n.refdir
    #n.run_mw(nranks)
    n.run_mw_python(nranks, 'nist.py')
    ok, msg = n.check_values(atol=0.0)
    self.assertTrue(ok, msg)
    ok, msg = n.check_xyz()
    self.assertTrue(ok, msg)
    ok, msg = n.check_pdb()
    self.assertTrue(ok, msg)
    ok, msg = n.check_lammps()
    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_nist1(self):
    self.run_conf(1)

  def test_conf_nist1_4MPI(self):
    self.run_conf(4)
