#! /usr/bin/env python
from mmtk import *
import sys
import Units
import Positioner
from Database import BlueprintHook
import pprint
from time import clock
from ChemicalObjects import PDBMolecule

from ConfigIO import *
from DelaunayTriangulation import DelaunayTriangulation

import sys
if 1:
    # Load the PDB sequence
    sequence = PDBFile(sys.argv[1]).readSequenceWithConfiguration()
    if len(sequence) > 0:
	# Construct the peptide chain objects
	chains = map(PeptideChain, sequence)
	# Construct positions of missing hydrogens
	for chain in chains:
	    chain.findHydrogenPositions()
	# Make protein object
	molecule = Protein(chains)
    else:
	molecule = PDBMolecule(sequence.other)
    molecule.normalizePosition()

    delaunay = DelaunayTriangulation(molecule, qhull = 0)
    delaunay.removeFacesBiggerThan(4*Units.Ang*Units.Ang, 1.4*Units.Ang)
    if 0:
	import RasMol
	ras_session = RasMol.RasMol(molecule, 1)
	coord_list = delaunay.coordList()
	for coords in coord_list:
	    ras_session.AddPolyLines(coords)
	ras_session.RefreshScreen()
	ras_session.Interact()
	raise RuntimeError

ff = GenericForceField('mm3')
block1 = Group('diamondoid_fragment')
join = Joiner(block1)
max_dim = 2.75
width = int(5*max_dim)
print 'building',width, width/2, width, 'unit system'
t0 = clock()
if len(sys.argv) > 2:
    if sys.argv[2][-3:] == 'pdb':
	sequence2 = PDBFile(sys.argv[2]).readSequenceWithConfiguration()
	result = PDBMolecule(sequence2.other, forcefield = ff)
	for i in range(len(result.atoms)):
	    a = result.atoms[i]
	    result.setIndex(a, i)
	    result.setAtomProperty(a, 'mm3_atom_type', 1)
	result.constructBondsToAtom()
	print 'constructBondsToAtom'
	sys.stdout.flush()
	result.calcDanglingBonds(ff)
	print 'calcDanglingBonds'
	sys.stdout.flush()
	save(result, 'ptestm.out')
    else:
	result = load(sys.argv[2])
    t1 = t2 = t0
else:
    join.repeatGroup('hookeast', 'hookwest', width)
    result = join.getResult()
    join = Joiner(result)
    join.repeatGroup('hooksouth', 'hooknorth', int(4*max_dim), 1,\
		     [180*Units.deg]*width)
    t1 = clock()
    result = join.getResult()
    join = Joiner(result)
    join.repeatGroup('hookup', 'hookdown', int(5*max_dim))
    t2 = clock()
    result = join.getResult()
    del join
    del block1
t3 = clock()
try:
    if not Utility.isDefinedPosition(result.atoms[0].position()):
	result.setPositions(ff)
    result.normalizePosition()
    t4 = clock()
    pos_list = []
    for a in result.atoms:
	pos_list.append(a.position())
        #if not delaunay.contains(a.position()):
	#    remove_atoms.append(a)
    keep_atoms_index = delaunay.listContains(pos_list)
    print 'keep',len(keep_atoms_index),'atoms'
    sys.stdout.flush()
    keep_dict = {}
    for i in keep_atoms_index:
	keep_dict[id(result.atoms[i])] = 1
    remove_atoms = []
    for a in result.atoms:
	if not keep_dict.has_key(id(a)):
	    remove_atoms.append(a)
    t5 = clock()
    result.removeAtoms(remove_atoms, None)
    print 'removed',len(remove_atoms),'#atoms left',len(result.atoms)
    t6 = clock()
    result.hydrogenateDanglingBonds(ff)
except ConformationError,msg:
    print 'ConformationError',msg
    t4 = t5 = t6 = clock()

from ConfigIO import OutputFile
fmt = OutputFile('PDB/ptest.pdb')
fmt.write(result)
fmt.close()
print 'otimes',t1-t0,t2-t1,t3-t2,t4-t3,'dc',t5-t4,'rm',t6-t5,'H',clock()-t6
sys.stdout.flush()
save(result, 'ptest.out')
if 1:
    import RasMol
    ras_session = RasMol.RasMol(result, 1)
    ras_session.RefreshScreen()
    ras_session.Interact()