from pymol import cmd
from pymol import stored

def list_hb(selection, selection2=None, cutoff=3.2,
        angle=55, mode=1, hb_list_name='hbonds'):
    """
    USAGE
    
    list_hb needs a selection, not an object!
    list_hb selection,
    [selection2 (default=None)],
    [cutoff (default=3.2)],
    [angle (default=55)],
    [hb_list_name]

    The script automatically adds a requirement that
    atoms in the selection (and selection2 if used)
    must be either of the elements N or O.

    If mode is set to 0, then no angle cutoff is used,
    otherwise the angle cutoff is used and defaults to
    55 degrees.

    e.g.
    To get a list of all H-bonds within chain A of an object
    list_hb 1abc & c. a &! r. hoh, cutoff=3.2,
    hb_list_name=abc-hbonds

    To get a list of H-bonds between chain B and everything
    else:
    list_hb 1tl9 & c. b, 1tl9 &! c. b 
    """

    #hb_data = open('./pymolScripts/hb_' + selection + '.dat', 'w')
    #hb_data.write('')
    #hb_data.close()

    #hb_data = open('./pymolScripts/hb_' + selection + '.dat', 'w')

    cutoff=float(cutoff)
    angle=float(angle)
    mode=float(mode)
    # ensure only N and O atoms are in the selection
    selection = selection + " & e. n+o"
    if not selection2:
        hb = cmd.find_pairs(selection, selection, mode=mode,
                cutoff=cutoff, angle=angle)
    else:
        selection2 = selection2 + " & e. n+o"
        hb = cmd.find_pairs(selection, selection2, mode=mode,
                cutoff=cutoff, angle=angle)

    # sort the list for easier reading
    hb.sort(lambda x, y: (cmp(x[0][1], y[0][1]))) 

    stored.listA = []
    stored.listB = []
    stored.listC = []

    for pairs in hb:
        cmd.iterate("%s and index %s" % (pairs[0][0], pairs[0][1]),
            'print "%1s/%3s`%s/%s/%i " % (chain, resn, resi, name, ID),') 
        cmd.iterate("%s and index %s" % (pairs[0][0], pairs[0][1]),
            'stored.listA.append( "%1s/%3s`%s/%s/%i " % (chain, resn, resi, name, ID),)')

        cmd.iterate("%s and index %s" % (pairs[1][0],pairs[1][1]),
                'print "%1s/%3s`%s/%s/%i " % (chain, resn, resi, name, ID),') 
        cmd.iterate("%s and index %s" % (pairs[1][0], pairs[1][1]),
            'stored.listB.append( "%1s/%3s`%s/%s/%i " % (chain, resn, resi, name, ID),)') 

        print "%.2f" % cmd.distance(hb_list_name, "%s and index %s" %
                (pairs[0][0], pairs[0][1]), "%s and index %s" % (pairs[1][0],
                    pairs[1][1]))
        stored.listC.append("%.2f" % cmd.distance(hb_list_name, "%s and index %s" %
                (pairs[0][0], pairs[0][1]), "%s and index %s" % (pairs[1][0],
                    pairs[1][1])))

    for line in enumerate(stored.listA):
        hb_data.write(stored.listA[line[0]]
                + '    '
                + stored.listB[line[0]]
                + '    '
                + stored.listC[line[0]]
                + '\n')

    #hb_data.close()
    
cmd.extend("list_hb", list_hb)