#
#
# $Header: /opt/cvs/python/packages/share1.5/AutoDockTools/Tests/compare_ligand_files.py,v 1.1 2004/05/10 18:27:04 rhuey Exp $
#
#

import os
from string import find, split, strip



class FileChecker:

    def __init__(self, criteria={}):
        # all have basic criterion:
        #   check number of lines
        self.criteria = {'number_lines':self.compare_number_lines}

    def check_files(self, test_file, std_file, start_key=None):
        #read lines from file to test
        fptr = open(test_file)
        test_lines = fptr.readlines()
        fptr.close()
        #remove lines in test_lines until first line is "REMARK"
        i = 0
        first = 0
        if start_key is not None:
            for i in range(len(test_lines)):
                line = test_lines[i]
                if find(line, start_key)==0:
                    first = i
                    break
        test_lines = test_lines[first:] 
        #read lines from file for comparison
        fptr = open(std_file)
        std_lines = fptr.readlines()
        first = 0
        if start_key is not None:
            for i in range(len(std_lines)):
                line = std_lines[i]
                if find(line, start_key)==0:
                    first = i
                    break
        std_lines = std_lines[first:] 
        fptr.close()
        return self.check(test_lines, std_lines)


    def check(self, test_lines, std_lines):
        result = {} 
        for k, v in self.criteria.items():
            result[k] = apply(v, (test_lines, std_lines,),{})
        return result


    def compare_number_lines(self, lines1, lines2):
        #returns True/False/'ERROR'
        try:
            result = len(lines1)==len(lines2)
        except:
            result = 'ERROR'
        return result


        
class LigandFileChecker(FileChecker):



    def __init__(self):
        FileChecker.__init__(self)
        self.criteria['numatoms'] = self.compare_number_atoms
        self.criteria['numbranches'] = self.compare_number_branches
        self.criteria['numactivetors'] = self.compare_number_activetors
        self.criteria['numaromaticCs'] = self.compare_number_aromaticCs
        self.criteria['rootnum'] = self.compare_root_number
        self.criteria['TORSDOF'] = self.compare_TORSDOF
        self.criteria['root'] = self.compare_root
        self.criteria['lines'] = self.compare_atom_lines
        #this is harder...
        #self.criteria['tors_atomnames'] = self.compare_tors_atomnames


    #def check(self, test_lines, std_lines):
    #    """
    #    compare according to criteria for this FileChecker
    #    """

    def compare_numbers(self, num, std_num):
        try:
            result = num==std_num
        except:
            result = 'ERROR'
        return result


    def compare_number_atoms(self, test_lines, std_lines):
        num = 0
        std_num = 0
        #RESET THESE HERE, ??
        self.atomlines = []
        self.std_atomlines = []
        for l in test_lines:
            if find(l, "HETATM")==0 or find(l, "ATOM")==0:
                num += 1
                self.atomlines.append(l)
        for std_l in std_lines:
            if find(std_l, "HETATM")==0 or find(std_l, "ATOM")==0:
                std_num += 1
                self.std_atomlines.append(std_l)
        return self.compare_numbers(num, std_num)


    def compare_number_branches(self, test_lines, std_lines):
        num = 0
        std_num= 0
        for l in test_lines:
            if find(l, "BRANCH")==0:
                num += 1
        for std_l in std_lines:
            if find(std_l, "BRANCH")==0:
                std_num += 1
        return self.compare_numbers(num, std_num)


    def compare_number_activetors(self, test_lines, std_lines):
        num = 0
        std_num = 0
        for l, std_l in zip(test_lines, std_lines):
            if find(l, "active torsions")>-1:
                testnum = int(split(l)[1])
                stdnum = int(split(std_l)[1])
                return self.compare_numbers(num, std_num)


    def compare_number_aromaticCs(self, test_lines, std_lines):
        num = 0
        std_num= 0
        for l in test_lines:
            if find(l, "HETATM")==0 or find(l, "ATOM")==0:
                num += 1
                ll = split(l)
                if ll[1][0]=='A':
                    nums += 1
        for std_l in std_lines:
            if find(std_l, "HETATM")==0 or find(std_l, "ATOM")==0:
                std_num += 1
                ll = split(std_l)
                if ll[1][0]=='A':
                    std_nums += 1
        return self.compare_numbers(num, std_num)


    def compare_root_number(self, test_lines, std_lines):
        num = 0
        std_num = 0
        ctroot = 0
        std_ctroot = 0
        for l in test_lines:
            if find(l, "ROOT")==0:
                ctroot = 1
            elif find(l, "ENDROOT")==0:
                ctroot = 0
            elif ctroot:
                num += 1
        for std_l in std_lines:
            if find(std_l, "ROOT")==0:
                std_ctroot = 1
            elif find(std_l, "ENDROOT")==0:
                std_ctroot = 0
            elif std_ctroot:
                std_num += 1
        return self.compare_numbers(num, std_num)


    def compare_root(self, test_lines, std_lines):
        rootname = ''
        std_rootname = ''
        for i in range(len(test_lines)):
            l = test_lines[i]
            if find(l, "ROOT")==0:
                rootname = split(test_lines[i+1])[1]
        for i in range(len(std_lines)):
            l = std_lines[i]
            if find(l, "ROOT")==0:
                std_rootname = split(std_lines[i+1])[1]
        #compare_numbers would work for strings, also..FIX THE NAME!!
        return self.compare_numbers(rootname, std_rootname)


    def compare_TORSDOF(self, test_lines, std_lines):
        num = 0
        std_num = 0
        for l in test_lines:
            if find(l, "TORSDOF")==0:
                num = int(split(l)[1])
        for std_l in std_lines:
            if find(std_l, "TORSDOF")==0:
                std_num = int(split(std_l)[1])
        return self.compare_numbers(num, std_num)


    def buildCoordDict(self, atomlines):
        #set up a dictionary where the keys are the coords of an atom line
        #and the values are the lines themselves
        crdD = {}
        for l in atomlines:
            #can't rely on split... need to use indicies
            crdD[(round(float(l[30:38]),3), round(float(l[38:46]),3), round(float(l[46:54])),3)] = l
        return crdD


    def compare_atom_lines(self, test_lines, std_lines):
        if not hasattr(self, 'atomlines'):
            tempresult = self.compare_number_atoms(test_lines, std_lines)
        #for each atom line, parse the coords + find same line in std_lines
        atomdict = self.buildCoordDict(self.atomlines)
        test_keys = atomdict.keys()
        test_key_dict = {}
        for k in atomdict.keys():
            test_key_dict[k] = 1
        std_atomdict = self.buildCoordDict(self.std_atomlines)
        std_keys = std_atomdict.keys()
        std_key_dict = {}
        for k in std_atomdict.keys():
            std_key_dict[k] = 1
        result = {}
        #check that there are no duplicates in test
        result['no_dups_test'] = self.compare_numbers(len(test_key_dict.keys()),
                                            len(test_keys))
        #check that there are no duplicates in std
        result['no_dups_std'] = self.compare_numbers(len(std_key_dict.keys()),
                                            len(std_keys))

        #check that there are no atoms which occur in only 1 of the files
        result['len(dups)=len(std)'] = self.compare_numbers(len(test_key_dict.keys()),
                                            len(std_key_dict.keys()))
        #what if keys differ??
        bad_keys=[]
        for k in test_keys:
            if k not in std_keys:
                bad_keys.append(k)
        std_bad_keys=[]
        for k in std_keys:
            if k not in test_keys:
                std_bad_keys.append(k)
        result['bad_keys'] = bad_keys
        result['std_bad_keys'] = std_bad_keys

        bad_lines = []
        for k in test_keys:
            #check that each line is the same, after the coords of the atom
            #check:
            #     name
            #     charge, if present
            #     solvation parameters, if present
            #check charge field, if present
            testline = atomdict[k]
            stdline = std_atomdict[k]
           
            #check the same format, roughly
            testlen = len(strip(testline))
            stdlen = len(strip(stdline))
            if testlen!=stdlen:
                bad_lines.append(testline)
            elif testline[13:15]!=stdline[13:15]:
                #check that the names are the same except for Hydrogens
                #check SPECIFICALLY for difference in Hydrogen name [yuk]
                if testline[13]=='H' and stdline[13]=='H':
                    continue
                else:
                    bad_lines.append(testline)
            elif testlen>=76:
                #check that the charges are the same +/- .001
                testcharge = float(strip(testline[70:76]))
                stdcharge = float(strip(stdline[70:76]))
                #allow some rounding error..
                if abs(testcharge-stdcharge)>.0015:
                    bad_lines.append(testline)
            elif testlen>=92:
                testAtVol = float(strip(testline[78:84]))
                stdAtVol = float(strip(stdline[78:84]))
                #allow some rounding error..
                if abs(testAtVol-stdAtVol)>.001:
                    bad_lines.append(testline)
                testAtSolPar = float(strip(testline[86:92]))
                stdAtSolPar = float(strip(stdline[86:92]))
                if abs(testAtSolPar-stdAtSolPar)>.001:
                    bad_lines.append(testline)
            
        result['bad_lines'] = bad_lines
        return result


    def report(self, result):
        #print "in report with ", result
        no_errors = True
        err_msg = ""
        for k,v in result.items():
            print "k=", k, " v=", v
            if k not in ['lines', 'bad_keys', 'std_bad_keys']:
                if v==True:
                    continue
                elif v==None:
                    continue
                elif v==False:
                    print k, ':', v
                    err_msg = err_msg + k + ':' + str(v) + "\n"
                    no_errors = False
                elif len(v)!=0:
                    err_msg = err_msg + k + ':'
                    for l in v:
                        err_msg = err_msg + l 
            elif k == 'lines':
                for lines_k, lines_v in v.items():
                    if lines_v not in [True, []]:
                        no_errors = False
                        print lines_k, ':', lines_v
                        err_msg = err_msg + lines_k + ':'
                        if lines_v == False:
                            err_msg = err_msg + 'False\n'
                        else:
                            err_msg = err_msg + str(len(lines_v)) + ' lines\n'
                            for l in lines_v:
                                err_msg = err_msg + l
            else:
                if len(v):
                    print k, ':', v
                    err_msg = err_msg + str(len(v)) + ' items\n'
        if no_errors:
            print "found no difference between files"
        return no_errors, err_msg
               


if __name__ == "__main__":
    lfc = LigandFileChecker()
    results = lfc.check_files('ind.out.pdbq', 'ind.out.pdbq')
    lfc.report(results)
    
