
import numpy as np
import os.path
import sys
import filecmp
import subprocess

out = open("benchmark.out",'w')
out.close()

glob_mpi_launcher = ""
glob_mw_exec = ""
glob_python_path = ""
glob_skip_long_test = False

class mwrun(object):
  def __init__(self, mw_exec, path_to_config):
    # Constants
    self.mw_exec = mw_exec
    self.default_rtol=1.0e-5
    self.default_atol=1.0e-7  # force au
    self.workdir = path_to_config

  def run_mw(self, nranks):
    cwd = os.getcwd()
    os.chdir(self.workdir)
    cmd_line = glob_mpi_launcher.split(" ")
    cmd_line.extend(["-n", str(nranks), self.mw_exec])
    subprocess.check_call(cmd_line)
    os.chdir(cwd)

  def run_mw_python(self, nranks, python_script):
    cwd = os.getcwd()
    os.chdir(self.workdir)
#    pypath = "PYTHONPATH=$PYTHONPATH:"+cwd+"/../build/python/:"+cwd+"/../src/"
#    cmd_line = pypath.split(" ")
    cmd_line = glob_mpi_launcher.split(" ")
    cmd_line.extend(["-n", str(nranks), glob_python_path, python_script])
    subprocess.check_call(cmd_line)
    os.chdir(cwd)

  def compare_datafiles(self, file_out, file_ref, rtol=None, atol=None):
    """Compare two datasets from 2 runs"""
    if rtol is None:
      rtol = self.default_rtol
    if atol is None:
      atol = self.default_atol

    # Load data from ref file
    path_ref = os.path.join(self.workdir, file_ref)
    data_ref = np.loadtxt(path_ref)

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

    #atol = 1.0e-7
    #m = np.abs(data_ref).mean()
    #if m < 1.0e-8:
    #  atol = m / 1000.0
    # compare data
    allclose = np.allclose(data_out, data_ref, rtol, 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]),])

    return allclose, msg

  def compare_columns_datafiles(self, file_out, file_ref, col_ind, rtol=None, atol=None ):
    """Compare a single column from datasets from 2 runs"""
    if rtol is None:
      rtol = self.default_rtol
    if atol is None:
      atol = self.default_atol
    

    # Load data from ref file
    path_ref = file_ref
    data_ref = np.loadtxt(path_ref)

    # Load data from output file
    path_out = file_out 
    data_out = np.loadtxt(path_out)

    first_col_ref=data_ref[:,col_ind]
    first_col_out=data_out[:,col_ind]
    #atol = 1.0e-7
    #m = np.abs(data_ref).mean()
    #if m < 1.0e-8:
    #  atol = m / 1000.0

    # compare data
    allclose = np.allclose(first_col_ref,first_col_out, rtol, 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]),])

    return allclose, msg    



  def diff_files(self, out, ref):
    allsame = filecmp.cmp(out, ref)
    if allsame:
      msg = "{:s}: Ok".format(out)
    else:
      msg = "\n%s different from %s\n"%(out, ref)

    return allsame, msg

  def check_xyz(self):
    return self.diff_files(self.xyz_file, self.xyz_ref)

  def check_pdb(self):
    return self.diff_files(self.pdb_file, self.pdb_ref)

  def check_lammps(self):
    return self.diff_files(self.lammps_file, self.lammps_ref)

  def check_forces(self, rtol=None, atol=None):
    return self.compare_datafiles(self.forces_files, self.forces_ref, rtol, atol)

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

    # Load data from ref file
    path_ref = os.path.join(self.workdir, file_ref)
    fil_ref = open(path_ref, 'rb')
    data = np.fromfile(fil_ref, dtype=np.float, count = -1)
    size = int(np.sqrt(np.shape(data)[0] - 1 - 3*10 - 2))
    data_ref = np.reshape(data[-size*size:], (size, size), order='F')
    fil_ref.close()

    # Load data from output file
    path_out = os.path.join(self.workdir, file_out)
    fil_out = open(path_out, 'rb')
    data = np.fromfile(fil_out, dtype=np.float, count = -1)
    size = int(np.sqrt(np.shape(data)[0] - 1 - 3*10 - 2))
    data_out = np.reshape(data[-size*size:], (size, size), order='F')
    fil_out.close()

    # compare data
    allclose = np.allclose(data_out, data_ref, rtol, 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]),])

    return allclose, msg

