import unittest
import sys
import os
import shutil

import mw_utils

runout={
'finished': """ Short range energy    36.1915995050419     
 Velocity convergence            4
 **** Dumped pot, elefld and density ****
 
 **** Radial distribution function written out ****
 
 Total time taken:   142.654646158218     
 IO time:  0.796273231506348     
 Three first steps:   11.0714640617371     
 Bye Bye""",
 
'cg_failed': """ Dispersion energy   0.000000000000000E+000
 vqeng   -2.90284568630675     
 Short range energy    30.4749414724436     
 Velocity convergence            3
    
 Step    15120439       61439
 Done           22  rattle iterations -1st part
 Only the two first species are considered electrode species
 which corresponds to         7298  atoms
 cg failed to converge""",
 
'unknown': """ Dispersion energy   0.000000000000000E+000
 vqeng   -2.90284568630675     
 Short range energy    30.4749414724436     
 Velocity convergence            3
    
 Step    15120439       61439
 Done           22  rattle iterations -1st part
 Only the two first species are considered electrode species
 which corresponds to         7298  atoms"""
}

class test_status(unittest.TestCase):


  
  def setup(self, dir_name, run_version):
    if dir_name != ".":
      os.makedirs(dir_name)
      
    self.file_runout = os.path.join(dir_name, 'run.out')
    self.dir_name = dir_name

    f = open(self.file_runout, 'w')
    f.write(runout[run_version])
    f.close()

  def tearDown(self):
    os.remove(self.file_runout)
    if self.dir_name != ".":
      os.rmdir(self.dir_name)
      
  
  def test_finished(self):
    test_dir= "."
    self.setup(test_dir, 'finished')
    res = mw_utils.status(test_dir)
    expected = 'finished'
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
  def test_cg_failed(self):
    test_dir= "."
    self.setup(test_dir, 'cg_failed')
    res = mw_utils.status(test_dir)
    expected = 'cg_failed'
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
  def test_unknown(self):
    test_dir= "."
    self.setup(test_dir, 'unknown')
    res = mw_utils.status(test_dir)
    expected = 'unknown'
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
  def test_dir(self):
    test_dir= "run098"
    self.setup(test_dir, 'finished')
    res = mw_utils.status(test_dir)
    expected = 'finished'
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
  def test_norunout(self):
    test_dir= "run098"
    self.setup(test_dir, 'finished')
    res = mw_utils.status(".")
    expected = "Error: not a valid run directory, [Errno 2] No such file or directory: './run.out'"
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
    
class test_restart(unittest.TestCase):
  
  def create_file(self, filename):
    f = open(filename,"w")
    f.write("%s\n"%filename)
    f.close()

  def setup(self, dir_name, bug_restart=False, bug_job=False):
    if dir_name != ".":
      os.makedirs(dir_name)
      
    self.dir_name = dir_name

    self.restart = []
    if not bug_restart:
      for i in range(0,4000,1000):
        self.restart.append( os.path.join(dir_name, "restart.1511%04d.out"%i) )
        self.create_file(self.restart[-1])

    self.runtime = os.path.join(dir_name, "runtime.inpt")
    self.create_file(self.runtime)

  def tearDown(self):
    for i in self.restart:
      os.remove(i)

    os.remove(self.runtime)

    if self.dir_name != ".":
      os.rmdir(self.dir_name)
      
  
  def test_build_run_to_1(self):
    test_dir= "run065"
    self.setup(test_dir)
    path, res = mw_utils.build_run_to(test_dir)
    expected = 'run066'
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
  def test_build_run_to_1b(self):
    test_dir= "test-yaml_065"
    self.setup(test_dir)
    path, res = mw_utils.build_run_to(test_dir)
    expected = 'test-yaml_066'
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
  def test_build_run_to_2(self):
    test_dir= "run065"
    self.setup(test_dir)
    cwd = os.getcwd()
    os.chdir(test_dir)
    path, res = mw_utils.build_run_to(".")
    expected = 'run066'
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    os.chdir(cwd)
    
  def test_build_run_to_3(self):
    test_dir= "run065oo"
    self.setup(test_dir)
    path, res = mw_utils.build_run_to(test_dir)
    expected = None
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
  def test_build_run_to_4(self):
    test_dir= "99run065"
    self.setup(test_dir)
    path, res = mw_utils.build_run_to(test_dir)
    expected = None
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
  def test_build_run_to_5(self):
    test_dir= "run065/"
    self.setup(test_dir)
    path, res = mw_utils.build_run_to(test_dir)
    expected = 'run066'
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
  def test_find_restart_1(self):
    test_dir= "run069"
    self.setup(test_dir)
    res = mw_utils.find_restart(test_dir)
    expected = os.path.join(test_dir,'restart.15113000.out')
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
  def test_find_restart_2(self):
    test_dir= "run069"
    self.setup(test_dir, bug_restart=True)
    res = mw_utils.find_restart(test_dir)
    expected = None
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    
  def test_restart_1(self):
    test_dir= "run071"
    self.setup(test_dir)
    res = mw_utils.restart(test_dir, None)
    res = os.path.basename(res)
    expected = "run072"
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    self.assertTrue(os.path.exists("run072")==True,  'Directory run072 not created')
    self.assertTrue(os.path.exists("run072/data.inpt")==True,  'restart file not copied')
    self.assertTrue(os.path.exists("run072/runtime.inpt")==True,  'runtime.inpt file not copied')
    f = open("run072/data.inpt", 'r')
    l = f.readlines()
    f.close()
    res = l[0]
    expected = "run071/restart.15113000.out\n"
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))

    shutil.rmtree("run072")
    
  def test_restart_2(self):
    test_dir= "run071"
    self.setup(test_dir)

    res = mw_utils.restart(None, None)
    res = os.path.basename(res)
    expected = "run072"
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))
    self.assertTrue(os.path.exists("run072")==True,  'Directory run072 not created')
    self.assertTrue(os.path.exists("run072/data.inpt")==True,  'restart file not copied')
    self.assertTrue(os.path.exists("run072/runtime.inpt")==True,  'runtime.inpt file not copied')
    f = open("run072/data.inpt", 'r')
    l = f.readlines()
    f.close()
    res = l[0]
    expected = "run071/restart.15113000.out\n"
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))

    shutil.rmtree("run072")
    
  
  def test_restart_3(self):
    test_dir= "run071"
    self.setup(test_dir)

    os.makedirs("run072")
    res = mw_utils.restart(test_dir, None)
    expected = None
    self.assertTrue(res==expected,  'Expecting %s but got %s'%(expected,  res))

    shutil.rmtree("run072")
    
  
def runTest(method_to_test):
  if len(method_to_test) == 0:
    suite = test_all()
  else:
    suite = unittest.TestSuite()
    for i in method_to_test:
      suite.addTest(test_restart(i))

  unittest.TextTestRunner(verbosity=2).run(suite)
  
def test_all():
  status_suite = unittest.TestLoader().loadTestsFromTestCase(test_status)
  restart_suite = unittest.TestLoader().loadTestsFromTestCase(test_restart)
  all_tests = unittest.TestSuite([status_suite, restart_suite])
  return all_tests
  
if __name__ == "__main__":
  """Without argument, all tests are executed otherwise only methods put as argument are executed"""
  runTest(sys.argv[1:])
  
