#! /usr/bin/env python

# setup the environment

import os, sys, string, time

# 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'

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

bmp = os.environ['BMP']
bmp_incl = os.environ['BMP_INCL']
CLUSTER=eval(os.environ['BMP_CLUSTER'])

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

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

import Matches, Target_info, Model
import Model_Queue, Results, Struct_dump
import Model_struct, Ensemble, PQS
import BMP_Phaser, BMP_Molrep, BMP_Refmac
import Developer

# 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 <target MTZ file> <NUM_MODEL> <NUM_MR> <NUM_ENSEM> <MR_PROGRAM>]
                   ...<target sequence file> <job_id> <E_value> <ignore PDB file>

                -m <NUM_MODEL>                    
                    --Model only (don't do molecular replacement)
                -p <target 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) 
           """

   print useStr 

# If correct, initialise variables

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

if len(sys.argv) == 11:

   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 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!='PHASER' and MR_PROGRAM!='MOLREP':
         print MR_PROGRAM
         print "MR_PROGRAM value must be 'PHASER/Phaser/phaser' or 'MOLREP/Molrep/molrep' (AMORE to be added later)" 
	 sys.exit(1)
      
      print ("Enter 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:

   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

   else:
      usage()
      sys.exit(1)

# Otherwise, echo the usage and exit

else:
   usage()
   sys.exit(1)

#sys.exit(-1)

# Setup the directory structure

os.system('mkdir -p ' + dir + '/search_' + job_id)
os.system('mkdir -p ' + mrsearchdir + '/input')
os.system('mkdir -p ' + mrsearchdir + '/logs')
os.system('mkdir -p ' + mrsearchdir + '/data')
os.system('mkdir -p ' + mrsearchdir + '/models')
os.system('mkdir -p ' + mrsearchdir + '/pdb_files')
os.system('mkdir -p ' + mrsearchdir + '/results')
os.system('mkdir -p ' + mrsearchdir + '/results/images')
os.system('mkdir -p ' + mrsearchdir + '/results/data')
os.system('mkdir -p ' + mrsearchdir + '/results/pdb_files')
os.system('mkdir -p ' + mrsearchdir + '/results/log_files')

# 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.

os.system('cp ' + seq_file + ' ' + mrsearchdir + '/input/')
os.system('cp ' + ignore_file + ' ' + mrsearchdir + '/input/')
seq_file=mrsearchdir + '/input/' + seq_filename
ignore_file = mrsearchdir + '/input/' + ignore_filename 
if MODEL_ONLY==False:
   os.system('cp ' + mtzdata + ' ' + mrsearchdir + '/input/')
   mtzdata = mrsearchdir + '/input/' + mtz_filename 

# Copy over header for results page

os.popen('cp ' + bmp_incl + '/web_template/*.jpg ' + mrsearchdir + '/results/images', 'r')

# Target sequence information

target_info=Target_info.TargetInfo()
target_info.setTargetInfo(seq_file, mtzdata, mrsearchdir, mtz_coldata, MODEL_ONLY)

# Search information

search=Matches.Read_Matches()

mstat=search.get_OCAmatches(target_info, mrsearchdir + '/results/data/matches.html')
mstat.setNUM_MODEL(NUM_MODEL)
mstat.setE_value(E_value)
if MODEL_ONLY == False:
   mstat.setNUM_MR(NUM_MR)
   mstat.setNUM_ENSEM(NUM_ENSEM)
   mstat.setMR_PROGRAM(MR_PROGRAM)

# Read the development options

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

search.get_seq_list(mstat)
search.get_seq_id(mstat)

search.get_match_stats(mstat)
if mstat.no_of_OCAhits == 0:
  print "No matches found using the OCA database. Try a higher e-value"
  sys.exit(0)

#os.system('cp /tmp/ssm_result.xml ' + mrsearchdir + '/results/ssm_result.xml')

search.get_ssm_matches(mstat, job_id, mrsearchdir + '/results/data/ssm_result.xml')
search.read_ssm_results(mstat)
search.sort_list(mstat, 'seqID', NUM_MODEL, NUM_MR)

# PDB information
if mr == False:
   pdbs=Matches.Get_PDB()
   pdbs.download_pdbs(mstat, mrsearchdir)
   pdbs.getPDBinfo(mstat, mrsearchdir)

# Do the modelling of the PDB model structures

   model=Model.Modelling()
   model.Clustalw_align(mstat, target_info, mrsearchdir)
   #model.modeller(target_info, mstat, mrsearchdir)
   model.molrep(target_info, mstat, mrsearchdir)
   model.chainsaw(target_info, mstat, mrsearchdir)

# Monitor the queue for the completion of all the jobs

   queue=Model_Queue.Queue_info()
   queue.qstat_check(mstat.jobid_array)
   mstat.resetJobid_Array()

# Output the results

   results=Results.Model_Results()
   results.dump_results(target_info, mstat, mrsearchdir)
   results.write_html(target_info, mstat, mrsearchdir, MODEL_ONLY)

# Dump the structure contents

   dump=Struct_dump.Dump_all()
   dump.write_out_models(mstat, mrsearchdir + '/results/data/mstat.dat')
   dump.write_out_target(target_info, mrsearchdir + '/results/data/target.dat')

#else:
   dump=Struct_dump.Dump_all()
   dump.read_in_models(mstat, mrsearchdir + '/results/data/mstat.dat')
   dump.read_in_target(target_info, mrsearchdir + '/results/data/target.dat')

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

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

# Setup an ensemble of the best results for Phaser

   if MR_PROGRAM=='PHASER':
      ens=Ensemble.Ensemble()
      ens.setup_ensemble(mstat, mrsearchdir, NUM_ENSEM)

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

   #target_info.no_of_mols = 4
   if target_info.no_of_mols >= 2:
      pqs=PQS.PQS()

      pqs.get_pqs_list()
      pqs.get_pqs_files(target_info, mstat, mrsearchdir)

# Do the molecular replacement using the model list

   print "BMP log message: Starting MR program for list of Matches"

   if MR_PROGRAM=='PHASER':
      phaser=BMP_Phaser.Phaser_submit()
      phaser.Phaser_setup(target_info, mstat, mrsearchdir)
      phaser.Phaser_start(target_info, mstat, mrsearchdir)

   if MR_PROGRAM=='MOLREP':
      molrep=BMP_Molrep.Molrep_submit()
      molrep.Molrep_setup(target_info, mstat, mrsearchdir)
      molrep.Molrep_start(target_info, mstat, mrsearchdir)

# Poll the queue to see if the jobs have finished
   queue.qstat_check(mstat.jobid_array)
   mstat.resetJobid_Array()

   mr_results=Results.MR_results()
   mr_results.collect_MR_scores(mstat, mrsearchdir)
   mr_results.sort_scores(mstat, MR_PROGRAM)
   mr_results.plot_scores(mstat, mrsearchdir, MR_PROGRAM)

# Refine the resulting models and collect Rfree values

   print "BMP log message: Starting refinement program for list of Matches"
   refmac=BMP_Refmac.Refmac_submit()
   refmac.Refmac_setup(target_info, mstat, mrsearchdir)
   refmac.Refmac_start(target_info, mstat, mrsearchdir)

# Poll the queue to see if the jobs have finished
   queue.qstat_check(mstat.jobid_array)
   mstat.resetJobid_Array()

   mr_results.collect_refmac_scores(mstat, mrsearchdir)
   mr_results.sort_scores(mstat, 'REFMAC')
   mr_results.plot_scores(mstat, mrsearchdir, 'REFMAC')
   mr_results.write_html(target_info, mstat, mrsearchdir, MODEL_ONLY)

# construct the results webpage

finish_time=time.ctime()
os.system('echo "Job finished: ' + finish_time + '" > ' + mrsearchdir + '/results/finished.txt')

sys.exit(0)
