# $CVSHeader: pal8e/src/Comparisons.py,v 1.5 2005/02/15 01:10:25 indraneel Rel $

## compares outputs (PSSD, DSSP, stride)

from Bin import Bin
from Sigma import Sigma
from Log import LogFH
from string import split, join
import Numeric, os
import string

from PSEAParser import PSEA as PSEAParserPSEA
from PSSDParser import PSSD as PSSDParserPSSD
from DSSPParser import DSSP as DSSPParserDSSP

class TypeTranslator:
    def __init__(self, logger, type_definition):
        self.type_definition = type_definition

    def getType(self, ss):
        if not self.type_definition: return ss
        if self.type_definition.has_key(ss):
            return self.type_definition[ss]
        return ' '

class ElemData(dict):
    def __repr__(self):
        line = ''
        line += self['Expt']+'  '+self['fileN']+'\n'
        line += 'status: '+self['status']+'  '
        line += 'found: '+str(self['found'])+'  '
        line += 'elemLen: '+str(self['elemLen'])+'  '
        line += 'resUsed: '+str(self['resUsed'])+'  '
        line += 'conseqMiss: '+str(self['conseqMiss'])+'  '+'\n'
        line += 'resIn: '+str(self['resIn'])+'\n'
        line += 'resOut: '+str(self['resOut'])+'\n'
        line += str(self['elem'])
        return line

class ElemDataList(list):
    def __repr__(self):
        line = ''
        for elem in self:
            line += str(elem)+'\n'
        return line

class Common:
    def __init__(self, logger, outfile):
        self.logger = logger
        self.outfile = outfile
        self.elemFoundData = ElemDataList()
        self.elemMissedData = ElemDataList()
        self.resPercentFound = []
        self.resPercentFoundE = [] # out of total residues
        self.resPercentFoundH = [] # out of total residues
        self.resBasePercentFoundE = [] # out of total residues
        self.resBasePercentFoundH = [] # out of total residues

    def setTranslator(self, translator):
        self.translator =  translator

    def compare(self, elemObj, resObj, percent):
        '''
        minimum percent of self element located by
        residues of other for element to be "found"
        '''
        self.elemObj = elemObj
        self.resObj = resObj
        expt = 'Elem: '+elemObj.typeclass+' Res: '+resObj.typeclass
        self.fileElemFoundData = ElemDataList()
        self.fileElemMissedData = ElemDataList()
        for element in elemObj.elements:
            #if element.type not in ['H', 'E']: continue
            #if element.type != 'E': continue
            #if len(element) < 3: continue
            if element.type != ' ': continue
            found = 0
            missed = 0
            total = 0
            resIn = []
            resOut = []
            conseqMiss = []
            conseq_miss = 0
            for res in element:
                if not resObj.has_key(res.chain): continue
                resID = (res.resid, res.icode)
                if not resObj[res.chain].has_key(resID):
                    continue
                total += 1
                resObjTypes = resObj[res.chain][resID]
                resSStypes = map(self.translator.getType, resObjTypes['SS'])
                if element.type in resSStypes:
                    found += 1
                    resIn.append(res)
                    conseqMiss.append(conseq_miss)
                    conseq_miss = 0
                else:
                    missed += 1
                    resOut.append(res)
                    conseq_miss += 1
            conseqMiss.append(conseq_miss)
            conseqMiss = max(conseqMiss)
            percentFound = (float(found)/total) * 100
            elemData = {'Expt':expt, \
                        'fileN':elemObj.fname + ' ' + resObj.fname, \
                        'elem':element, \
                        'found':percentFound, \
                        'elemLen':len(element), \
                        'resUsed':total, \
                        'resIn':resIn, \
                        'resOut':resOut, \
                        'conseqMiss':conseqMiss
                        }
            elemData = ElemData(elemData)
            if (percentFound >= percent):
                elemData['status'] = 'Found'
                self.fileElemFoundData.append(elemData)
            else:
                elemData['status'] = 'Missed'
                self.fileElemMissedData.append(elemData)
            #print elemData
        #print self.fileElemMissedData
        self.elemFoundData += self.fileElemFoundData
        self.elemMissedData += self.fileElemMissedData

    def compareReliability(self, objBase, objTest):
        '''
        residue to residue match with base definition
        '''
        total = 0
        match = 0
        baseH = 0
        baseE = 0
        testH = 0
        testE = 0
        for chain in objBase.keys():
            if not objTest.has_key(chain): continue
            for res in objBase[chain].keys():
                if not objTest[chain].has_key(res): continue
                baseSS = objBase[chain][res]['SS']
                baseSS.sort()
                testSS = objTest[chain][res]['SS']
                testSS.sort()
                if baseSS == testSS: match += 1
                total += 1
                if 'E' in baseSS: baseE += 1
                if 'H' in baseSS: baseH += 1
                if 'E' in testSS: testE += 1
                if 'H' in testSS: testH += 1
	#print total
        if total > 0:
            self.resPercentFound.append(float(match)/total*100)
            self.resPercentFoundE.append(float(testE)/total*100)
            self.resPercentFoundH.append(float(testH)/total*100)
            self.resBasePercentFoundE.append(float(baseE)/total*100)
            self.resBasePercentFoundH.append(float(baseH)/total*100)

    def compareReliabilityFound(self, objBase, objTest):
        '''
        residue to residue match with base definition from DSSP
        overpredictions are neglected
        '''
        total = 0
        match = 0
        baseH = 0
        baseE = 0
        testH = 0
        testE = 0
        for chain in objBase.keys():
            if not objTest.has_key(chain): continue
            for res in objBase[chain].keys():
                if not objTest[chain].has_key(res): continue
                baseSS = objBase[chain][res]['SS'][0]
                testSS = objTest[chain][res]['SS']
                if 'b' in testSS: testSS.append('E')
                if 'a' in testSS: testSS.append('H')
                if baseSS in ['H', 'I', 'G']: baseSS = 'H'
                total += 1
                if 'E' == baseSS: baseE += 1
                if 'H' == baseSS: baseH += 1
                if 'E' in testSS: testE += 1
                if 'H' in testSS: testH += 1
                if baseSS in ['H', 'E']:
                    if baseSS in testSS: match += 1
        if total > 0:
            #self.resPercentFound.append(float(match)/total*100)
            self.resPercentFound.append(float(match)/(baseE+baseH)*100)
            self.resPercentFoundE.append(float(testE)/total*100)
            self.resPercentFoundH.append(float(testH)/total*100)
            self.resBasePercentFoundE.append(float(baseE)/total*100)
            self.resBasePercentFoundH.append(float(baseH)/total*100)

    def genInsightInput(self, fileID, insLogDir, insRelPath):
        if not self.fileElemMissedData: return
        relFname = insRelPath+fileID+'.pdb'
        oid = 'P_'+fileID
        lines = ['Get Molecule PDB User '+relFname+' '+oid+'\n']
        lines += ['Display Molecule Only Atoms Trace '+oid+'\n']
        lines += ['Color Molecule Atoms '+oid+ \
                  ' Specified Specification 247,247,247\n']
        lines += ['Clip -Auto 400 0 -Perspective INSIGHTII\n']
        elemcolor = 'green'
        totalResOut = 0
        conseqMiss = []
        percentFound = []
        percentFoundElemLen = []
        for elemData in self.fileElemMissedData:
            elem = elemData['elem']
            stch=elem[0].chain
            stres=elem[0].resid
            endres=elem[-1].resid
            sticode=elem[0].icode
            endicode=elem[-1].icode
            if stch == ' ': stch = ''
            if sticode == ' ': sticode = ''
            if endicode == ' ': endicode = ''
            line = 'Color Molecule Atoms '+oid+':'
            line += stch+str(stres)+sticode+'-'+stch+str(endres)+endicode
            line += ' Specified Specification '+elemcolor+'\n'
            lines += [line]
            for res in elemData['resOut']:
                chn = res.chain
                ricode = res.icode
                SS = self.resObj[chn][(res.resid, ricode)]['SS']
                if ('H' in SS) and ('E' in SS):
                    rescolor = 'magenta'
                elif 'H' in SS:
                    rescolor = 'cyan'
                elif 'E' in SS:
                    rescolor = 'yellow'
                else: rescolor = 'red'
                if chn == ' ': chn = ''
                if ricode == ' ': ricode = ''
                line = 'Color Molecule Atoms '+oid+':'
                line += chn+str(res.resid)+ricode
                line += ' Specified Specification '+rescolor+'\n'
                lines += [line]
            totalResOut += len(elemData['resOut'])
            conseqMiss.append(elemData['conseqMiss'])
            percentFound.append(elemData['found'])
            percentFoundElemLen.append(len(elemData['elem']))
        totalres = "%03d" % totalResOut
        conseqMiss = "%03d" % (max(conseqMiss))
        m = min(percentFound)
        percentfound = "%06.2f" % (m)
        idx = percentFound.index(m)
        percentfoundElemLen = "%03d" % (percentFoundElemLen[idx])
        fh=file(insLogDir+'/'+percentfoundElemLen+'_'+percentfound+'_'+ \
                fileID+'.log', 'w')
        fh.writelines(lines)
        fh.close()
        
    def report(self):
        lines = []
        for elemData in self.elemMissedData:
            lines.append(str(elemData)+'\n')
        fh=file(self.outfile,'a')
        fh.writelines(lines)
        fh.close()

    def setHandle(self, handle):
        self.handle = handle

    def reportReliability(self, fh=None):
        resPercentFound = self.resPercentFound
	#print resPercentFound
        items = len(resPercentFound)
        total = sum(resPercentFound)
        mean = total/items
        sigma = Sigma(resPercentFound)
        s = sigma.sigma()
        line = "%7.2f" % (mean)+'\t'+ "%6.2f" % (s)+'\t'+self.handle
        line += '\t'+ "%7.2f" % (sum(self.resPercentFoundE)/items)
        line += '\t'+ "%7.2f" % (sum(self.resPercentFoundH)/items)
        line += '\t'+ "%7.2f" % (sum(self.resBasePercentFoundE)/items)
        line += '\t'+ "%7.2f" % (sum(self.resBasePercentFoundH)/items)
        line += '\n'
        if fh == None:
            fh=file(self.outfile,'a', 1) # line buffered
            fh.write(line)
            fh.close()
        else: fh.write(line)

class Stride(dict):
    def __init__(self):
        self.typeclass = 'Stride'

class PSEA(PSEAParserPSEA):
    def __init__(self, logger):
        self.typeclass = 'PSEA'
        self.logger = logger
        PSEAParserPSEA.__init__(self, logger)
        
    def store(self, fname):
        return PSEAParserPSEA.store(self, fname)
        
class DSSP(DSSPParserDSSP):
    def __init__(self, logger):
        self.typeclass = 'DSSP'
        self.logger = logger
        DSSPParserDSSP.__init__(self, logger)

    def store(self, fname):
        return DSSPParserDSSP.store(self, fname)

    def compare(self, other, common):
        if other.typeclass == 'PSSD':
            type_translator = TypeTranslator(self.logger, None)
            common.setTranslator(type_translator)
            #percent = 80 # for H or E regions
            percent = 50 # for not H nor E regions
        common.compare(self, other, percent)

class PSSD(PSSDParserPSSD):
    def __init__(self, logger):
        self.typeclass = 'PSSD'
        self.logger = logger
        PSSDParserPSSD.__init__(self, logger)
        
    def store(self, fname):
        return PSSDParserPSSD.store(self, fname)
        
    def compare(self, other, common):
        if other.typeclass == 'DSSP':
            typedef={'H':'H', 'G':'H', 'I':'H', 'E':'E'}
            type_translator = TypeTranslator(self.logger, typedef)
            common.setTranslator(type_translator)
            percent = 20
        common.compare(self, other, percent)

class Runner:
    def __init__(self, PSSDdir, outfile, logfile):
        self.pssddir = PSSDdir
        self.outfile = outfile
        self.logger = LogFH(logfile)
        if os.path.isfile(self.outfile):
            os.unlink(self.outfile)
        self.cmptypes = []
        
    def set_cmptype(self, TypeClass, TypeInDir, TypeFileExt):
        self.cmptypes.append([TypeClass, TypeInDir, TypeFileExt])
        
    def compare(self, elemObj, resObj, common):
        elemObj.compare(resObj, common)
        
    def runall(self):
        dssp_cmp_pssd = Common(self.logger, self.outfile)
        for (i, fname) in enumerate(os.listdir(self.pssddir)):
            fileID = string.join(fname.split('.')[:-1], '')
            self.logger.log(str(i)+'\t'+fileID)
            pssd = PSSD(self.logger)
            fname = self.pssddir+fileID+'.ssd'
            if (pssd.store(fname) == 'Err'): continue
            for cmptype in self.cmptypes:
                [TypeClass, TypeInDir, TypeFileExt] = cmptype
                typeInstance = TypeClass(self.logger)
                fname = TypeInDir+fileID+TypeFileExt
                if (typeInstance.store(fname) == 'Err'): continue
                if typeInstance.typeclass == 'DSSP':
                    self.compare(typeInstance, pssd, dssp_cmp_pssd)
                    insLogDir = '../../inslog2/'
                    insRelPath = 'wrk/dat/nrchn2787/'
                    dssp_cmp_pssd.genInsightInput(fileID, \
                                                  insLogDir, insRelPath)
        dssp_cmp_pssd.report()

if __name__ == '__main__':
    outfile = 'comp-out'
    logfile = 'comp-log'
    
    PSSDdir = '../../dat/cull100.pssd/'
    DSSPdir = '../../dat/cull100.dssp/'
    PSEAdir = '../../dat/cull100.psea/'

    runner = Runner(PSSDdir, outfile, logfile)
    runner.set_cmptype(DSSP, DSSPdir, '.dssp')
    runner.runall()
