from modeller import *
from modeller.automodel import *

log.verbose()    # for debugging

env = environ(rand_seed=-12312) # random seed

env.io.atom_files_directory = '.'  # Working directory

def defsym(mdl, aln, seg1, seg2):
    """Constrain segments |seg1| and |seg2| to be identical"""
    for (set,seg) in [(2,seg1), (3,seg2)]:
        mdl.pick_atoms(aln, pick_atoms_set=set, selection_segment=seg,
                       atom_types='CA', selection_status='INITIALIZE',
                       selection_search='SEGMENT')
    mdl.symmetry.define(symmetry_weight=1.0, add_symmetry=(True, False))


def makealpha(mdl, aln, residue_ids):
    """Enforce alpha-helical structure on |residue_ids|"""
    mdl.restraints.make(aln, restraint_type='ALPHA', residue_ids=residue_ids,
                        spline_on_site=False)


class mymodel(automodel):
    def special_restraints(self, aln):
        segs = [('1', '99'), ('100', '198'), ('199', '297'), ('298', '396')]
        for n in range(len(segs)):
            defsym(self, aln, segs[n-1], segs[n])

        makealpha(self, aln, ('5', '30'))  	#M1 region (consensus)
        makealpha(self, aln, ('40', '53')) 	#P region (jalview p/out)
        makealpha(self, aln, ('64', '95')) 	#M2 region (consensus)
        makealpha(self, aln, ('104', '129'))
        makealpha(self, aln, ('139', '152'))
        makealpha(self, aln, ('163', '194'))
	makealpha(self, aln, ('203', '228'))
        makealpha(self, aln, ('238', '251'))
        makealpha(self, aln, ('262', '293'))
	makealpha(self, aln, ('302', '327'))
        makealpha(self, aln, ('337', '350'))
        makealpha(self, aln, ('361', '392'))


a = mymodel(env,
            alnfile  = 'x_vs_y.ali',  # The alignment filename
            knowns   = ('template'),  # The structure filenames, the pdb
                                     # is called filename.atm
            sequence = 'target')      # The sequence filename
a.starting_model= 1                  # Starting and ending model indices
a.ending_model  = 4
a.deviation = 4.0                    # Deviation in models

# Call the routine model which in turn calls all the other routines required.
a.make()

mdl = model(env, file = ???)
mdl.energy()
molpdf = mdl.energy(output='VERY_LONG', residue_span_range=(1,9999), normalize_profile=True, edat=energy_data(dynamic_sphere=True))
mdl.seq_id = 90.00
score = mdl.assess_ga341()