#! /usr/bin/python -u
# -*- coding: utf-8 -*-
#

import sys, os, datetime
from optparse import OptionParser

try:
    import numpy # Just to test if we have numpy installed.
    from modeller import *
    from modeller.scripts import complete_pdb
    modeller_installed = True
except ImportError:
    modeller_installed = False

import sphinx

# Location of Sphinx package
scriptpath = os.path.dirname(sphinx.__file__)

VERSION = "Sphinx - (build "+str(datetime.date.fromtimestamp(os.path.getmtime(sys.argv[0])))+")"
longVersion = """
                              
################################################################################
#                                                                              #
#                                    SPHINX                                    #
#                                                                              #
#                       Hybrid Loop Structure Prediction                       #
#                                                                              #
################################################################################


                For feedback, bug reports and questions contact:                 
                     Claire Marks <marks@stats.ox.ac.uk>                                            

""".format(version=VERSION)


USAGE="""
Sphinx [OPTIONS] <db_dir> <pdb_file> <start_residue> <end_residue> <chain> 
                 <loop_sequence> <results_dir>

<db_dir>
    Root directory of a fragment database

<pdb_file>
    Query PDB file with or without the missing loop. If the loop is present, 
    its coordinates are ignored, except for reporting each decoy's RMSD to
    this structure.

<start_residue>
    The number of the residue BEFORE the beginning of the loop. This residue
    is always present in the input model. 
    
<end_residue>
    The number of the LAST residue of the loop. This does not have to be 
    present in the input model.
    
<chain>
    The chain on which the loop to be modelled is found.

<loop_sequence>
    Literally, the amino acid sequence of the loop, e.g. "LTPVA"
    
<results_dir>
    The location of the folder in which Sphinx should place its output.


"""

outputHelp = """

################################################################################
                                     OUTPUT                                   

Sphinx outputs the following:
    
    all_decoy_structures.pdb - literally, all the decoy structure that were
                generated by Sphinx during the modelling process
                
    logfile.txt - a file containing the job information and details of each
                fragment used (alignment, ESS score etc.)
                
    ranking_RAPDF.csv - the scores for all the decoys generated, calculated
                using Sphinx's inbuilt ranking function
                
    top_500_decoy_structures.pdb - the structures of the top 500 decoys, as
                ranked using Sphinx's inbuilt ranking function
                
    CompleteModels - a folder containing the full PDB files for the top 500 
                decoys
                
    ranking_final.csv - a ranking of the top 500 decoys.
                

################################################################################
"""


def main(dbdir, pdb_file, start_of_loop, end_of_loop, chain, loop_sequence, results_directory, ignore_list = [], no_fragments = -1, sphinx_version = "general", contact_file = None, ranking_method = "soaploop", sidechain_method="pears"):
    s_file = ""
    try:
        #KK: We need to indicate here if we crashed or not
        status_file = os.path.join(results_directory, "status.txt")
        logfile = os.path.join(results_directory, "logfile.txt")
        if os.path.exists(results_directory) is False:
          os.makedirs(results_directory)
        print("Starting Sphinx")
        loop_sequence = loop_sequence.strip().upper()
        
        frd = Fread()
        
        frd.open_errstream("-")
        
        frd.set_subst_tables(os.path.join(scriptpath, 'dat/esst.txt'))
        
        frd.set_structure(pdb_file)
        
        results_file = logfile
        decoy_directory = results_directory
        
        if opts.ignore_list != "":
            ignore_list = [pdbcode.lower() for pdbcode in opts.ignore_list.split(":")]
        
        if opts.no_fragments == -1:
            if len(loop_sequence)%2 == 0:
                if sphinx_version == "general":
                    no_fragments = int((len(loop_sequence)-2) * 25 / 2)
                else:
                    no_fragments = (len(loop_sequence)-2) * 25
            else:
                if sphinx_version == "general":
                    no_fragments = int((len(loop_sequence)-1) * 25 / 2)
                else:
                    no_fragments = (len(loop_sequence)-1) * 25

        if contact_file:
            if os.path.exists(contact_file) is False:
                print("Contact file given does not exist!")
                sys.exit(1)                

        if ranking_method == "soaploop" and modeller_installed:
            print("Ranking method will be SOAP-Loop.")
        elif ranking_method == "soaploop" and modeller_installed is False:
            print("MODELLER not installed; ranking method will be dDFIRE.")
            ranking_method = "ddfire"
        elif ranking_method == "ddfire":
            print("Ranking method will be dDFIRE.")
        elif ranking_method == "rwplus":
            print("Ranking method will be RWplus.")

        if sidechain_method == "pears":
            print("Sidechains will be modelled using PEARS.")
        elif sidechain_method == "scwrl4":
            print("Sidechains will be modelled using SCWRL4.")


        # Generate decoys
        frd.automodel_loop(start_of_loop, 
            end_of_loop, 
            loop_sequence, 
            pdb_file, 
            results_file, 
            decoy_directory,
            scriptpath,
            no_fragments,
            resnum = True,
            chain = chain,
            dbdir = dbdir,
            ignore_list = ignore_list,
            sphinx_version = sphinx_version,
            ranking_method = ranking_method,
            sidechain_method = sidechain_method)
           
 
        # Run Ranking to get 500 best decoys
        print("Selecting top 500 decoys... (using %s data)" %(countsfile.split("/")[-1].split("_")[0]))

        #frd.decoy_file = os.path.join(decoy_directory, "all_decoy_structures.pdb")
        #frd.actual_loop_start = int(start_of_loop)+1

        print("Running RAPDF...")
        RAPDF_output = os.path.join(decoy_directory, "ranking_RAPDF.csv")        
        os.system("%s -i %s -o %s -pdb %s -chain %s -maxdist 8 -start %s -end %s -countsfile %s"
                %(RAPDF_exec, frd.decoy_file, RAPDF_output, pdb_file, chain, frd.actual_loop_start, end_of_loop, countsfile))
                
        if contact_file:
            print("Running contact ranking...")
            contact_output = os.path.join(decoy_directory, "ranking_contacts.csv")
            success = rank_using_contacts(contact_output, contact_file, frd.decoy_file, pdb_file, chain, frd.actual_loop_start, end_of_loop)
            if success:
                top500 = select_top_500_with_contacts(RAPDF_output, contact_output)
            else:
                top500 = select_top_500(RAPDF_output)
        else:
            top500 = select_top_500(RAPDF_output) 
       
        #if modeller_installed:
        #    print("\nRanking top 500 decoys with SOAP-Loop...") 
        #else:
        #    print("\nMODELLER is not installed. Decoys will not be ranked further. Writing full PDB files for each decoy...")
        
        print("\nRanking top 500 decoys...")
        write_rank_top500(chain, frd.actual_loop_start, end_of_loop, top500, frd.decoy_file, pdb_file, scriptpath, loop_sequence, dbdir, ranking_method=ranking_method, sidechain_method=sidechain_method)
        
        print("\nFinished!!")
        
        frd.note("Finished")
        
        frd.close_errstream()
        #Compress the tpo 500 results
        os.system("tar -zcf %s/CompleteModels.tar.gz %s/CompleteModels/* >/dev/null 2>&1"%(results_directory,results_directory))
        #Indicate that we have hereby finished.
        s_file = open(status_file,'w')
        s_file.write('Done')
        s_file.close()
    except:
        raise
        s_file = open(status_file,'w')
        s_file.write('Crashed')
        s_file.close()

parser = OptionParser(version=VERSION, usage=USAGE)
	    
parser.add_option("--copyright", dest="copyright", action="store_true",
                  help="Display copyright message and exit.")

parser.add_option("--output_help", dest="output_help", action="store_true",
                  help="Display a description of the program's output, then exit.")
                  
parser.add_option("--sphinx_version", dest = "sphinx_version", choices=["general", "h3", "membrane"],
                  help="Which version of Sphinx to use - each version uses data specific to that protein type", 
                  default="general")

parser.add_option("--ignore", dest="ignore_list", type="string", default = "",
                  help="A list of PDB codes separated by colons; fragments from these will be ignored.")
                  
parser.add_option("--no_fragments", dest="no_fragments", type="int", default=-1,
                  help="Override the default number of fragments from the database to be used in loop building. More may be used if multiple fragments give the same ESSS score.")

parser.add_option("--contact_file", dest="contact_file", type="str", default=None,
                  help="Path to a file containing predicted contact information that Sphinx should use when ranking decoys. This file should be in csv format, with each line containing the information for one contact (residue number 1, residue number 2, score). All contacts provided will be used (so if only contacts with a particular score are required, only include these!) and will be assumed to be on the same chain as the target loop.")

parser.add_option("--ranking_method", dest="ranking_method", default="soaploop", choices=["soaploop", "ddfire", "rwplus"],
                  help="Specify the ranking method that Sphinx should use to score the top 500 decoys. Choices are soaploop (default), ddfire, or rwplus. SOAP-Loop can only be used when MODELLER is installed.")

parser.add_option("--sidechain_method", dest="sidechain_method", default="pears", choices=["pears", "scwrl4"],
                  help="Specify the algorithm Sphinx should use to model the loop sidechains. Choices are PEARS (default) or SCWRL4 (must be in your PATH).")

(opts, args) = parser.parse_args()

if opts.copyright:
    print(longVersion)
    sys.exit(1)
    
if opts.output_help:
    print(outputHelp)
    sys.exit(1)

if len(args) != 7:
    print(longVersion)
    parser.print_help()
    print(outputHelp)
    sys.exit(1)

# Model H3-specific
#if opts.sphinx_version != "general":
#    from sphinx.fragments_h3.fread import Fread
#    from sphinx.rapdf_h3.ranking_functions import *
#    RAPDF_exec   = "RAPDF_H3"
#    countsfile   = os.path.join(scriptpath, "dat/H3_counts_matrix.csv")
#else:
from sphinx.fragments.fread import Fread
from sphinx.ranking.ranking_functions import *
from sphinx.contacts.contact_functions import rank_using_contacts, select_top_500_with_contacts

RAPDF_exec   = os.path.join(scriptpath, "ranking", "ranking_execs", "RAPDF")
if opts.sphinx_version == "h3":
    countsfile = os.path.join(scriptpath, "dat/h3_counts_matrix.csv")
elif opts.sphinx_version == "membrane":
    countsfile = os.path.join(scriptpath, "dat/membrane_counts_matrix.csv")
else:
    countsfile = os.path.join(scriptpath, "dat/general_counts_matrix.csv")

dbdir, pdb_file, start_of_loop, end_of_loop, chain, loop_sequence, results_directory = args

results_directory = os.path.abspath(results_directory)

main(dbdir, pdb_file, start_of_loop, end_of_loop, chain, loop_sequence, results_directory, opts.ignore_list, opts.no_fragments, sphinx_version = opts.sphinx_version, contact_file = opts.contact_file, ranking_method=opts.ranking_method, sidechain_method=opts.sidechain_method)
