"""
nist_validation.py

This script validate MW2 output against nist reference data
"""
import unittest
import argparse
import sys

import mwrun
from nist import test_nist
from tosi_fumi import test_tosi_fumi
from xft import test_xft
from pim import test_pim
from dip_plus_elec import test_dip_plus_elec
from dihedrals import test_dihedrals
from aim import test_aim
from four_site_model import test_four_site_model
from maze import test_maze
from matrix_inversion import test_matrix_inversion
from benchmark import test_benchmark
from plumed import test_plumed
from steele import test_steele
from unwrap import test_unwrap
from dump_per_species import test_dump_per_species
from charge_neutrality import test_charge_neutrality
from python_interface import test_python_interface
from external_field import test_external_field
from thomas_fermi import test_thomas_fermi
from non_neutral import test_non_neutral
from piston import test_piston


def runTest(args):
  if args.test_suite == "all":
    suite_list = test_all(args.reduced)
  elif args.test_suite == "long":
    suite_list = test_long()
  else:
    # Transforms string given by user into test object
    try:
      test_object = globals()["test_"+args.test_suite]
    except KeyError:
      print("Error: %s not a valid test suite"%args.test_suite)
      sys.exit(1)

    if len(args.tests) == 0:
      simple_suite = unittest.TestLoader().loadTestsFromTestCase(test_object)
    else:
      simple_suite = unittest.TestSuite()
      for i in args.tests:
        simple_suite.addTest(test_object(i))
    suite_list = [simple_suite]

  suite = unittest.TestSuite(suite_list)

  res = unittest.TextTestRunner(stream=sys.stdout, verbosity=2).run(suite)
  if not (res.errors == [] and  res.failures == []):
    sys.exit(1)

def test_long():
  suite_maze = unittest.TestLoader().loadTestsFromTestCase(test_maze)
  suite_matrix_inversion = unittest.TestLoader().loadTestsFromTestCase(test_matrix_inversion)
  suite_plumed = unittest.TestLoader().loadTestsFromTestCase(test_plumed)
  suite_dip_plus_elec = unittest.TestLoader().loadTestsFromTestCase(test_dip_plus_elec)
  alltests = [suite_maze, suite_matrix_inversion, suite_plumed, suite_dip_plus_elec]

  return alltests

def test_all(minimal_suite):
  suite_nist = unittest.TestLoader().loadTestsFromTestCase(test_nist)
  suite_tosi_fumi = unittest.TestLoader().loadTestsFromTestCase(test_tosi_fumi)
  suite_xft = unittest.TestLoader().loadTestsFromTestCase(test_xft)
  suite_pim = unittest.TestLoader().loadTestsFromTestCase(test_pim)
  suite_dip_plus_elec = unittest.TestLoader().loadTestsFromTestCase(test_dip_plus_elec)
  suite_dihedrals = unittest.TestLoader().loadTestsFromTestCase(test_dihedrals)
  suite_aim = unittest.TestLoader().loadTestsFromTestCase(test_aim)
  suite_four_site_model = unittest.TestLoader().loadTestsFromTestCase(test_four_site_model)
  suite_maze = unittest.TestLoader().loadTestsFromTestCase(test_maze)
  suite_matrix_inversion = unittest.TestLoader().loadTestsFromTestCase(test_matrix_inversion)
  suite_benchmark = unittest.TestLoader().loadTestsFromTestCase(test_benchmark)
  suite_charge_neutrality = unittest.TestLoader().loadTestsFromTestCase(test_charge_neutrality)
  suite_external_field = unittest.TestLoader().loadTestsFromTestCase(test_external_field)
  suite_thomas_fermi = unittest.TestLoader().loadTestsFromTestCase(test_thomas_fermi)
  suite_non_neutral = unittest.TestLoader().loadTestsFromTestCase(test_non_neutral)
  suite_piston = unittest.TestLoader().loadTestsFromTestCase(test_piston)
  suite_steele = unittest.TestLoader().loadTestsFromTestCase(test_steele)
  suite_plumed = unittest.TestLoader().loadTestsFromTestCase(test_plumed)
  suite_python_interface = unittest.TestLoader().loadTestsFromTestCase(test_python_interface)
  suite_unwrap = unittest.TestLoader().loadTestsFromTestCase(test_unwrap)
  suite_dump_per_species = unittest.TestLoader().loadTestsFromTestCase(test_dump_per_species)

  alltests = [suite_nist, suite_tosi_fumi, suite_xft, suite_pim, suite_dip_plus_elec, suite_dihedrals,
              suite_aim, suite_steele, suite_maze, suite_matrix_inversion, suite_benchmark,
              suite_charge_neutrality, suite_external_field, suite_thomas_fermi,
              suite_non_neutral, suite_piston, suite_unwrap, suite_four_site_model, suite_dump_per_species]

  if minimal_suite:
    print("Plumed and Python test suite skipped")
  else:
    alltests.append(suite_plumed)
    alltests.append(suite_python_interface)

  return alltests

def set_global_vars(args):
  mwrun.glob_mpi_launcher = args.mpi_launcher
  mwrun.glob_mw_exec = args.mw_exec
  mwrun.glob_python_path = args.python_path
  mwrun.glob_skip_long_test = args.reduced
  mwrun.glob_compare_hessian = args.intel_proc


def build_parser():
  # create the top-level parser
  parser = argparse.ArgumentParser(prog='test_MW')
  parser.add_argument('-r', '--reduced', action="store_true", help='Run a reduced test suite that skip the long tests')
  parser.add_argument('-m', '--mpi_launcher',  type=str, nargs='?', default="mpirun", help='MPI launcher to be used')
  parser.add_argument('-e', '--mw_exec',  type=str, nargs='?', default="../mw", help='Path to Metalwalls binary')
  parser.add_argument('-py', '--python_path',  type=str, nargs='?', default="python", help='Path to Python')
  parser.add_argument('-s', '--test_suite',  type=str, nargs='?', default="all", help='Test suite to be executed (default all except python_interface).\nAvailable test suites: nist, tosi_fumi, xft, pim, unwrap, dump_per_species, dihedrals, dip_plus_elec, aim, four_site_model, steele, maze, benchmark, plumed, charge_neutrality, external_field, thomas_fermi, non_neutral, piston, python_interface, long')
  parser.add_argument('-intel', '--intel_proc', action="store_true", help='Flag for Intel architecture')
  parser.add_argument('tests',  type=str, nargs=argparse.REMAINDER, default=[], help='Lists of tests to execute (separated with white spaces). If empty, executes all the tests of the test_suite defined with --test_suite')
  return parser

if __name__ == "__main__":
  """Without argument, all tests are executed otherwise only methods put as argument are executed"""
  p =  build_parser()
  args = p.parse_args(sys.argv[1:])
  set_global_vars(args)
  runTest(args)
