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

import mwrun

class test_plumed(unittest.TestCase):

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

  def run_conf_compare_MW_plumed(self, dirname1, dirname2, nranks):
    path_to_config = os.path.join(self.test_path, dirname2)
    n = mwrun.mwrun(self.mw_exec, path_to_config)
    self.workdirs = [n.workdir]
    n.run_mw(nranks)


    path_to_config = os.path.join(self.test_path, dirname1)
    n = mwrun.mwrun(self.mw_exec, path_to_config)
    self.workdirs.append(n.workdir)
    n.run_mw(nranks)

    path1=os.path.join(self.test_path, dirname1,"energies.out")
    path2=os.path.join(self.test_path, dirname2,"energies.out")

    ok, msg = n.compare_columns_datafiles(path1,path2,1)
    self.assertTrue(ok, msg)

  def run_conf(self, dirname, nranks):
    path_to_config = os.path.join(self.test_path, dirname)
    n = mwrun.mwrun(self.mw_exec, path_to_config)
    self.workdirs = [n.workdir]
    n.run_mw(nranks)

    ok, msg = n.compare_datafiles("forces.out", "forces.ref")
    self.assertTrue(ok, msg)

    ok, msg = n.compare_datafiles("energies.out", "energies.ref")
    self.assertTrue(ok, msg)

    ok, msg = n.compare_datafiles("energies_breakdown.out", "energies_breakdown.ref")
    self.assertTrue(ok, msg)

    ok, msg = n.compare_datafiles("colvar","colvar.ref")
    self.assertTrue(ok, msg)

  def tearDown(self):
    for w in self.workdirs:  
      for f in glob.glob(os.path.join(w, "*.out")):
        os.remove(f)
      exists = os.path.isfile(os.path.join(w, "colvar"))
      if (exists):
          os.remove(os.path.join(w, "colvar"))

  def test_plumed_spring_serial(self):
    self.run_conf("spring", 1)

  def test_plumed_spring_parallel(self):
    self.run_conf("spring", 4)

  def test_plumed_compare_plumed_noplumed_parallel(self):
      self.run_conf_compare_MW_plumed( "spring", "spring_noplumed", 4)

  def test_plumed_compare_plumed_noplumed_serial(self):
      self.run_conf_compare_MW_plumed( "spring", "spring_noplumed", 1)
