#!/usr/bin/env python

import os, sys, optparse, getpass
from multiprocessing import Process, Pipe
from Utils import Base, Fasta

#################################################
def options():
    parser = optparse.OptionParser('usage: %prog -i "proteins1.fa proteins2.fa ... proteinsN.fa" -l "lab1 lab2 ... labN" -p "1 3 ... 1" -e 1e-5 -s 0.6')
    parser.add_option('-i', '--filenames', dest='filenames', help='Names of the files of species containing the proteins', metavar='FILES', default='')
    parser.add_option('-l', '--labels', dest='labs', help="Labels for each species", metavar='LABELS', default='')
    parser.add_option('-p', '--positions', dest='positions', help="Default positions of unique identifier in FASTA header separated by |. Default position is 1 for all.", metavar='POSITIONS', default='')
    parser.add_option('-e', '--evalue', dest='evalue', help="E-value used at blast. Default is 1e-5. Use 1e-X format only!", metavar='EVALUE', default='1e-5')
    parser.add_option('-s', '--similarity', dest='sim', help="Required similarity (0 .. 1). Default if 0.5", metavar='SIM', default='0.5')
    parser.add_option('-m', '--minlen', dest='minlen', help="Allowed minimum lenght of a protein. Default is 20.", metavar='MINLEN', default='20')
    parser.add_option('-b', '--noblast', dest='skipBlast', action='store_true', help="Skip BLAST", default=False)
    parser.add_option('-a', '--add', dest='add', action='store_true', help="Adds given species and labels to existing BLAST", default=False)
    parser.add_option('-u', '--uname', dest='uname', help='Username', metavar='Unix username at the remove server.', default='')
    options, args = parser.parse_args()
    if options.filenames == '' or options.labs == '':
        parser.print_help()
        print '\nE.g.: orthoMcl -i "proteome1.fa proteome2.fa" -l "Tax Tvi" -p "4 4" -e 1e-5 -s 0.5'
        print "Results will be found in TmpOrthoMcl directory in groups.txt file."
        print "Note! The labels must be exactly 3 characters long and preferrably start with an upper case character."
        sys.exit(1)
    return options


#################################################
def checkResidue(fastaFile):
    '''
    '''
    retVal = "nucleotides"
    try:
        limit = 100
        fasta = Fasta(fastaFile)
        for i in xrange(len(fasta.headers)):
            if i > limit: break
            seq = fasta.seqs[i].upper()
            for item in seq:
                if item not in ['A', 'T', 'C', 'G', 'N']:
                    retVal = "amino acids"
                    break
    except IOError:
        print "Fatal error: file %s not found. Exiting..." %fastaFile
        sys.exit()
    return retVal


#################################################
def checkUniqueIds(fastaFile):
    '''
    '''
    fasta = Fasta(fastaFile)
    if len(fasta.headers) != len(set(fasta.headers)):
        print "Fatal error: FASTA sequence identifiers are not unique in %s. Exiting..." %fastaFile
        print "Probably position for this file is given wrong..."
        sys.exit()


#################################################
def createOrthoMclConfigFile(wd, userName, eValue, similarity):
    '''
    '''
    eValue = eValue.split('e')[1]
    similarity = int(float(similarity) * 100.0)
    handle = open("%s/orthomcl.config" %wd, 'w')
    handle.write("# this config assumes a mysql database named 'orthomcl'.  adjust according\n")
    handle.write("# to your situation.\n")
    handle.write("dbVendor=mysql\n")
    handle.write("dbConnectString=dbi:mysql:ortho%s\n" %userName)
    handle.write("dbLogin=ortho%s\n" %userName)
    handle.write("dbPassword=password\n")
    handle.write("similarSequencesTable=SimilarSequences\n")
    handle.write("orthologTable=Ortholog\n")
    handle.write("inParalogTable=InParalog\n")
    handle.write("coOrthologTable=CoOrtholog\n")
    handle.write("interTaxonMatchView=InterTaxonMatch\n")
    handle.write("percentMatchCutoff=%d\n" %similarity)
    handle.write("evalueExponentCutoff=%s\n" %eValue)
    handle.write("oracleIndexTblSpc=NONE\n")
    handle.close()


#################################################
def createMySqlScripts(wd, userName):
    '''
    '''
    handle = open("%s/createDb.sql" %wd, 'w')
    handle.write("CREATE DATABASE ortho%s;\n" %userName)
    handle.write("GRANT SELECT,INSERT,UPDATE,DELETE,CREATE VIEW,CREATE, INDEX, DROP on ortho%s.* TO ortho%s@localhost;\n" %(userName, userName))
    handle.write("set password for ortho%s@localhost = password('password');\n" %userName)
    handle.close()
    handle = open("%s/dropDb.sql" %wd, 'w')
    handle.write("drop database if exists ortho%s;\n" %userName)
    handle.close()


#################################################
def callShell(base, cmdStr, dummy = None):
    '''
    '''
    base.shell(cmdStr)

#################################################
def main():
    '''
    '''
    opts = options() # files contains exactly two PE files

    eValue = opts.evalue
    similarity = opts.sim
    minlen = opts.minlen
    files = opts.filenames.split()
    labels = opts.labs.split()
    if len(labels) != len(set(labels)):
        print "Fatal error: duplicate labels found. Exiting..."
        sys.exit(0)
    if len(files) != len(set(files)):
        print "Fatal error: duplicate fasta file names found. Exiting..."
        sys.exit(0)
    positions = None
    if opts.positions != "":
        positions = opts.positions.split()
    if positions == None:
        positions = []
        for i in xrange(len(files)):
            positions.append("1")
    if len(files) != len(labels):
        print "Fatal error: number of files does not match with the number of labels. Exiting..."
        sys.exit(0)
    if len(positions) != len(labels):
        print "Fatal error: number of labels does not match with the number of positions of the ids. Exiting..."
        sys.exit(0)
    for lab in labels:
        if len(lab) != 3:
            print "Fatal error: labels have to be exactly three characters long. Exiting..."
            sys.exit(0)

    base = Base()
    wd = "TmpOrthoMcl"
    wdFasta = "%s/Fasta" %wd
    wdAdds = "%s/Adds" %wd
    base.createDir(wd)
    logHandle = open("%s/log.txt" %wd, 'w')
    base.setHandle(logHandle)
    base.createDir(wdFasta)
    userName = getpass.getuser()
    createOrthoMclConfigFile(wd, userName, eValue, similarity)
    createMySqlScripts(wd, userName)

    if opts.add == True and opts.skipBlast == False:
        base.createDir(wdAdds)
        # Check the existence of BLAST file
        if os.path.exists("%s/goodProteins.blast" %wd) == True:
            # Extract the existing labels
            myLabels = set()
            handle = open("%s/goodProteins.blast" %wd, 'r')
            for line in handle:
                label = line[0:3]
                myLabels.add(label)
            handle.close()
            # Copy to '.' and backup
            base.shell("cp %s/goodProteins.blast ." %wd)
            base.shell("cp %s/goodProteins.blast %s/goodProteins.blast.backup" %(wd, wd))
            print "## Existing labels are %s." %(' '.join(sorted(myLabels)))
            print "## Labels to add are %s." %(' '.join(sorted(labels)))
            for label in labels:
                if label in myLabels:
                    print "Fatal error: added label exists already in BLAST file. Exiting..."
                    sys.exit(-1)
    '''
    prevMolType = None
    for myFile in files:
        molType = checkResidue(myFile)
        if prevMolType == None:
            prevMolType = molType
        elif prevMolType != molType:
            print "Fatal error: files have to all be either amino acids or nucleotides. Exiting..."
            print "File %s failed and was %s." %(myfile, molType)
            sys.exit(0)
    '''

    requiredMolType = "amino acids"
    for myFile in files:
        molType = checkResidue(myFile)
        if requiredMolType != molType:
            print "Fatal error: files have to all be amino acids. Exiting..."
            print "File %s failed and was %s." %(myFile, molType)
            sys.exit(0)

    base.shell("rm -f %s/*.fasta" %wd)
    if opts.add == False: # Do not delete existing fasta files when adding new labels
        base.shell("rm -f %s/*.fasta" %wdFasta)
    else:
        base.shell("rm -f %s/*.fa*" %wdAdds) # Delete the ones to add if there were remnants
    for i in xrange(len(files)):
        myLab, myFile, myPos = labels[i], files[i], positions[i]
        if myFile == "%s.fasta" %myLab:
            print "Fatal error: orthoMCL produces same filenames what you already have. Please renate your fasta files to .fa instead of .fasta. Exiting..."
            sys.exit(0)
        base.shell("orthomclAdjustFasta %s %s %s" %(myLab, myFile, myPos))
        checkUniqueIds("%s.fasta" %myLab)
        if opts.add == False:
            base.shell("mv -f %s.fasta %s" %(myLab, wdFasta))
        else:
            base.shell("cp -f %s.fasta %s" %(myLab, wdAdds))
            #base.shell("mv -f %s.fasta %s" %(myLab, wdFasta))

    #min_length:          minimum allowed length of proteins.  (suggested: 10)
    #max_percent_stop:    maximum percent stop codons.  (suggested 20)
    if opts.skipBlast == False:
        if opts.add == True:
            base.shell("orthomclFilterFasta %s %s 20" %(wdAdds, minlen))
            base.shell("mv -f goodProteins.fasta goodProteinAdds.fasta")
            base.shell("mv -f poorProteins.fasta %s/poorProteinAdds.fasta" %wd)
        base.shell("orthomclFilterFasta %s %s 20" %(wdFasta, minlen))
        base.shell("mv -f poorProteins.* %s" %wd)

    # Blast all against all
    if opts.skipBlast == False:
        if opts.add == True:
            base.shell("makeblastdb -in goodProteinAdds.fasta -dbtype prot")
            #base.shell("mpiformatdb -a F -o F --nfrags=10 -p T -i goodProteinAdds.fasta")
        if opts.uname == "":
            base.shell("makeblastdb -in goodProteins.fasta -dbtype prot")
        #base.shell("mpiformatdb -a F -o F --nfrags=10 -p T -i goodProteins.fasta")
    #base.shell("blastall -a 32 -d goodProteins.fasta -i goodProteins.fasta -m 8 -e %s -p blastp -o goodProteins.blast" %eValue)
    blastEvalue = eValue
    if float(blastEvalue) < 1e-5: blastEvalue = "1e-5"
    if opts.skipBlast == False and opts.add == False:
        #base.shell("blastall -a 32 -d goodProteins.fasta -i goodProteins.fasta -m 8 -e %s -p blastp -o goodProteins.blast" %blastEvalue)
        if opts.uname == '':
            base.shell("blastp -db goodProteins.fasta -query goodProteins.fasta -outfmt 6 -evalue %s -num_threads 32 > goodProteins.blast" %blastEvalue)
        else:
            base.shell("blast.py blastp -query goodProteins.fasta -outfmt 6 -evalue %s -uname %s > goodProteins.blast" %(blastEvalue, opts.uname))
    base.shell("awk '{if ($11<=%s) print $0}' goodProteins.blast > %s/filtered.blast" %(eValue, wd))
    base.shell("mv -f goodProteins.* %s" %wd)
    if opts.skipBlast == False and opts.add == True:
        myP = []
        #cmdStr = "blastall -a 32 -d goodProteins.fasta -i goodProteinAdds.fasta -m 8 -e %s -p blastp -o goodProteinAdds.blast" %blastEvalue
        cmdStr = "blastp -db goodProteins.fasta -query goodProteinAdds.fasta -outfmt 6 -evalue %s -num_threads 32 > goodProteinAdds.blast" %blastEvalue
        p = Process(target=callShell, args=(base, cmdStr, None))
        p.start()
        myP.append(p)
        #cmdStr = "blastall -a 32 -d goodProteinAdds.fasta -i goodProteins.fasta -m 8 -e %s -p blastp -o goodProteinOlds.blast" %blastEvalue
        cmdStr = "blastp -db goodProteinAdds.fasta -query goodProteins.fasta -outfmt 6 -evalue %s -num_threads 32 > goodProteinOlds.blast" %blastEvalue
        p = Process(target=callShell, args=(base, cmdStr, None))
        p.start()
        myP.append(p)
        #cmdStr = "blastall -a 32 -d goodProteinAdds.fasta -i goodProteinAdds.fasta -m 8 -e %s -p blastp -o goodProteinSelf.blast" %blastEvalue
        cmdStr = "blastp -db goodProteinAdds.fasta -query goodProteinAdds.fasta -outfmt 6 -evalue %s -num_threads 32 > goodProteinSelf.blast" %blastEvalue
        p = Process(target=callShell, args=(base, cmdStr, None))
        p.start()
        myP.append(p)
        for p in myP:
            p.join()
        base.shell("cat %s/goodProteins.blast goodProteinOlds.blast goodProteinAdds.blast goodProteinSelf.blast > goodProteins.blast" %wd)
        base.shell("awk '{if ($11<=%s) print $0}' goodProteins.blast > %s/filtered.blast" %(eValue, wd))
        base.shell("mv -f goodProteins*.* %s" %wd)
        base.shell("cp %s/* %s/" %(wdAdds, wdFasta))
        base.shell("cat %s/*.fasta > %s/goodProteins.fasta" %(wdFasta, wd))

    base.shell("orthomclBlastParser %s/filtered.blast %s > %s/similarSequences.txt" %(wd, wdFasta, wd))
    # Prepare database
    base.shell("mysql --user=root --password=password < %s/dropDb.sql" %wd)
    base.shell("mysql --user=root --password=password < %s/createDb.sql" %wd)
    base.shell("orthomclInstallSchema %s/orthomcl.config" %wd)
    base.shell("orthomclLoadBlast %s/orthomcl.config %s/similarSequences.txt" %(wd, wd))
    # Identify potential orthologs
    base.shell("orthomclPairs %s/orthomcl.config %s/orthomclPairs.log cleanup=no" %(wd, wd))
    base.shell("rm -rf pairs")
    base.shell("rm -rf %s/pairs" %wd)
    base.shell("orthomclDumpPairsFiles %s/orthomcl.config" %wd)
    base.shell("mv -f pairs %s" %wd)
    # Group the orthologs
    base.shell("mcl mclInput --abc -I 2.0 -o mclOutput")
    base.shell("orthomclMclToGroups OWN_ 1 < mclOutput > %s/groups.txt" %wd)
    base.shell("mv -f mclInput %s" %wd)
    base.shell("mv -f mclOutput %s" %wd) 
    logHandle.close()
       

if __name__ == "__main__":
    main()
