from modeller.automodel import *
log.verbose()
log.level(1, 1, 1, 1, 0)
env = environ()
env.io.atom_files_directory = './:../atom_files'

a = automodel(env,
              alnfile      = '*.pir',
              knowns   = 'xxx',
              sequence = 'vvv')

a.starting_model= 1
a.endind_model  = 1
a.make()

def defsym(mdl, aln, seg1, seg2):
    for (set, seg) in [(2, seg1), (3, seg2)]:
        mdl.pick_atoms(aln, pick_atoms_set=set, selection_segment=seg,
                       atom_types='ALL', selection_status='INITIALIZE',
                       selection_search='SEGMENT')
    mdl.symmetry.define(symmetry_weight=1.0, add_symmetry=(True, False))

def defsym(mdl, aln, seg1, seg3):
    for (set, seg) in [(2, seg1), (4, seg3)]:
        mdl.pick_atoms(aln, pick_atoms_set=set, selection_segment=seg,
                       atom_types='ALL', selection_status='INITIALIZE',
                       selection_search='SEGMENT')
    mdl.symmetry.define(symmetry_weight=1.0, add_symmetry=(True, False))

def defsym(mdl, aln, seg1, seg4):
    for (set, seg) in [(2, seg1), (5,seg4)]:
        mdl.pick_atoms(aln, pick_atoms_set=set, selection_segment=seg,
                       atom_types='ALL', selection_status='INITIALIZE',
                       selection_search='SEGMENT')
    mdl.symmetry.define(symmetry_weight=1.0, add_symmetry=(True, False))

def defsym(mdl, aln, seg1, seg5):
    for (set, seg) in [(2, seg1), (6, seg5)]:
        mdl.pick_atoms(aln, pick_atoms_set=set, selection_segment=seg,
                       atom_types='ALL', selection_status='INITIALIZE',
                       selection_search='SEGMENT')
    mdl.symmetry.define(symmetry_weight=1.0, add_symmetry=(True, False))


# Generate five copies of a segment:
code = 'xxx'
mdl = model(env, file=code, model_segment=('1:A', '211:E'))
aln = alignment(env)
aln.append_model(mdl, align_codes=code, atom_files=code)
aln.append_model(mdl, align_codes=code+'_ini', atom_files=code+'_ini')
mdl.generate_topology(aln, sequence=code+'_ini')
mdl.transfer_xyz(aln)
mdl.build(initialize_xyz=False, build_method='INTERNAL_COORDINATES')
mdl.rename_segments(segment_ids=('A', 'B'), renumber_residues=(1, 1))
mdl.rename_segments(segment_ids=('B', 'C'), renumber_residues=(1, 1))
mdl.rename_segments(segment_ids=('C', 'D'), renumber_residues=(1, 1))
mdl.rename_segments(segment_ids=('D', 'E'), renumber_residues=(1, 1))

myedat = energy_data(dynamic_sphere = False)
mdl.energy(edat=myedat)
mdl.randomize_xyz(deviation=6.0)

# Define the two segments (chains in this case) to be identical:
defsym(mdl, aln, seg1=('1:A', '211:A'), seg2=('1:B', '211:B'))
defsym(mdl, aln, seg1=('1:A', '211:A'), seg3=('1:C', '211:C'))
defsym(mdl, aln, seg1=('1:A', '211:A'), seg4=('1:D', '211:D'))
defsym(mdl, aln, seg1=('1:A', '211:A'), seg5=('1:E', '211:E'))

# Make them identical by optimizing the initial randomized structure
# without any other restraints:
mdl.energy(edat=myedat)
mdl.write(file='define_symmetry-1.atm')
mdl.optimize(max_iterations=300, edat=myedat)
mdl.write(file='define_symmetry-2.atm')
mdl.energy(edat=myedat)

# Now optimize with stereochemical restraints so that the
# result is not so distorted a structure (still distorted
# because optimization is not thorough):
myedat.dynamic_sphere = True
mdl.restraints.make(aln, restraint_type='stereo', spline_on_site=False,
                    edat=myedat)
mdl.randomize_xyz(deviation=3.0)
for method in (1, 3, 1):   # 1 = conjugate gradients, 3 = molecular dynamics
    mdl.optimize(max_iterations=300, md_return='FINAL', edat=myedat,
                 optimization_method=method, output='REPORT')
mdl.write(file='define_symmetry-3.atm')
mdl.energy(edat=myedat)

# Create a blank alignment so that superpose uses its 1:1 default
aln = alignment(env)

mdl = model(env, file='define_symmetry-3.atm', model_segment=('1:A', '211:A'))
mdl2 = model(env, file='define_symmetry-3.atm', model_segment=('1:B', '211:B'))
mdl3 = model(env, file='define_symmetry-3.atm', model_segment=('1:C', '211:C'))
mdl4 = model(env, file='define_symmetry-3.atm', model_segment=('1:D', '211:D'))
mdl5 = model(env, file='define_symmetry-3.atm', model_segment=('1:E', '211:E'))
mdl.pick_atoms(aln, atom_types='MNCH')
mdl.superpose(mdl2, mdl3, mdl4, mdl5, aln)