#!/usr/bin/env python

from mmCif import mmcifIO
import os
import cPickle as pickle

PWD = os.path.dirname(os.path.realpath(__file__))
DICT_FNAME = os.path.join(PWD, "mmcif_pdbx_v40.dic")
PICKLE_FNAME = DICT_FNAME[:DICT_FNAME.rfind(".")] + ".pkl"

residue_map_1to3 = {
    'A': 'ALA', 'C': 'CYS', 'D': 'ASP', 'E': 'GLU', 'F': 'PHE', 'G': 'GLY',\
    'H': 'HIS', 'I': 'ILE', 'K': 'LYS', 'L': 'LEU', 'M': 'MET', 'N': 'ASN',\
    'P': 'PRO', 'Q': 'GLN', 'R': 'ARG', 'S': 'SER', 'T': 'THR', 'V': 'VAL',\
    'W': 'TRP', 'X': 'UNK', 'Y': 'TYR',
}

residue_map_3to1 = {v: k for k, v in residue_map_1to3.items()}

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

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

    return f


def get_mmcif(fname):
    infile = mmcifIO.CifFileReader(input="data")
    f = infile.read(fname, output="cif_file")
    return f

def get_category_items(dictionary, category):
    results = []

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

    return results

def get_category(cif_file, category):
    for db in cif_file.getDataBlocks():
        if category in db.getCategoryIds():
            return db.getCategory(category)

    return None

def first_available(entities):
    for idx in range(1, 100000):
        if idx not in entities:
            return idx

    return None

def get_seq_alias_map(aliases, mmcif):
    out = {}

    # get non-polymers only
    _, entities = split_entities(mmcif)

    for a in aliases:
        if a not in out:
            ent = first_available(entities)
            entities.append(ent)
            out[a] = ent

    return out

def split_entities(mmcif):
    cat = get_category(mmcif, "entity")
    types = cat.getItem("type").value
    poly = []
    nonpoly = []

    for idx, t in enumerate(types):
        if t == "polymer":
            poly.append(int(cat.getItem("id").value[idx]))
        else:
            nonpoly.append(int(cat.getItem("id").value[idx]))

    return (poly, nonpoly)

def get_atom_site_seq(mmcif, entity):
    cat = get_category(mmcif, "atom_site")
    entities = cat.getItem("label_entity_id").value
    seq = {}

    for idx, e in enumerate(entities):
        if e != str(entity):
            continue

        seq_id = int(cat.getItem("label_seq_id").value[idx])

        if seq_id not in seq.keys():
            threeL = cat.getItem("label_comp_id").value[idx]
            if threeL in residue_map_3to1:
                seq[seq_id] = residue_map_3to1[threeL]
            else:
                seq[seq_id] = "X"

    return "".join([seq[key] for key in sorted(seq.keys())])

    #for key in sorted(seq.keys()):
    #    print key,
    #    threeL = seq[key]
    #    if threeL in residue_map_3to1:
    #        print residue_map_3to1[seq[key]]
    #    else:
    #        print 'X'
    
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)

def print_rows(cat):
    size = len(cat.getItems()[0].value)
    names = cat.getItemNames()
    result = [names]

    for i in range(size):
        tmp = []
        for n in names:
            tmp.append(cat.getItem(n).value[i])
        result.append(tmp)

    pretty_print(result)