import sys

# START:readPdb
def readPdb(inputName):
    input = open(inputName, 'r')
    result = []
    for line in input:
        if line[:4] == 'ATOM':
            fields = line.split()
            atom = [fields[2]] + fields[4:7]
            result.append(atom)
    input.close()
    return result
# END:readPdb

def normalizeSymbols(atoms):
    for record in atoms:
        record[0] = record[0].capitalize()

def translateName(inputName):
    return inputName[:-4] + '.w3d'

def truncateName(inputName):
    inputName = inputName.capitalize()
    return inputName[:-4]

Lookup = {
    'Br' : (2, 0.651, 0.161, 0.161),
    'C'  : (3, 0.0, 1.0, 0.0),
    'Cl' : (8, 0.122, 0.941, 0.122),
    'Fe' : (13, 0.878, 0.4, 0.2),
    'H'  : (6, 1.0, 1.0, 1.0),
    'N'  : (17, 0.0, 0.0, 1.0),
    'O'  : (19, 1.0, 0.0, 0.0),
    'S'  : (17, 1.0, 1.0, 0.0)
}

def writeHtml(outputName, atoms):
    output = open('mol_name.sed', 'w')
    print >> output, "s/__mol_name__/{}/".format(molName)
    output.close()

    output = open(outputName, 'w')
#
#           <transform translation='-3.194 1.453 -0.615'> 
#           <shape> 
#               <appearance> 
#               <material diffuseColor='0 1 0'></material> 
#               </appearance> 
#               <sphere></sphere> 
#           </shape> 
#           </transform> 
#
    scale=0.07
    for (symbol, X, Y, Z) in atoms:
        if symbol not in Lookup:
            print >> sys.stderr, 'Unknown atom "%s"' % symbol
            sys.exit(1)
        radius,red,grn,blu = Lookup[symbol]
        print >> output, "            <transform translation=\'",X, Y, Z,"\'>"
        print >> output, "            <shape>"
        print >> output, "                <appearance>"
        print >> output, "                <material diffuseColor=\'",red,grn,blu,"\'></material>"
        print >> output, "                </appearance>"
        print >> output, "                <sphere radius=\'",radius*scale,"\'></sphere>"
        print >> output, "            </shape>"
        print >> output, "            </transform>"
    output.close()

for inputName in sys.argv[1:]:
    atoms = readPdb(inputName)
    normalizeSymbols(atoms)
    outputName = translateName(inputName)
    molName = truncateName(inputName)
    writeHtml(outputName, atoms)