#  ============================================================================
#
#  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 =============================
#  =======================  Latest update: 6/8/09 ============================= 

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 = 'model1.pdb'           # input model for optimization
code = 'mdl1'                           # 4 letter code of the structure
em_map_file = 'map.mrc'                 # name of EM density map (mrc)
format='MRC'                            # map format: MRC or XPLOR
apix=1.5                                # voxel size: A/pixel
box_size=90                             # size of the density map (cubic)
resolution=10.0                         # resolution
x=-31.00; y=-36.0; z=-11.00             # origin of the map
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 = 4                         # number of iterations in each MD run 
          				# (default=4) 
rigid_filename = 'rigid.txt'            # rigid bodies file name


############### 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, box_size, 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, box_size, resolution, x, y, z, rigid_filename, num_of_iter)