from pymol import cmd
from pymol.cgo import *
 
def dump2surfaceCGO():
    CGOobj = []
    dumpedFile = open("dump.tmp").read()
    for block in dumpedFile.split('\n\n'):
        CGOobj.append(BEGIN)
        CGOobj.append(TRIANGLES)
 
        for line in block.split('\n'):
            if line == '':
                continue
 
            vals = line.split()
            CGOobj.append(NORMAL)
            CGOobj.append(float(vals[3]))
            CGOobj.append(float(vals[4]))
            CGOobj.append(float(vals[5]))
            CGOobj.append(VERTEX)
            CGOobj.append(float(vals[0]))
            CGOobj.append(float(vals[1]))
            CGOobj.append(float(vals[2]))
 
        CGOobj.append(END)
    return CGOobj
 
def dump2meshCGO():
    CGOobj = []
    dumpedFile = open("dump.tmp").read()
    for block in dumpedFile.split('\n\n'):
        CGOobj.append(BEGIN)
        CGOobj.append(LINE_STRIP)
 
        for line in block.split('\n'):
            if line == '':
                continue
 
            CGOobj.append(VERTEX)
            vals = line.split()
 
            CGOobj.append(float(vals[0]))
            CGOobj.append(float(vals[1]))
            CGOobj.append(float(vals[2]))
 
        CGOobj.append(END)
    return CGOobj
 
def getType(objname):
    session = cmd.get_session()['names']
    for obj in session:
        if obj == None:
            continue
        if obj[0] != objname:
            continue
        return obj[4]
    return -1
 
 
def dump2CGO(obj):
    cmd.dump("dump.tmp", obj)
    type = getType(obj)
    cgo = []
    if (type == 3): # Mesh
        cgo = dump2meshCGO()
    elif (type == 7): #Surface
        cgo = dump2surfaceCGO()
    else:
        print "Unknown type"
        return
 
    cmd.load_cgo(cgo, "CGO " + obj)
 
cmd.extend('dump2CGO', dump2CGO)
cmd.auto_arg[0]['dump2CGO'] = [cmd.object_sc, 'object', '']