#  ============================================================================
#
#  a. Define the type of oprimization: conjugate gradients minimization (CG)
#     or molecular dynamics (MD)
#  b. Edit input parameters: density map parameters and probe structure
#
#  =======================  Maya Topf, 4 Dec 2007 =============================
#  ============== Maya Topf, Agnel Joseph, Latest update: 15/5/18 =============

from modeller import *
import shutil
import sys, os, os.path
import string
import math
from CG import opt_cg
from MD import opt_md

env = environ()

############### INPUT PARAMETERS ##################
optimization = 'MD'                     # type of optimization: CG / MD
input_pdb_file = 'mdl1.pdb'          # input model for optimization
code = 'mdl1'                           # 4 letter code of the structure
em_map_file = '1ake_molmap45_pad10_cubic.mrc'     # name of EM density map (mrc)
format='MRC'                            # map format: MRC or XPLOR
apix=1.5                                  # voxel size: A/pixel
resolution=4.5                         # resolution (A)
x=-11.994; y=7.881; z=-10.910                           # origin of the map (A)
path = './'	                        # path to work directory
init_dir = 1                            # number of the initial directory 
num_of_runs = 1                         # number of flex-em CG runs 
num_of_iter = 2                         # number of iterations in each MD run 
          				# (default=4) 
rigid_filename = 'rigid.txt'            # rigid bodies file name
cap_shift = 0.39			# change this for high resolution refinement (sse restrained (0.15), all-atom refinement (0.10))

############### RUN OPTIMIZATION ##################
path = os.path.abspath(path)

#  CG
# ----
if optimization == 'CG':
    for i in range(init_dir,init_dir+num_of_runs):
        scratch = path + '/' + str(i) + '_cg/'
        os.system("mkdir -p " + scratch)
        os.system("cp " + path + '/'+ input_pdb_file + " " + scratch)
        os.chdir(scratch)
        opt_cg(path, code, str(i), 55*i, em_map_file, input_pdb_file,
                format, apix, resolution, x, y, z, rigid_filename)

#  MD
# ----
elif optimization == 'MD':
        scratch = path + '/' + str(init_dir) + '_md/'
        os.system("mkdir -p " + scratch)
        os.system("cp " + path + '/'+ input_pdb_file + " " + scratch)
        os.chdir(scratch)
        opt_md(path, code, str(init_dir), 10*init_dir, em_map_file, input_pdb_file,
                format, apix, resolution, x, y, z, rigid_filename, num_of_iter,cap_shift)