import re, random, os

# open dir.des.scop.txt and dir.cla.scop.txt files for perusing SCOP info
# open file with pdbs containing just one chain in the ASU

# Typical usage below:
"""
from phaser import SCOPdistiller
myobj = SCOPdistiller.SCOPdistiller("SCOP_MR\\dir.des.scop.txt","SCOP_MR\\dir.cla.scop.txt", "SCOP_MR\\PDBsXray1chainASUto2.5A.txt", "SCOP_MR\\All_PDB.txt")
pdbs = myobj.GetPDBsWithSCOPids(0)
myobj.MakeBatchFile(pdbs,"BLASTMRbat.cmd")

or

pdbs = eval(open("SCOP_MR\\newscoppdbs.txt","r").read())
myobj.MakeBatchFile(pdbs,"NewBLASTMRbat.cmd", False)




mfile = open("scoppdbs.txt","w")
mfile.write(str(pdbs))
mfile.close()
# edit scoppdbs.txt to cut away those pdb and fold ids that haven't yet been calculated.
myobj.CheckResultsFile("newcombinedscoppdbs.txt", "Getmyresults.cmd", "Strider")

"""

formatstr = """
cd %s
title BLAST MR runs on %s - %d out of %d
mklink /H %s\Sculpt12prtcl%s_DoneMRJobs.txt %s\\Strider%s_DoneMRJobs.txt
call phenix.python %%%%PHASER_DIST%%%%\\phaser\\BLASTrunner.py d %s Sculpt12prtcl B:\\Users\\oeffner\\HomologySearch\\sculpt12.db N N
del %s\\Sculpt12prtcl%s_DoneMRJobs.txt
cd B:\\Users\\oeffner\\HomologySearch


"""


class SCOPdistiller():
    def __init__(self, SCOPdescriptionfname, SCOPclassesfname, pdbs1chainfname, totalpdbfname):
        desfile = open(SCOPdescriptionfname, "r")
        self.SCOPdescription = desfile.read()
        desfile.close()

        clsfile = open(SCOPclassesfname, "r")
        self.SCOPclasses = clsfile.read()
        clsfile.close()

        asufile = open(pdbs1chainfname,"r")
        self.pdbs1chain = asufile.read()
        asufile.close()

        pdbIDfile = open(totalpdbfname,"r")
        txt = pdbIDfile.read()
        pdbIDfile.close()
        self.pdbcodes = txt.split("\n")
        self.pdbdict = {}
        self.notmakingdictionary = True

        self.classnamedict = {}
        self.foldnamedict = {}
        self.sfnamedict = {}
        self.famdict = {}
# make dictionary of classnames
        self.classnames = re.findall("\tcl\t([a-z]+)\t-\t(.+)",
            self.SCOPdescription)

        for classname in self.classnames:
            self.classnamedict[classname[0]] = classname[1]

# make dictionary of foldnames
        self.foldnames = re.findall("\tcf\t([a-z0-9\.]+)\t-\t(.+)",
          self.SCOPdescription)

        for foldname in self.foldnames:
            self.foldnamedict[foldname[0]] = foldname[1]

# make dictionary of superfamilynames
        self.sfamnames = re.findall("\tsf\t([a-z0-9\.]+)\t-\t(.+)",
          self.SCOPdescription)

        for sfamname in self.sfamnames:
            self.sfnamedict[sfamname[0]] = sfamname[1]

# make dictionary of familynames
        self.familynames = re.findall("\tfa\t([a-z0-9\.]+)\t-\t(.+)",
          self.SCOPdescription)

        for famname in self.familynames:
            self.famdict[famname[0]] = famname[1]

        self.MakePDBidDictionary()
        self.classIDrecmp = re.compile(r"([a-z]+)")
        self.foldIDrecmp = re.compile(r"([a-z]\.[0-9]+)")
        self.sfamIDrecmp = re.compile(r"([a-z]\.[0-9]+\.[0-9]+)")





    def ListClassNames(self):
        return self.classnames




    def ListFoldNames(self):
        foldnames = []
        for foldname in self.foldnames:
            classID = re.search(r"([a-z]+)", foldname[0]).group()

            foldnames.append([foldname[0], self.classnamedict[classID], foldname[1]])

        return foldnames




    def ListSuperFamilyNames(self):
        superfamilynames = []
        for sfamname in self.sfamnames:
            classID = re.search(r"([a-z]+)", sfamname[0]).group()
            foldID = re.search(r"([a-z]\.[0-9]+)", sfamname[0]).group()

            superfamilynames.append([sfamname[0], self.classnamedict[classID],
               self.foldnamedict[foldID], sfamname[1]])

        return superfamilynames




    def ListFamilyNames(self):
        familynames = []
        for famname in self.familynames:
            classID = re.search(r"([a-z]+)", famname[0]).group()
            foldID = re.search(r"([a-z]\.[0-9]+)", famname[0]).group()
            sfamID = re.search(r"([a-z]\.[0-9]+\.[0-9]+)", famname[0]).group()

            familynames.append([famname[0], self.classnamedict[classID],
              self.foldnamedict[foldID], self.sfnamedict[sfamID], famname[1]])

        return familynames




    def MakePDBidDictionary(self):
        self.notmakingdictionary = False
        regex = re.compile( r"^d(....)\w+\s+\1\s+(\S+)\s+(\S+)\s", re.MULTILINE )
        it = regex.finditer( self.SCOPclasses )
        for rec in it:
            ( pdbid, chainid, scopid ) = rec.groups()
            self.pdbdict.setdefault( pdbid, set() ).add( ( pdbid, chainid, scopid ) )
# add empty elements for those pdb codes that haven't been SCOP classified
        for pdbid in self.pdbcodes:
            if pdbid.lower() not in self.pdbdict:
                self.pdbdict[pdbid.lower()] = [('','','')]

        self.notmakingdictionary = True




    def GetFamilyIdFromPDB(self, pdb, chainid = ""):
# returns array with pdb, chain and res id, scop id and residue length
        if len(self.pdbdict.items()) > 0 and self.notmakingdictionary:
            if pdb.lower() in self.pdbdict:
                famids = list(self.pdbdict[pdb.lower()])
                if chainid != "": # make dictionary according to chainid and select it
                    FamidsWithCorrectChainid = []
                    for famid in famids:
                        if famid[1].find(chainid.upper(),0) > -1:
                            FamidsWithCorrectChainid.append(famid)

                    if len(FamidsWithCorrectChainid) > 0:
                        match = FamidsWithCorrectChainid
                    else:
                        match = [('','','','')]
                else:
                    match = famids
            else:
                match = [('','','','')]
        else:
            chainresidueID = chainid + "\S+"
            match = re.findall("\s(" + pdb.lower() + ")\s+(" + chainresidueID + ")\s(\S+)\s",
             self.SCOPclasses)

# if pdb entry is present then match yields something like
# [('1coy', 'A:319-450', 'd.16.1.1'), ('1coy', 'A:4-318,A:451-506', 'c.3.1.2')]
        res = [('','','','')]
        try:
# If the length of the pdb id is > 0 then the pdb id exists in the SCOP and match isn't empty
            if len(match[0][0]) > 0:
# want to return something like
# [('1coy', 'A:319-450', 'd.16.1.1', 131), ('1coy', 'A:4-318,A:451-506', 'c.3.1.2', 369)]
# where 131 and 369 are the residue length belonging to the respective domains
                res =[]
                for m in match:
                    m = list(m)
                    residueranges = m[1].split(",")
                    #m[1] = residueranges
                    resdomainlength = 0
                    for rng in residueranges:
                        resrng = rng.split(":")[1]
                        if len(resrng) > 0:
# deal with negative residue numbers, i.e. peptide markers and so
                            if resrng[0] != "-":
                                resdomainlength += abs(int(resrng.split("-")[1]) - int(resrng.split("-")[0]))+1
                            else: # negative residue numbers are ignored and cast to zero
                                nnresrng = "0" + resrng[2:]
                                resdomainlength += abs(int(nnresrng.split("-")[1]) - int(nnresrng.split("-")[0]))+1

                    m.append(resdomainlength)

                    m = tuple(m)
                    res.append(m)

        except  Exception, e:
            print "SCOPdistiller: " + str(e)
            res =  [('','','','')]

        finally:
            return res




    def GetFamilyNameFromPDB(self, pdb, chainid = ""):
        compositefamname = []
        famids = self.GetFamilyIdFromPDB(pdb, chainid)
        if famids != [('','','','')]:
            for famid in famids:
                classID = self.classIDrecmp.search(famid[2]).group()
                foldID = self.foldIDrecmp.search(famid[2]).group()
                sfamID = self.sfamIDrecmp.search(famid[2]).group()
                """
                classID = re.search(r"([a-z]+)", famid[2]).group()
                foldID = re.search(r"([a-z]\.[0-9]+)", famid[2]).group()
                sfamID = re.search(r"([a-z]\.[0-9]+\.[0-9]+)", famid[2]).group()
                """

                familyname = (famid[2], self.classnamedict[classID],
                  self.foldnamedict[foldID], self.sfnamedict[sfamID], self.famdict[famid[2]])

                compositefamname.append(familyname)
        else:
            compositefamname = [('UnSCOPed','UnSCOPed','UnSCOPed','UnSCOPed','UnSCOPed')]
            #compositefamname = [('','','','','')]

        return compositefamname




    def ListPDBsofIDs(self, id):
        dot_count = id.count( "." )

        if dot_count < 3:
            id += "."
        escaped = re.escape(id) # escape it since there will be dots in the id
        pdbs = re.findall("\s([a-z0-9]{4})\s.+\s("+ escaped +"[\S]*)",
            self.SCOPclasses)

        return pdbs




    def PrintListofPDBsofIDs(self, id, s=""): # get all pdbs branching below this hierachy level
        names = ""
        ndots = len(re.findall("\.",id))# count number of dots in the id
        if ndots == 0: # it's a class id
            names = self.ListClassNames()

        if ndots == 1: # it's a fold id
            names = self.ListFoldNames()

        if ndots == 2: # it's a superfamily id
            names = self.ListSuperFamilyNames()

        if ndots == 3: # it's a family id
            names = self.ListFamilyNames()

        for i in range(len(names)):
            if names[i][0] == id:
                break

        if names[i][0] != id: # didn't find it so don't print anything
            return

        pdbs = self.ListPDBsofIDs(names[i][0])

        print names[i]
        if s== "v": # verbose
            for pdb in pdbs:
                print "\t\t%s, %s" %pdb

        return pdbs




    def ListPDBs1ChainInASU(self, pdbs, s = ""):
        selectedpdbs = []
        for pdb in pdbs:
            if self.pdbs1chain.find(pdb[0].upper(),0) >= 0:
                selectedpdbs.append(pdb)

        if s== "v": # verbose
            for pdb in selectedpdbs:
                print "\t\t%s, %s" %pdb

        return selectedpdbs




# gather all the above into what we really want
    def GetPDBsAndSCOPids(self, x, pdbsfname="", maxpdbs = 100000):
        pdbsmodresshuffled = []
        random.seed(x) # so that we can reproduce lists seeded with the same number

        alreadyusedset = set([])
        if pdbsfname != "":
            mfile= open(pdbsfname,"r")
            txt = mfile.read()
            mfile.close()
            alreadyusedpdbs = eval(txt)
            for pdb in alreadyusedpdbs:
                alreadyusedset.add(pdb[0])


        for cls in self.ListClassNames():
            pdbset = set([])
            foldset = set([])
            uniquepdbs = []
            uniquefolds = []
            pdbs = self.PrintListofPDBsofIDs(cls[0])
            pdbs1chain = self.ListPDBs1ChainInASU(pdbs)
# Shuffle list so that if more pdb ids have the same fold, i.e. all but the first pdb
# would be discarded below, a list generated with one random seed
# will not be the same as when generated with a different random number seed.
            #random.shuffle(pdbs1chain)
            notusedset = set([])

            for pdb in pdbs1chain:
                if not pdb[0] in alreadyusedset:
                    notusedset.add(pdb)
                    alreadyusedset.add(pdb)
                else:
                    print "%s is already in alreadyusedset" % pdb[0]

# skip duplicate folds
            for pdb in notusedset:
                if not pdb[1] in foldset:
                    uniquefolds.append(pdb)
                    foldset.add(pdb[1])
                else:
                    print "%s is already in foldset" % pdb[1]


#            count = 0
# skip duplicate pdb codes
#            for pdb in uniquefolds:
#                if not pdb[0] in pdbset and count < maxpdbs:
#                    uniquepdbs.append(pdb)
#                    pdbset.add(pdb[0])
#                    print pdb
#                    count = count+1


            pdbsmodresshuffled.extend(notusedset)
            #pdbsmodresshuffled.extend(uniquefolds)
            #pdbsmodresshuffled.extend(pdbs1chain)
            print "len(notusedset) = %d" %len(notusedset)

        return pdbsmodresshuffled





    def GetPDBsWithSCOPids(self, x, pdbsfname=""):
        pdbsmodresshuffled = []
        random.seed(x) # so that we can reproduce lists seeded with the same number

        alreadyusedset = set([])
        notusedset = set([])

        if pdbsfname != "":
            mfile= open(pdbsfname,"r")
            txt = mfile.read()
            mfile.close()
            alreadyusedpdbs = eval(txt)
            for pdb in alreadyusedpdbs:
                alreadyusedset.add(pdb[0].lower())


        pdblist = self.pdbs1chain.split()
        #random.shuffle(pdblist)
        for pdb in pdblist:
            if not pdb.lower() in alreadyusedset and len(pdb) == 4:
                print "pdb = %s" %pdb
                scopid = self.GetFamilyIdFromPDB(pdb)[0][2]
                if scopid == "":
                    scopid = "UnSCOPed"
                notusedset.add((pdb, scopid))
                alreadyusedset.add(pdb)
            else:
                print "%s is already in alreadyusedset" % pdb

        return notusedset




    def MakeBatchFile(self,pdbs, batfname, win=True):
        mstr = ""
        i = 0
        total = len(pdbs)
        for pdb in pdbs:
            i += 1
# compile batch file for doing BLAST MR calculations
# only want directories named with SCOP fold id
            if len(pdb[1].split(".")) > 1: # if it's a proper SCOPid
                foldID = pdb[1].split(".")[0] + "." + pdb[1].split(".")[1]
            else: # otherwise it says "UnSCOPed"
                foldID = pdb[1]
# for windows
            if win:
                mstr += "if exist StopThreadsGracefully goto veryend\n"
                #mstr += "if not exist %s mkdir %s && cd %s && mklink /J SCOP_MR ..\SCOP_MR\n" \
                #  %(foldID, foldID, foldID)
                #mstr += "if exist %s cd %s\n" %(foldID, foldID)
                mstr += "cd %s\n" %foldID
                mstr += "title BLAST MR runs on %s\\%s - %d out of %d\n" %(foldID,pdb[0].upper(), i, total)
                #mstr += "python E:\\Users\\oeffner\\Phenix\\cctbx_sources\\Develop\\phaser\\phaser\\BLASTrunner.py c %s 90 15 20000 Y e:\mtz_files rdo20 Y 8 Strider\n" %pdb[0].upper()
                mstr += "call phenix.python %%%%PHASER_DIST%%%%\\phaser\\BLASTrunner.py f %s 60 15 20000 Y P:\\MTZ_files rdo20 1 8 Olga 10\n" %pdb[0].upper()
                mstr += "@echo BLAST MR run on %s, %s,   %d out of %d " %(pdb[0].upper(),pdb[1],i,total)
                mstr += "returned %errorlevel% >> B:\\Users\\oeffner\\HomologySearch\\MRscriptresults.txt\n"
                mstr += "cd B:\\Users\\oeffner\\HomologySearch\n\n"
# for bash
            else:
                mstr += "if [ -f StopThreadsGracefully ]; then break; fi\n"
                mstr += "cd %s\n" %foldID
                mstr += "phenix.python $PHENIX/phaser/phaser/BLASTrunner.py b %s 4 Strider\n" %pdb[0].upper()
                mstr += "echo BLAST MR run on %s, %s,   %d out of %d " %(pdb[0].upper(),pdb[1],i,total)
                mstr += "returned $? >> %sresults.txt\n" %batfname
                mstr += "cd $HOME/HomologySearch/\n\n"



        mstr += ":veryend\n"
        mstr += "echo all over now\n\n"

        print "Writing out %s to disk." %batfname
        blastmrfile = open(batfname,"w")
        blastmrfile.write(mstr)
        blastmrfile.close()





    def CheckResultsFile(self,pdbscopfname, getresbatfname, calcname):
        mstr = ""
        i = 0
        mfile= open(pdbscopfname,"r")
        txt = mfile.read()
        mfile.close()
        pdbs = eval(txt)
        total = len(pdbs)
# to make scripts the same speed when dividing this one into several smaller ones we randomly shuffle
# pdbs as to evenly spreading out large as well as short MR calculations
        random.shuffle(pdbs)
        print "Skimming through %d calculations" %total

        mstr = ["","","","","","","",""]

        topdir = os.getcwd()
        calcs_in_file= total/8
        n = 0
        eigths = calcs_in_file
        m = 0
        for (i,pdb) in enumerate(pdbs):
            m += 1
            if i > eigths:
                n += 1
                eigths = calcs_in_file * (n+1)
                m = 0

# only want directories named with SCOP fold id
            if pdb[1] == "UnSCOPed":
                foldID = "UnSCOPed"
            else:
                foldID = pdb[1].split(".")[0] + "." + pdb[1].split(".")[1]
            pdbcalcdir = os.path.join(foldID, pdb[0].upper())
            blastMRdonefile = os.path.join(pdbcalcdir, "Strider" + pdb[0].upper() + "_DoneMRJobs.txt")
            if os.path.exists(blastMRdonefile):
                print "%d, Calculation %d on %s seems OK" %(n, i, pdb[0].upper())
# compile batch file for extracting results
                mstr[n] += "title Extracting BLAST MR results from the %d. calculation out of %d \n" %(m, calcs_in_file)
                mstr[n] += "if exist %s cd %s\n" %(foldID, foldID)
                mstr[n] += "python %PHASER_DIST%\\phaser\\BLASTrunner.py d " + pdb[0].upper()
# location of database file
                mstr[n] += " " + calcname + " " + topdir + "\\Completeresults" + str(n+1) + ".db N\n"
                mstr[n] += "@echo Results from %s, %s,  %d , %d " %(foldID, pdb[0].upper(), i, m)
                mstr[n] += "returned %errorlevel% >> "
                mstr[n] += "E:\\Users\\oeffner\\HomologySearch\\GatherDBaseMRresults%d.txt\n" %(n+1)
                mstr[n] += "cd " + topdir + "\n\n"

            else:
                print "No results for the %d run on %s, %s" %(i, pdb[0].upper(), foldID)

        for i in range(0,8):
            blastmrresfile = open(getresbatfname + str(i+1) + ".cmd","w")
            blastmrresfile.write(mstr[i])
            blastmrresfile.close()




    def CheckResultsFile2(self,pdbsfname, getresbatfname, calcname):
        mstr = ""
        i = 0
        mfile= open(pdbsfname,"r")
        txt = mfile.read()
        mfile.close()
        pdbs = eval(txt)
        total = len(pdbs)
# to make scripts the same speed when dividing this one into several smaller ones we randomly shuffle
# pdbs as to evenly spreading out large as well as short MR calculations
        random.shuffle(pdbs)
        print "Skimming through %d calculations" %total

        mstr = ["","","","","","","",""]

        def writestrings(foldid, tpdb):
            pdbcalcdir = os.path.join(foldid, tpdb[0].upper())
            blastMRdonefile = os.path.join(pdbcalcdir, "Strider" + tpdb[0].upper() + "_DoneMRJobs.txt")
            if os.path.exists(blastMRdonefile):
                print "%d, Calculation %d on %s seems OK" %(n, i, tpdb[0].upper())
# compile batch file for extracting results
                mstr[n] += "title Extracting BLAST MR results from the %d. calculation out of %d \n" %(m, calcs_in_file)
                mstr[n] += "if exist %s cd %s\n" %(foldid, foldid)
                mstr[n] += "python %PHASER_DIST%\\phaser\\BLASTrunner.py d " + tpdb[0].upper()
# location of database file
                mstr[n] += " " + calcname + " " + topdir + "\\Completeresults" + str(n+1) + ".db N N\n"
                mstr[n] += "@echo Results from %s, %s,  %d , %d " %(foldid, tpdb[0].upper(), i, m)
                mstr[n] += "returned %errorlevel% >> "
                mstr[n] += "E:\\Users\\oeffner\\HomologySearch\\GatherDBaseMRresults%d.txt\n" %(n+1)
                mstr[n] += "cd " + topdir + "\n\n"
                return 1
            else:
                return 0

        topdir = os.getcwd()
        calcs_in_file= total/8
        n = 0
        eigths = calcs_in_file
        m = 0
        for (i,pdb) in enumerate(pdbs):
            m += 1
            if i > eigths:
                n += 1
                eigths = calcs_in_file * (n+1)
                m = 0

            scopids = self.GetFamilyIdFromPDB(pdb[0].upper())
# only want directories named with SCOP fold id
            ret = 0
            if scopids[0][2] == "":
                foldID = "UnSCOPed"
                ret += writestrings(foldID, pdb)
            else:
                for scopid in scopids:
                    foldID = scopid[2].split(".")[0] + "." + scopid[2].split(".")[1]
                    ret += writestrings(foldID, pdb)
            if ret < 1:
                print "No results for the %d run on %s" %(i, pdb[0].upper())

        for i in range(0,8):
            blastmrresfile = open(getresbatfname + str(i+1) + ".cmd","w")
            blastmrresfile.write(mstr[i])
            blastmrresfile.close()



    def CheckResultsFile3(self,pdbscopfname, getresbatfname, calcname, sqlfname):
        mstr = ""
        mfile= open(pdbscopfname,"r")
        txt = mfile.read()
        mfile.close()
        pdbs = eval(txt)
        total = len(pdbs)
# to make scripts the same speed when dividing this one into several smaller ones we randomly shuffle
# pdbs as to evenly spreading out large as well as short MR calculations
        print "Skimming through %d calculations" %total

        mstr = ""

        topdir = os.getcwd()
        for (i,pdb) in enumerate(pdbs):
# only want directories named with SCOP fold id
            if pdb[1] == "UnSCOPed":
                foldID = "UnSCOPed"
            else:
                foldID = pdb[1].split(".")[0] + "." + pdb[1].split(".")[1]
            pdbcalcdir = os.path.join(foldID, pdb[0].upper())

            blastMRdonefile = os.path.join(pdbcalcdir, "Strider" + pdb[0].upper() + "_DoneMRJobs.txt")
            if os.path.exists(blastMRdonefile):
                print "Calculation %d on %s seems OK" %(i, pdb[0].upper())
# compile batch file for extracting results
                mstr += "title MLAD calculations from the %d. calculation\n" %i
                mstr += "if exist %s cd %s\n" %(pdbcalcdir, pdbcalcdir)
                mstr += "call phenix.python M:\\SourceFiles\\Python\\CalcTemplateMLADs.py "\
                  + calcname + " " + pdb[0] + ".pdb " + sqlfname
                mstr += "\ncd " + topdir + "\n\n"

            else:
                print "No results for the %d run on %s, %s" %(i, pdb[0].upper(), pdbcalcdir)

        blastmrresfile = open(getresbatfname + ".cmd","w")
        blastmrresfile.write(mstr)
        blastmrresfile.close()





    def CheckResultsFile4(self,pdbscopfname, getresbatfname):
        mstr = ""
        mfile= open(pdbscopfname,"r")
        txt = mfile.read()
        mfile.close()
        pdbs = eval(txt)
        total = len(pdbs)
# to make scripts the same speed when dividing this one into several smaller ones we randomly shuffle
# pdbs as to evenly spreading out large as well as short MR calculations
        print "Skimming through %d calculations" %total

        mstr = ""

        topdir = os.getcwd()
        for (i,pdb) in enumerate(pdbs):
# only want directories named with SCOP fold id
            targpdb = pdb[0].upper()
            if pdb[1] == "UnSCOPed":
                foldID = "UnSCOPed"
            else:
                foldID = pdb[1].split(".")[0] + "." + pdb[1].split(".")[1]
            pdbcalcdir = os.path.join(foldID, targpdb)

            blastMRdonefile = os.path.join(pdbcalcdir, "Strider" + targpdb + "_DoneMRJobs.txt")
            if os.path.exists(blastMRdonefile):
                print "Calculation %d on %s seems OK" %(i, targpdb)
# compile batch file for extracting results
                mstr += formatstr %(foldID, pdbcalcdir, i, total, targpdb, targpdb,
                 targpdb, targpdb, targpdb, targpdb, targpdb)
            else:
                print "No results for the %d run on %s, %s" %(i, targpdb, pdbcalcdir)

        blastmrresfile = open(getresbatfname + ".cmd","w")
        blastmrresfile.write(mstr)
        blastmrresfile.close()