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

import mwrun

class test_benchmark(unittest.TestCase):

  @classmethod
  def setUpClass(self):
    self.test_path = "benchmark"
    self.mw_exec = os.path.abspath(mwrun.glob_mw_exec)
    self.au2kcalmolang = 1185.8210465697064
    self.tol = {}
    self.tol["rtol_forces"] = 1.0e-7
    self.tol["atol_forces"] = 1.0e-5
    self.tol["rtol_charges"] = 1.0e-7
    self.tol["atol_charges"] = 1.0e-2

  def compare_lammpsforces(self, file_out, file_ref, rtol, atol):
    """Compare two datasets from 2 runs"""

    # Load data from ref file
    path_ref = os.path.join(self.workdir, "../lammps/", file_ref)
    data_ref = np.loadtxt(path_ref, usecols=(5,6,7), skiprows=9)
    data_ref /= self.au2kcalmolang

    # Load data from output file
    path_out = os.path.join(self.workdir, file_out)
    data_out = np.loadtxt(path_out, usecols=(0,1,2))

    # compare data
    allclose = np.allclose(data_out, data_ref, rtol, abs(data_ref).mean() * atol)

    if allclose:
      msg = "{:s}: Ok".format(file_out)
    else:
      msg = '\n'.join(
        ["{:s}: Error".format(file_out),
         "{:s}: ".format(file_out), str(data_out[:10]),
         "{:s}: ".format(file_ref), str(data_ref[:10]),
         "{:s}: ".format("relative + absolute error"), str(np.abs((data_out[:10]-data_ref[:10])/(1+data_ref[:10]))),
         "{:s}: ".format("max relative error"), str(np.nanmax(np.abs((data_out[:]-data_ref[:])/data_ref[:]))),
         "{:s}: ".format("max relative + absolute error"), str(np.nanmax(np.abs((data_out[:]-data_ref[:])/(data_ref[:]+1))))])

#    with open('benchmark.out', 'a') as out:
#      report = '\n'.join(
#        ["{:s}: Report".format(self.workdir),
#         "{:s}: ".format(file_out), str(data_out[:10]),
#         "{:s}: ".format(file_ref), str(data_ref[:10]),
#         "{:s}: ".format("relative error"), str(np.abs((data_out[:10]-data_ref[:10])/(data_ref[:10]))),
#         "{:s}: ".format("max relative error"), str(np.nanmax(np.abs((data_out[:]-data_ref[:])/data_ref[:]))),
#         "{:s}: ".format("max relative + absolute error"), str(np.nanmax(np.abs((data_out[:]-data_ref[:])/(data_ref[:]+1)))),
#         "{:s}: ".format("mean value"), str(np.mean(np.abs(data_ref[:])))])
#      out.write(report)
 #     out.write("\n\n")

    return allclose, msg

  def compare_lammps_elecforces(self, file_out, file_ref, rtol, atol):
    """Compare two datasets from 2 runs"""

    # Load data from ref file
    path_ref = os.path.join(self.workdir, "../lammps/", file_ref)
    data_ref = np.loadtxt(path_ref, usecols=(2,3,4,), skiprows=9)
    data_ref /= self.au2kcalmolang

    # Load data from output file
    path_out = os.path.join(self.workdir, file_out)
    data_out = np.loadtxt(path_out, usecols=(0,1,2,))

    # compare data
    allclose1 = np.allclose(data_out[0, :], data_ref[0, :], rtol, atol)
    allclose2 = np.allclose(data_out[-1, :], data_ref[-1, :], rtol, atol)

    if allclose1 and allclose2:
      msg = "{:s}: Ok elec1".format(file_out)
    else:
      msg = '\n'.join(
        ["{:s}: Error".format(file_out),
         "{:s}: ".format(file_out), str(data_out),
         "{:s}: ".format(file_ref), str(data_ref)])

#    with open('benchmark.out', 'a') as out:
#      report = '\n'.join(
#        ["{:s}: Report".format(self.workdir),
#         "{:s}: ".format(file_out), str(data_out),
#         "{:s}: ".format(file_ref), str(data_ref)])
#      out.write(report)
#      out.write("\n\n")

    return (allclose1 and allclose2), msg

  def compare_lammpscharges(self, file_out, file_ref, rtol, atol):
    """Compare two datasets from 2 runs"""

    # Load data from ref file
    path_ref = os.path.join(self.workdir, "../lammps/", file_ref)
    data_ref = np.loadtxt(path_ref, usecols=(2,), skiprows=9)

    # Load data from output file
    path_out = os.path.join(self.workdir, file_out)
    data_out = np.loadtxt(path_out, usecols=(0,))

    start = 64
    stop = 127

    # compare data
    allclose = np.allclose(data_out[start:stop+1], data_ref[start:stop+1], rtol, abs(data_ref[start:stop+1]).mean() * atol)

    if allclose:
      msg = "{:s}: Ok".format(file_out)
    else:
      msg = '\n'.join(
        ["{:s}: Error".format(file_out),
         "{:s}: ".format(file_out), str(data_out[start:stop+1]),
         "{:s}: ".format(file_ref), str(data_ref[start:stop+1]),
         "{:s}: ".format("relative + absolute error"), str(np.abs((data_out[start:stop+1]-data_ref[start:stop+1])/(1+data_ref[start:stop+1]))),
         "{:s}: ".format("max relative error"), str(np.nanmax(np.abs((data_out[start:stop+1]-data_ref[start:stop+1])/data_ref[start:stop+1]))),
         "{:s}: ".format("max relative + absolute error"), str(np.nanmax(np.abs((data_out[start:stop+1]-data_ref[start:stop+1])/(data_ref[start:stop+1]+1))))])

#    with open('benchmark.out', 'a') as out:
#      report = '\n'.join(
#        ["{:s}: Report".format(self.workdir),
#         "{:s}: ".format(file_out), str(data_out[start:stop+1]),
#         "{:s}: ".format(file_ref), str(data_ref[start:stop+1]),
#         "{:s}: ".format("relative error"), str(np.abs((data_out[start:stop+1]-data_ref[start:stop+1])/(data_ref[start:stop+1]))),
#         "{:s}: ".format("max relative error"), str(np.nanmax(np.abs((data_out[start:stop+1]-data_ref[start:stop+1])/data_ref[start:stop+1]))),
#         "{:s}: ".format("max relative + absolute error"), str(np.nanmax(np.abs((data_out[start:stop+1]-data_ref[start:stop+1])/(data_ref[start:stop+1]+1)))),
#         "{:s}: ".format("mean value"), str(np.mean(np.abs(data_ref[start:stop+1])))])
#      out.write(report)
#      out.write("\n\n")

    return allclose, msg

  def run_conf(self, confID, nranks, check_forces=True, check_charges=False, check_elec=False, tol=None):
    path_to_config = os.path.join(self.test_path, confID, "mw2")
    n = mwrun.mwrun(self.mw_exec, path_to_config)
    self.workdir = n.workdir
    n.run_mw(nranks)
    if tol == None:
      tol = self.tol

    if check_elec:
      ok, msg = self.compare_lammps_elecforces("electrode_forces.out", "elec_forces.ref", tol["rtol_forces"], tol["atol_forces"])
      self.assertTrue(ok, msg)

    if check_forces:
       ok, msg = self.compare_lammpsforces("forces.out", "forces.ref", tol["rtol_forces"], tol["atol_forces"])
       self.assertTrue(ok, msg)

    if check_charges:
      ok, msg = self.compare_lammpscharges("charges.out", "charges.ref", tol["rtol_charges"], tol["atol_charges"])
      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"))

  def test_conf_3D_bulkNaCl(self):
    self.run_conf("1-3D_bulkNaCl", 8)

  def test_conf_2D_NaCl_lv_lj(self):
    self.run_conf("2-2D_NaCl_lv/1-lj-noelec", 8)

  def test_conf_2D_NaCl_lv(self):
    self.run_conf("2-2D_NaCl_lv/2-lj-elec", 8)

  def test_conf_2D_capacitor(self):
    self.run_conf("3-capacitor", 12, check_charges=True)

  def test_conf_2D_elec_forces(self):
    self.run_conf("4-piston", 12, check_elec=True)

  def test_conf_2D_elec_forces_charged(self):
    tol = {}
    tol["rtol_forces"] = 1.0e-7
    tol["atol_forces"] = 1.0e-4
    self.run_conf("5-piston-charged", 12, tol=tol, check_elec=True)

  def test_conf_2D_easy_charged(self):
    self.run_conf("6-easy-charged", 4, check_elec=True, check_forces=False)

