#!/usr/bin/env python

# setup the environment

import os, sys, string, time
import shutil
import urllib

# Test for environment variables

if not os.environ.has_key('CCP4'):
    raise RuntimeError, 'CCP4 not found'
if not os.environ.has_key('BMP'):
    raise RuntimeError, 'BMP not found'
if not os.environ.has_key('http_proxy'):
    raise RuntimeError, 'http_proxy environment variable not found'

# Add the include directories to the system path and import the 
# necessary scripts

bmp = os.environ['BMP']
bmp_incl = os.environ['BMP_INCL']
bmp_data = os.environ['BMP_DATA']

scratch_dir=os.environ['CCP4_SCR']

BMP_MA_PROGRAM=os.environ['BMP_MA_PROGRAM']
BMP_USER=os.environ['USER']
BMP_USERHOME=os.environ['HOME']
BMP_SETUP_DIR=BMP_USERHOME + '/.ccp4_bmp'

CLUSTER=eval(os.environ['BMP_CLUSTER'])
BMP_DEBUG=eval(os.environ['BMP_DEBUG'])
BMP_SSM=eval(os.environ['BMP_SSM'])
BMP_SCOP=eval(os.environ['BMP_SCOP'])
BMP_PW_ALIGN=eval(os.environ['BMP_PW_ALIGN'])
BMP_FASTA_LOCAL=eval(os.environ['BMP_FASTA_LOCAL'])
BMP_RUN_LOCAL=eval(os.environ['BMP_RUN_LOCAL'])

os.environ["MODEL_ONLY"]="False"

sys.path.append(bmp_incl + '/dev')
sys.path.append(bmp_incl + '/file_info')
sys.path.append(bmp_incl + '/modelling')
sys.path.append(bmp_incl + '/mr')
sys.path.append(bmp_incl + '/output')
sys.path.append(bmp_incl + '/seq_align')
sys.path.append(bmp_incl + '/ccp4')
sys.path.append(bmp_incl + '/structures')
sys.path.append(bmp_incl + '/tools')

import Matches, BMP_target_info, Model
import Model_Queue, Struct_dump
import Model_struct, Ensemble 
import PQS 
import BMP_MR
import Developer
import Write_seq_align 
import Align_score 
import Update_DB
import Cleanup

# Testing: set mr to be true or false 

mr=False

# Define the usage and test to see if input is correct

def usage():
   useStr= """Usage: model_maker -[m <NUM_MODEL> --or-- p <MTZ file> <NUM_MODEL> <NUM_MR> <NUM_ENSEM> <MR_PROGRAM>]
                   ...<SEQ file> <job_id> <E_value> <IGNORE file>

                -m <NUM_MODEL>                    
                    --Model only (don't do molecular replacement)
                -p <MTZ file> <NUM_MOLEL> <NUM_MR> <NUM_ENSEM> <MR_PROGRAM>
                    --Do full modelling and molecular replacement
 
                NUM_MODEL  = Number of search matches to create models from 
                NUM_MR     = Number of models to use in MR step (less or equal to NUM_MODEL)
                NUM_ENSEM  = Number of models to be used in Ensemble construction (less or equal to NUM_MR)
                MR_PROGRAM = Molecular Replacement program to use (PHASER, MOLREP or AMORE) 
                E_value    = E value for OCA search (higher = more flexible = more hits) 

                ------------
                Input Files:
                ------------
 
                MTZ file - relection file. Should have a FreeR_flag column.
                SEQ file - fasta format file containing target sequence.
                IGNORE file - file containing PDB codes to be ignored in the search for similar 
                              structures. Each code should entered on a single line.
               
           """

   print useStr 

# If correct, initialise variables

# mtz_coldata is an array to hold the mtz column labels
mtz_coldata=[]

if len(sys.argv) == 11 or len(sys.argv) == 12:

   if sys.argv[1] == '-p':
      dir=os.getcwd()
      mtzdata=sys.argv[2]
      mtz_filename=string.split(mtzdata,'/')[-1]

      NUM_MODEL=int(sys.argv[3])
      NUM_MR=int(sys.argv[4])
      NUM_ENSEM=int(sys.argv[5])
      MR_PROGRAM=sys.argv[6]

      seq_file=sys.argv[7]
      seq_filename=string.split(seq_file,'/')[-1]

      job_id = sys.argv[8] 
      E_value = float(sys.argv[9]) 

      ignore_file = sys.argv[10] 
      ignore_filename = string.split(ignore_file,'/')[-1] 

      mrsearchdir = dir + '/search_' + job_id      
      MODEL_ONLY=False

      if len(sys.argv) == 12:
         results_dir = sys.argv[11]
      else:
         results_dir = mrsearchdir + '/results'

      if NUM_MR > NUM_MODEL:
         print "NUM_MR (%d) should be <= NUM_MODEL (%d)" % (NUM_MR, NUM_MODEL)
	 sys.exit(1)
      if NUM_ENSEM > NUM_MR:
         print "NUM_ENSEM (%d) should be <= NUM_MR (%d)" % (NUM_ENSEM, NUM_MR)
	 sys.exit(1)

      if MR_PROGRAM=='Phaser' or MR_PROGRAM=='phaser':
         MR_PROGRAM='PHASER'
      if MR_PROGRAM=='Molrep' or MR_PROGRAM=='molrep':
         MR_PROGRAM='MOLREP'
      if MR_PROGRAM=='Amore' or MR_PROGRAM=='amore':
         MR_PROGRAM='AMORE'

      if MR_PROGRAM!='PHASER' and MR_PROGRAM!='MOLREP' and MR_PROGRAM!='AMORE':
         print MR_PROGRAM
         print "MR_PROGRAM value must be 'PHASER/Phaser/phaser' or 'MOLREP/Molrep/molrep' or 'AMORE/Amore/amore'" 
	 sys.exit(1)
      
      print ("\nEnter MTZ column data(F, SIGF, FREE): ")
      mtz_coldata.append(raw_input("F: "))
      mtz_coldata.append(raw_input("SIGF: "))
      mtz_coldata.append(raw_input("FREE: "))

   else:
      usage()
      sys.exit(1)

# If we are just running a modelling job then we only need four inputs

elif len(sys.argv) == 7 or len(sys.argv) == 8:

   if sys.argv[1] == '-m':
      NUM_MODEL=int(sys.argv[2])
      NUM_MR=0
      dir=os.getcwd()
      mtzdata='/tmp/null.mtz'

      seq_file=sys.argv[3]
      seq_filename=string.split(seq_file,'/')[-1]

      job_id = sys.argv[4] 
      E_value = float(sys.argv[5]) 

      ignore_file = sys.argv[6] 
      ignore_filename = string.split(ignore_file,'/')[-1] 

      mrsearchdir = dir + '/search_' + job_id
      MODEL_ONLY=True

      if len(sys.argv) == 8:
         results_dir = sys.argv[7]
      else:
         results_dir = mrsearchdir + '/results'

   else:
      usage()
      sys.exit(1)

# Otherwise, echo the usage and exit

else:
   usage()
   sys.exit(1)

# Setup the directory structure

if os.access(dir + '/search_' + job_id, os.F_OK):
   print "\n"
   print "A directory for a job with this name already exists:\n " +  dir + "/search_" + job_id
   print "Please delete it or choose a different job name"
   sys.exit(0)

os.mkdir(dir + '/search_' + job_id)
os.mkdir(mrsearchdir + '/input')
os.mkdir(mrsearchdir + '/logs')
os.mkdir(mrsearchdir + '/data')
os.mkdir(mrsearchdir + '/sequences')

if not os.access(results_dir, os.F_OK): 
   os.mkdir(results_dir)
os.mkdir(results_dir + '/images')
os.mkdir(results_dir + '/data')
os.mkdir(results_dir + '/data/sequences')
os.mkdir(results_dir + '/pdb_files')
os.mkdir(results_dir + '/log_files')
os.mkdir(results_dir + '/marginal_solns')

# Copy in the input files to the search directory - .seq, .mtz and ignore.txt (and the col data file if needed).
# After copying, reassign the file pointers.

shutil.copy(seq_file, mrsearchdir + '/input/')
shutil.copy(ignore_file, mrsearchdir + '/input/')
seq_file=mrsearchdir + '/input/' + seq_filename
ignore_file = mrsearchdir + '/input/' + ignore_filename 
if MODEL_ONLY==False:
   shutil.copy(mtzdata, mrsearchdir + '/input/')
   mtzdata = mrsearchdir + '/input/' + mtz_filename 

# Copy over header for results page

shutil.copyfile(bmp_incl + '/web_template/ehtpx-logo.jpg',results_dir + '/images/ehtpx-logo.jpg')
shutil.copyfile(bmp_incl + '/web_template/cclrc-logo.jpg', results_dir + '/images/cclrc-logo.jpg')
shutil.copyfile(bmp_incl + '/web_template/slide_bar.gif', results_dir + '/images/slide_bar.gif')
shutil.copyfile(bmp_incl + '/web_template/results.html', results_dir + '/results.html')

# Check that the DB's are up to date. Only do this if a local search is not required.
# If a local search is set then the SSM search cannot be carried out.

print "\n-------------------"
print "Database Check:"
print "-------------------\n"

db_tools=Update_DB.DB_tools()
if BMP_RUN_LOCAL == False:
   db_tools.test_first_run(BMP_SETUP_DIR)
   db_tools.test_databases(BMP_SETUP_DIR)
else:
   db_tools.test_first_run(BMP_SETUP_DIR)
   BMP_SSM = False
   BMP_PQS = False
   BMP_FASTA_LOCAL = False

# Create an instance of mstat
mstat=Matches.Match_struct()
mstat.setSearchDir(mrsearchdir)
mstat.setResultsDir(results_dir)
mstat.setScratchDir(scratch_dir)

# Target sequence information

print "\n-------------------"
print "Target information:"
print "-------------------\n"

target_info=Target_info.TargetInfo()
mstat.InitError = target_info.setTargetInfo(seq_file, mtzdata, mrsearchdir, mtz_coldata, job_id, MODEL_ONLY)
target_info.printTargetInfo()

# Set the DB's for local-based searches
mstat.setFastaDB(bmp + '/data/pdb_seqres.txt')

# Set the output results html page
mstat.setResultsHTML(mstat.results_dir + "/results.html")

# Make a note of the start time for the job
mstat.setStartTime(time.time())

# Set the model preparation and MR parameters
mstat.setNUM_MODEL(NUM_MODEL)
if MODEL_ONLY == False:
   mstat.setNUM_MR(NUM_MR)
   mstat.setNUM_ENSEM(NUM_ENSEM)
   mstat.setMR_PROGRAM(MR_PROGRAM)

# initialize the results page
#w=Write_Search_results.Write_html()
mstat.search_results.write_header(mstat.results_html)
mstat.search_results.write_footer(mstat, mstat.results_html)

# write the target details
mstat.search_results.write_target_details(mstat, target_info, mstat.results_html)

# If there was an error in the initialisation exit the program
if mstat.job_end == True:
   sys.exit()

# Search information

print "\n-----------"
print "Fasta search:"
print "-----------\n"

seq_search=Matches.SEQ_match()

seq_search.get_Fasta_file(mstat, target_info, E_value, results_dir + '/data/matches', BMP_SETUP_DIR, BMP_FASTA_LOCAL)

# Read the development options

dev=Developer.Developer()
dev.ignore_PDBcodes(mstat, ignore_file)

seq_search.get_Fasta_matches(mstat, BMP_FASTA_LOCAL)

# Exit and finish html if there were no matches
if mstat.no_of_OCAhits == 0:
   mstat.search_results.write_SEQ_search_details(mstat, mstat.results_html)
   sys.exit()

ps=Matches.Process_Score()

print "\n----------------------------------------"
print "Fasta: Downloading & Processing PDB files:"
print "----------------------------------------\n"

ps.PDB_process(mstat, 'OCA')
ps.alignment(target_info, mstat, BMP_MA_PROGRAM)
ps.get_top_hit(mstat)

# write the OCA search details to the html page
mstat.search_results.write_SEQ_search_details(mstat, mstat.results_html)

#shutil.copy('/tmp/ssm_result.xml', mrsearchdir + '/results/ssm_result.xml')

if BMP_SSM == True:

   print "\n-----------"
   print "SSM search:"
   print "-----------\n"

   ssm_search=Matches.SSM_match()

   ssm_search.get_SSM_file(mstat, job_id, mstat.results_dir + '/data/ssm_result.xml')
   ssm_search.get_SSM_matches(mstat)

   print "\n----------------------------------------"
   print "SSM: Downloading & Processing PDB files:"
   print "----------------------------------------\n"

   ps.PDB_process(mstat, 'SSM')
   ps.alignment(target_info, mstat, BMP_MA_PROGRAM)
   ps.get_top_hit(mstat)

   # write the SSM search details to the html page
   mstat.search_results.write_SSM_search_details(mstat, mstat.results_html)
 
# Search the SCOP database to see if any domains exist in the template structures

if BMP_SCOP == True:

   print "\n-----------------------------"
   print "Searching for Domains (SCOP):"
   print "-----------------------------\n"

   scop_search=Matches.SCOP_match()
   scop_search.get_SCOP_matches(mstat, BMP_SETUP_DIR)
   
   print "\n---------------------------"
   print "SCOP: Processing PDB files:"
   print "---------------------------\n"

   ps.PDB_process(mstat, 'SCOP')
   ps.alignment(target_info, mstat, BMP_MA_PROGRAM)
   ps.get_top_hit(mstat)

   # write the SSM search details to the html page
   mstat.search_results.write_domains_search_details(mstat, mstat.results_html)
 
# Prune the list
ps.prune_list(mstat)
ps.print_list(mstat)

# write the overall search details to the html page
mstat.search_results.write_ALL_search_details(mstat, target_info, mstat.results_html)
 
# Do the modelling of the PDB model structures

print "\n--------------------"
print "Template Preperation:"
print "--------------------\n"

model=Model.Modelling()
model.MAFFT_pir(mstat, target_info)
model.molrep(target_info, mstat)

# Do the pruning of the template structures using Chainsaw 
# Only use pairwise if set to True
if BMP_PW_ALIGN:
   model.chainsaw(target_info, mstat, 'pairwise')

# Multiple alignment input to Chainsaw
model.chainsaw(target_info, mstat, 'multiple')

# Monitor the queue for the completion of all the jobs

if CLUSTER == True:
   queue=Model_Queue.Queue_info()
   queue.qstat_check(mstat, target_info, 'MODEL_PREP')
   mstat.resetJobid_Array()
   mstat.resetJobid_Dict()

# write the overall search details to the html page
mstat.search_results.write_model_prep_details(mstat, target_info, mstat.results_html)

if MODEL_ONLY == False:
# Setup the models for MR

   print "\n-------------------------------------"
   print "Molecular Replacement and Refinement:"
   print "-------------------------------------\n"

   mr_setup=Model_struct.MR_setup()
   mr_setup.setTypes(mstat, mrsearchdir)

   # Write the header for the MR secion
   mstat.mr_results.write_MR_header(mstat)

# If the number of molecules in our target is greater than 1 consult PQS for multimers

   if target_info.no_of_mols >= 2:

      print "\n------------------------------"
      print "Searching for Multimers (PQS):"
      print "------------------------------\n"

      pqs=PQS.PQS()

      pqs.set_pqs_list(BMP_SETUP_DIR)
      pqs.get_pqs_files(target_info, mstat, mrsearchdir)

      # write the PQS search details to the html page
      mstat.search_results.write_PQS_search_details(mstat, target_info, mstat.results_html)
 
# Do the molecular replacement using the model list
 
   MR=BMP_MR.MR_submit()
   MR.MR_setup(target_info, mstat, mrsearchdir)
   MR.MR_start(target_info, mstat, mrsearchdir)

# Poll the queue to see if the jobs have finished
   if CLUSTER == True:

      print "\n---------------------------"
      print "Results (post-refinement):"
      print "---------------------------\n"
        
      queue.qstat_check(mstat, target_info, 'MOL_REPLACE')
      mstat.resetJobid_Array()
      mstat.resetJobid_Dict()

# If Amore failed to produce a solution we will try Molrep on each of the models

   if mstat.solution_found == False and mstat.MR_PROGRAM == 'AMORE':

      mstat.setMR_PROGRAM('MOLREP')

      MR.MR_start(target_info, mstat, mrsearchdir)

      if CLUSTER == True:

         print "\n---------------------------"
         print "Results (post-refinement):"
         print "---------------------------\n"
        
         queue.qstat_check(mstat, target_info, 'MOL_REPLACE')
         mstat.resetJobid_Array()
         mstat.resetJobid_Dict()

# If Molrep failed to produce a solution we will try Phaser on each of the models

   if mstat.solution_found == False and mstat.MR_PROGRAM == 'MOLREP':

      mstat.setMR_PROGRAM('PHASER')

      print "\n------------------------------"
      print "Setting up ensemble for Phaser:"
      print "------------------------------\n"

      ens=Ensemble.Ensemble()
      ens.setup_ensemble(mstat, mrsearchdir, NUM_ENSEM)

      # write the ensemble details to the html page
      mstat.search_results.write_ensemble_details(mstat, target_info, mstat.results_html)
   
      MR.MR_start(target_info, mstat, mrsearchdir)

      if CLUSTER == True:

         print "\n---------------------------"
         print "Results (post-refinement):"
         print "---------------------------\n"
        
         queue.qstat_check(mstat, target_info, 'MOL_REPLACE')
         mstat.resetJobid_Array()
         mstat.resetJobid_Dict()

   # write the ensemble details to the html page
   mstat.mr_results.write_final_soln(mstat, target_info)
        
   finish_time = time.ctime()
   finish_file = open(mstat.results_dir + '/finished.txt', 'w')
   finish_file.write('Job finished: ' + finish_time)
   finish_file.write('')
   finish_file.write('Information only, please do not reply to this email.')
   finish_file.close()

   # Remove the data directory to save memory space
   if BMP_DEBUG == False:
      print "Cleanup: Removing search data directory"
      cleanup=Cleanup.Clean()
      cleanup.rm_dir(mstat.search_dir + '/data')

sys.exit(0)

