import sys
import numpy as np
import os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

BACKBONE = ['N', 'CA', 'C']

def load_backbone(fname):
    """
    Return a list of (ResNum, Atom Name, (x, y, z))
       for all 'ATOM' entries in FNAME.
    """
    residues = {}
    with open(fname) as fp:
        for line in fp:
            # PDB files have fixed width columns!
            if len(line) < 4 or line[:4] != 'ATOM': continue
            atom = line[12:16].strip()
            if atom not in BACKBONE: continue
            resnum = int(line[22:26].strip())
            x = float(line[30:38])
            y = float(line[38:46])
            z = float(line[46:54])
            if resnum not in residues: residues[resnum] = {}
            residues[resnum][atom] = np.array([x, y, z])
    return residues


def dihedral(p0, p1, p2, p3):
    # FROM: https://stackoverflow.com/questions/20305272/
    """Praxeolitic formula
    1 sqrt, 1 cross product"""

    b0 = -1.0*(p1 - p0)
    b1 = p2 - p1
    b2 = p3 - p2

    # normalize b1 so that it does not influence magnitude of vector
    # rejections that come next
    b1 /= np.linalg.norm(b1)

    # vector rejections
    # v = projection of b0 onto plane perpendicular to b1
    #   = b0 minus component that aligns with b1
    # w = projection of b2 onto plane perpendicular to b1
    #   = b2 minus component that aligns with b1
    v = b0 - np.dot(b0, b1)*b1
    w = b2 - np.dot(b2, b1)*b1

    # angle between v and w in a plane is the torsion angle
    # v and w may not be normalized but that's fine since tan is y/x
    x = np.dot(v, w)
    y = np.dot(np.cross(b1, v), w)
    return np.degrees(np.arctan2(y, x))

# Load backbone atoms
residues = load_backbone(sys.argv[1])
psi, phi = [], []
for resnum, coords in residues.items():
    if resnum-1 in residues and resnum+1 in residues:
        prev = residues[resnum-1]
        next = residues[resnum+1]

        psi += [dihedral(coords['N'], coords['CA'], coords['C'], next['N'])]
        phi += [dihedral(prev['C'], coords['N'], coords['CA'], coords['C'])]

pdbname = os.path.basename(sys.argv[1])
plt.title(pdbname + ' Ramachandran Plot')
plt.scatter(phi, psi)
plt.gca().set_aspect('equal', adjustable='box')
plt.ylim(-180, 180)
plt.xlim(-180, 180)
plt.xlabel('PHI')
plt.ylabel('PSI')
plt.savefig(pdbname + '.rama.png')
