'''
$Id: viewhits.py,v 1.1 2003/02/15 17:08:23 mliang Exp $
Copyright (c) 2003 Mike Liang.  All rights reserved.
Modified by Shirley Wu, May 2008.


DESCRIPTION

    "viewhits" will visualize hitfiles from Feature.  By default, it
    will display the top 100 scores.  Color can be specified by name
    or parenthesized, comma separted RGB tuple.  There is no way to
    specify "top X" scores yet, only cutoff scores.  

USAGE

    viewhits <hitfile> [, <cutoff>] [, <color>] [, <radius>]

'''

from pymol import cmd
from pymol.cgo import *
from pymol import viewing
from pymol import querying
import os


# =================================================================
class Hit:
    # =========================================================
    def __init__(self):
        self.x = None
        self.y = None
        self.z = None
        self.score = None

    # =========================================================
    def __str__(self):
        return '((%g %g %g) %g)' % (self.x, self.y, self.z, self.score)
    
    # =========================================================
    def __cmp__(self, other):
        if isinstance(other, Hit):
            return cmp(self.score, other.score)
        return cmp(self.score, other)

    # =========================================================
    def getScore(self):
        return self.score

    # =========================================================
    def getLocation(self):
        return (self.x, self.y, self.z)


# approximate floating point expression
FLOAT_RE = re.compile('[0-9.-]+')

# =================================================================
class HitFile:
    # =========================================================
    def __init__(self, filename = None):
        self.hits = []

        if filename:
            self.loadFile(filename)

    # =========================================================
    def loadFile(self, filename):
        for line in open(filename).xreadlines():
            hit = Hit()
	    fields = line.strip().split()
	    data = (fields[1], fields[2], fields[3], fields[4])
            hit.score, hit.x, hit.y, hit.z = map(float,data)
            self.hits.append(hit)
        
    # =========================================================
    def __str__(self):
        sio = cStringIO.StringIO()
        for hit in self.hits:
            print >>sio, hit
        return sio.getvalue()


# =================================================================
def _viewhits(hitfile, cutoff=None, color=(1.0,0.0,0.0), radius=0.3 ):
    hit_objs = []
    try:
        hf = HitFile(hitfile)
    except IOError:
        print "Error opening file: %s" % hitfile
        return
    if not hf: return

    print 'hitfile',hitfile,'cutoff',cutoff,'color',color,'radius',radius

    if cutoff == None:    
        hf.hits.sort()
        hits = hf.hits[-100:]
        print 'Cutoff Score:',hits[0].score
    else:
        hits = filter(lambda h,c=cutoff:h.score>=c, hf.hits)

    print "Number of Hits:", len(hits)
    hit_objs += [ COLOR, color[0], color[1], color[2] ]    
    for hit in hits:
        hit_objs += [ SPHERE, hit.x, hit.y, hit.z, radius ]

    pdbid = os.path.splitext(os.path.basename(hitfile))[0].lower()
    obj_name = '%s hits' % pdbid

    # supress view reset on load
    cur_view = viewing.get_view(0)
    cmd.delete(obj_name)
    cmd.load_cgo(hit_objs, obj_name)
    viewing.set_view(cur_view)
    
    if hf.hits:
        histogram([h.score for h in hf.hits], 10)


# =================================================================
def histogram(alist, numbins):
    minval = min(alist)
    maxval = max(alist)
    binsize = float(maxval-minval)/numbins
    bins = [0]*numbins

    for item in alist:
        idx = int((item-minval)/binsize)
        if idx == numbins: idx = numbins-1
        bins[idx] += 1

    print 'Min:',minval, 'Max:',maxval
    for idx in range(numbins):
        print '%6.2f-%6.2f: %d' % (minval+idx*binsize, minval+(idx+1)*binsize, bins[idx])

    return bins, minval, binsize


# =================================================================
def viewhits(hitfile, cutoff='', color='(1.0,0.0,0.0)', radius='0.3'):
    '''
    DESCRIPTION

    "viewhits" will visualize hitfiles from Feature.  By default, it
    will display the top 100 scores.  Color can be specified by name
    or parenthesized, comma separted RGB tuple.  There is no way to
    specify "top X" scores yet, only cutoff scores.  

    USAGE

    viewhits <hitfile> [, <cutoff>] [, <color>] [, <radius>]
    '''
    
    if not cutoff: cutoff = None
    else: cutoff = eval(cutoff)

    # Handle color lookup
    # color_tuple = querying.get_color_tuple(color)
    color_tuple = cmd._cmd.get_color(color,0)
    if color_tuple != None:
        color = color_tuple
    else:
        try:
            color = eval(color)
        except NameError:
            print "Unknown color: %s" % color
            return
    
    radius = eval(radius)
    _viewhits(hitfile, cutoff, color, radius)


# =================================================================
if __name__ == 'pymol':
    cmd.extend('viewhits',viewhits)