#!/usr/bin/env python
# -*- coding:utf-8 -*-

u"""Gō plot of residue-to-residue distances in a protein structure.

Usage:
    go1981.py <PDB_file> [<chain_ID>] [<junction_position> ...]

Example:
    ./go1981.py 3ONZ.pdb B 30 68 104

See:

Mitiko Gō
Correlation of DNA exonic regions with protein structural units in haemoglobin
Nature 291, 90 - 92 (07 May 1981)
doi:10.1038/291090a0

"""

import sys

from Bio.PDB import PDBParser

try:
    import matplotlib.pyplot as plt
except ImportError:
    import pylab as plt


def posn(residue):
    atoms = residue.child_dict
    # Beta carbon is a better representative of side chain position, but
    # glycine has no beta carbon atom
    return (atoms['CB'] if 'CB' in atoms else atoms['CA'])


def go1981plot(fname, chain=None, threshold=None, junctions=(), do_show=True):
    struct = PDBParser().get_structure(fname, fname)
    thischain = (struct[0][chain] if chain
            else struct[0].child_list[0])
    # Filter for amino acids (remove hetatms)
    residues = tuple(r for r in thischain if 'CA' in r.child_dict)
    distances = {}
    for y, resA in enumerate(residues):
        distances.update(
                # Lower-left triangle
                dict(((x,y), posn(resA) - posn(resB))
                     for x, resB in enumerate(residues[:y])))

    # ENH: use shading to indicate greater distances
    if not threshold:
        threshold = max(distances.itervalues())/2.0 + 2.5
        sys.stderr.write("Calculated radius: %f\n" % threshold)
    coords = ((xy[0] + 1, xy[1] + 1)
            for xy, dist in distances.iteritems()
            if dist > threshold)
    for x, y in coords:
        plt.plot(x, y, 'ko', alpha=.2)

    # Draw the diagonal
    plt.plot([1,len(residues)], [1,len(residues)], 'k-')
    # Draw the junctions
    for junc in junctions:
        plt.vlines(junc, junc, len(residues), color='k')
        plt.hlines(junc, 1, junc, color='k')
        plt.text(junc, junc, str(junc), fontsize=12)
    # Set the axes
    plt.xlim(1, len(residues))
    # Invert the y-axis (1 at the top)
    plt.ylim(len(residues), 1)
    # ENH: remove the right and top axis lines & ticks
    plt.title(fname)
    plt.xlabel('Residue number')
    plt.ylabel('Residue number')
    if do_show:
        plt.show()


if __name__ == '__main__':
    fname = sys.argv[1]
    if len(sys.argv) == 2:
        go1981plot(sys.argv[1])
    elif sys.argv[2].isalpha():
        go1981plot(sys.argv[1], sys.argv[2], junctions=map(int, sys.argv[3:]))
    else:
        go1981plot(sys.argv[1], junctions=map(int, sys.argv[2:]))