#!/usr/bin/env python

import os
import sys
import json
from os import path
import radical.utils.logger as rul
from repex_utils.replica_cleanup import *
from repex_utils.parser import parse_command_line
from pilot_kernels.pilot_kernel_pattern_s         import PilotKernelPatternS
from pilot_kernels.pilot_kernel_pattern_s_multi_d import PilotKernelPatternSmultiD

from amber_kernels_tex.kernel_pattern_s_tex       import KernelPatternStex
from amber_kernels_us.kernel_pattern_s_us         import KernelPatternSus
from amber_kernels_salt.kernel_pattern_s_salt     import KernelPatternSsalt
from amber_kernels_3d_tsu.kernel_pattern_s_3d_tsu import KernelPatternS3dTSU
from amber_kernels_3d_tuu.kernel_pattern_s_3d_tuu import KernelPatternS3dTUU


#-------------------------------------------------------------------------------

if __name__ == '__main__':
    
    name = 'amber'
    logger  = rul.getLogger ('radical.repex', name)

    work_dir_local = os.getcwd()
    params = parse_command_line()
    
    # get remd_input file
    json_data=open(params.remd_input)
    inp_file = json.load(json_data)
    json_data.close()

    # get resource config
    json_data=open(params.resource_config)
    rconfig = json.load(json_data)
    json_data.close()


    if (inp_file['remd.input']['re_pattern'] == 'S') and (inp_file['remd.input']['exchange'] == 'T-REMD'):

		logger.info("*********************************************************")
		logger.info("*           RepEx: T-REMD with AMBER kernel             *")
		logger.info("*********************************************************")
        
	    # initializing kernels
		md_kernel = KernelPatternStex( inp_file, rconfig, work_dir_local )
		pilot_kernel = PilotKernelPatternS( inp_file, rconfig )
        
		# initializing replicas
		replicas = md_kernel.initialize_replicas()
        
		try:
			pilot_manager, pilot_object, session = pilot_kernel.launch_pilot()
	        
			# now we can run RE simulation
			pilot_kernel.run_simulation( replicas, pilot_object, session, md_kernel )
            
			# this is a quick hack
			base = md_kernel.inp_basename + ".mdin"
		except:
			raise
		try:
			# finally we are moving all files to individual replica directories
			move_output_files(work_dir_local, base, replicas ) 
			session.close(cleanup=False)
            
			logger.info("Simulation successfully finished!")
			logger.info("Please check output files in replica_x directories.")
            
		except:
			logger.info("Unexpected error: {0}".format(sys.exc_info()[0]) )
			raise 
            
		finally :
			logger.info("Closing session.")
			session.close (cleanup=False)  
            
    if (inp_file['remd.input']['re_pattern'] == 'S') and (inp_file['remd.input']['exchange'] == 'US-REMD'):
        
		logger.info("*********************************************************")
		logger.info("*           RepEx: US-REMD with AMBER kernel            *")
		logger.info("*********************************************************")

		# initializing kernels
		md_kernel = KernelPatternSus( inp_file, rconfig, work_dir_local )
		pilot_kernel = PilotKernelPatternS( inp_file, rconfig )

		# initializing replicas
		replicas = md_kernel.initialize_replicas()

		try:
			pilot_manager, pilot_object, session = pilot_kernel.launch_pilot()

			# now we can run RE simulation
			pilot_kernel.run_simulation( replicas, pilot_object, session, md_kernel)

			# this is a quick hack
			base = md_kernel.inp_basename + ".mdin"

			# finally we are moving all files to individual replica directories
			move_output_files(work_dir_local, base, replicas ) 

			logger.info("Simulation successfully finished!")
			logger.info("Please check output files in replica_x directories.")
            
		except:
			logger.info("Unexpected error: {0}".format(sys.exc_info()[0]) )
			raise
            
		finally :
			logger.info("Closing session.")
			session.close (cleanup=False) 
            
    if (inp_file['remd.input']['re_pattern'] == 'S') and (inp_file['remd.input']['exchange'] == 'S-REMD'):
        
		logger.info("*********************************************************")
		logger.info("*           RepEx: S-REMD with AMBER kernel             *")
		logger.info("*********************************************************")

		# initializing kernels
		md_kernel = KernelPatternSsalt( inp_file, rconfig, work_dir_local )
		pilot_kernel = PilotKernelPatternS( inp_file, rconfig )

		# initializing replicas
		replicas = md_kernel.initialize_replicas()
        
		try:
			pilot_manager, pilot_object, session = pilot_kernel.launch_pilot()

			# now we can run RE simulation
			pilot_kernel.run_simulation( replicas, pilot_object, session, md_kernel )

			# this is a quick hack
			base = md_kernel.inp_basename + ".mdin"

			# finally we are moving all files to individual replica directories
			move_output_files(work_dir_local, base, replicas ) 
			session.close(cleanup=False)

			logger.info("Simulation successfully finished!")
			logger.info("Please check output files in replica_x directories.")

		except:
			logger.info("Unexpected error: {0}".format(sys.exc_info()[0]) )
			raise 

		finally:
			logger.info("Closing session.")
			session.close (cleanup=False)    
            
    if (inp_file['remd.input']['re_pattern'] == 'S') and (inp_file['remd.input']['exchange'] == 'TSU-REMD'):
        
		logger.info("*********************************************************")
		logger.info("*           RepEx: TSU-REMD with AMBER kernel           *")
		logger.info("*********************************************************")
        
		# initializing kernels
		md_kernel = KernelPatternS3dTSU( inp_file, rconfig, work_dir_local )
		pilot_kernel = PilotKernelPatternSmultiD( inp_file, rconfig )

		# initializing replicas
		replicas = md_kernel.initialize_replicas()
        
		try:
			pilot_manager, pilot_object, session = pilot_kernel.launch_pilot()

			# now we can run RE simulation
			pilot_kernel.run_simulation( replicas, pilot_object, session, md_kernel )

			# this is a quick hack
			base = md_kernel.inp_basename + ".mdin"

			# finally we are moving all files to individual replica directories
			move_output_files(work_dir_local, base, replicas ) 

			logger.info("Simulation successfully finished!")
			logger.info("Please check output files in replica_x directories.")
            
		except:
			logger.info("Unexpected error: {0}".format(sys.exc_info()[0]) )
			raise
            
		finally:
			logger.info("Closing session.")
			session.close (cleanup=False)
            
    if (inp_file['remd.input']['re_pattern'] == 'S') and (inp_file['remd.input']['exchange'] == 'TUU-REMD'):
        
		logger.info("*********************************************************")
		logger.info("*           RepEx: TUU-REMD with AMBER kernel           *")
		logger.info("*********************************************************")

		# initializing kernels
		md_kernel    = KernelPatternS3dTUU( inp_file, rconfig, work_dir_local )
		pilot_kernel = PilotKernelPatternSmultiD( inp_file, rconfig )

		# initializing replicas
		replicas = md_kernel.initialize_replicas()
        
		try:
			pilot_manager, pilot_object, session = pilot_kernel.launch_pilot()

			# now we can run RE simulation
			pilot_kernel.run_simulation( replicas, pilot_object, session, md_kernel )

			base = md_kernel.inp_basename + ".mdin"

			# finally we are moving all files to individual replica directories
			move_output_files(work_dir_local, base, replicas )

			logger.info("Simulation successfully finished!")
			logger.info("Please check output files in replica_x directories.")
             
		except:
			logger.info("Unexpected error: {0}".format(sys.exc_info()[0]) )
			raise
        
		finally:
			# always clean up the session, no matter if we caught an exception or
			# not.
			logger.info("Closing session.")
			session.close (cleanup=False)

