#! /usr/bin/env python
import collections
import sys

Atom = collections.namedtuple('Atom', ['x', 'y', 'z'])
Residue = collections.namedtuple('Residue', ['index', 'name', 'atoms'])

with open(sys.argv[1]) as pdbfile:
    residues = dict()

    interesting = {
        'A': ('N1', 'N6'),
        'C': ('N3', 'O2', 'N4'),
        'G': ('N1', 'N2', 'O6'),
        'U': ('N3', 'O2')
    }

    for line in pdbfile.readlines():
        if line.startswith('ATOM') or line.startswith('HETATM'):
            name = line[12:16].strip()
            resname = line[17:20].strip()

            if not resname in interesting or not name in interesting[resname]:
                continue

            x = float(line[30:38].strip())
            y = float(line[38:46].strip())
            z = float(line[46:54].strip())
            atom = Atom(x, y, z)

            index = int(line[22:26].strip())
            if not index in residues:
                residues[index] = Residue(index, resname, dict())

            atoms = residues[index].atoms
            atoms[name] = atom

    print(residues)