#!/usr/bin/env python

from mmCif import mmcifIO
import sys
import cPickle as pickle
import os

DICT = "mmcif_pdbx_v40.dic"
PICKLE = DICT[:DICT.rfind(".")] + ".pkl"
CATEGORY = sys.argv[1]

def pretty_print(data):
    padding = 2
    col_width = max(len(word) for row in data for word in row) + padding

    for row in data:
        print "".join(word.ljust(col_width) for word in row)

results = []

if os.path.isfile(PICKLE):
    with open(PICKLE, "rb") as inputf:
        f = pickle.load(inputf)
else:
    infile = mmcifIO.CifFileReader(input='dictionary')
    f = infile.read(DICT)

    with open(PICKLE, "wb") as output:
        pickle.dump(f, output, pickle.HIGHEST_PROTOCOL)

for f in f.getDataBlocks()[0].getSaveFrames():
    x = f.setCategory("item").setItem("category_id")
    if x and x.value == CATEGORY:
        item = f.getCategory("item")
        results.append((item.getItem("name").value,
                        item.getItem("mandatory_code").value))

pretty_print(results)