from __future__ import print_function
from modeller import *
from modeller.scripts import complete_pdb

log.verbose()
env = Environ()
env.libs.topology.read('${LIB}/top_heav.lib')
env.libs.parameters.read('${LIB}/par.lib')

# cryo-EM map details
map='TvLDH.10A.mrc'  
resolution=10.0  
box_size=48 
apix=1.88
x=-26.742; y=-9.5205; z=-10.375 #origin

# Read in already-fitted PDB file and select all C-alpha atoms
m1 = Model(env, file='TvLD_1_1.pdb')
sel = Selection(m1).only_atom_types('CA')

scal = physical.Values(default=0., em_density=1.0)
results = []

# Loop over all five output PDB files
for i in range(1, 6):
    # Read in the output PDB file
    fname = '../step4_model/TvLDH.B9999%04d.pdb' % i
    m2 = complete_pdb(env, fname)

    # Make a 1:1 alignment of the already-fitted file and this model
    aln = Alignment(env)
    aln.append_model(m1, align_codes='fitted')
    aln.append_model(m2, align_codes='model')

    # Structural superposition using sel (all C-alpha positions)
    sel.superpose(m2, aln)

    # Read in cryo-EM map
    den = Density(env, file=map, em_density_format='MRC',
                  voxel_size=apix, resolution=resolution, em_map_size=box_size,
                  density_type='GAUSS', px=x,py=y,pz=z)

    # Add the EM density fit score (negative of CCF) to the MODELLER
    # scoring function for this model
    m2.env.edat.density = den

    # Calculate scoring function (energy) for all atoms in m2, with all
    # scoring function terms except the EM density fit score turned off
    # (scaled to zero)
    molpdf, terms = Selection(m2).energy(schedule_scale=scal)
    results.append((fname, -molpdf))

# Print all results
print("\n%-40s %12s" % ('Filename', 'CCF'))
print("\n".join(['%-40s %12.4f' % x for x in results]))
