## \example pmi/symmetry.py
"""Clone molecules and use a symmetry constraint
"""

import IMP
import IMP.atom
import IMP.rmf
import IMP.pmi
import IMP.pmi.topology
import IMP.pmi.dof
import IMP.pmi.macros
import IMP.pmi.restraints.stereochemistry
import math
import sys

IMP.setup_from_argv(sys.argv, "Symmetry constraint example")

# Create System and State
mdl = IMP.Model()
s = IMP.pmi.topology.System(mdl)
st = s.create_state()

# Create a simple all-bead molecule
mol = st.create_molecule("mymol", sequence='A'*10, chain_id='A')
mol.add_representation(mol, resolutions=[1])

# Clone the molecule multiple times
# Calling molecule.create_clone makes a new molecule with the same name,
# sequence, initial structure, and choice of representations
# Note: another function, molecule.create_copy(), just copies the name
# and sequence
mols = [mol]
chains = 'BCDEFGHI'
for nc in range(7):
    clone = mol.create_clone(chains[nc])
    mols.append(clone)

hier = s.build()

# Create a symmetry constraint
#  A constraint is invariant: IMP will automatically move all clones to
#  match the reference
#  If instead you want some more flexibility, consider
#  IMP.pmi.restraints.stereochemistry.SymmetryRestraint
dof = IMP.pmi.dof.DegreesOfFreedom(mdl)
center = IMP.algebra.Vector3D([50, 0, 0])
for nc in range(7):
    rot = IMP.algebra.get_rotation_about_axis([0, 0, 1], 2*math.pi*(nc+1)/8)
    transform = IMP.algebra.get_rotation_about_point(center, rot)
    dof.constrain_symmetry(mols[0], mols[nc+1], transform)
mdl.update()  # propagates coordinates


# ########### Make stuff look cool with restraints ###########

# set up the original molecule as flexible beads
dof.create_flexible_beads(mols[0])

# Create a connectivity restraint for the first molecule
cr = IMP.pmi.restraints.stereochemistry.ConnectivityRestraint(objects=mol)
cr.add_to_model()

# Create excluded volume for all particles
evr = IMP.pmi.restraints.stereochemistry.ExcludedVolumeSphere(
    included_objects=mols)
evr.add_to_model()

# Quickly move all flexible beads into place
dof.optimize_flexible_beads(100)

# write a single-frame RMF to view the helix
out = IMP.pmi.output.Output()
out.init_rmf("example_symmetry.rmf3", hierarchies=[hier])
out.write_rmf("example_symmetry.rmf3")
