# Example for: model.restraints.add(), model.restraints.unpick()

# This will enforce cis conformation for Pro-56.

# Make a model and stereochemical restraints:

def cispeptide(rsr, ids1, ids2):
    # Delete the old restraint on the same atoms:
    rsr.unpick(atom_ids=ids1)
    # Add the new restraint:
    rsr.add(atom_ids=ids1, restraint_parameters=(3, 1, 3, 3, 4, 2, 0,
                                                 3.141593, 0.087))

    rsr.unpick(atom_ids=ids2)
    rsr.add(atom_ids=ids2, restraint_parameters=(3, 1, 3, 3, 4, 2, 0, 0.0,
                                                 0.087))

log.level(output=1, notes=1, warnings=1, errors=1, memory=0)
env = environ()

env.libs.topology.read(file='$(LIB)/top_heav.lib')
env.libs.parameters.read(file='$(LIB)/par.lib')

aln = alignment(env)
code = '1fas'
mdl = model(env, file=code)
aln.append_model(mdl, align_codes=code, atom_files=code)
aln.append_model(mdl, align_codes='1fas-ini', atom_files='1fas.ini')

mdl.generate_topology(aln=aln, sequence='1fas-ini', add_segment=False)
mdl.transfer_xyz(aln=aln)
mdl.build(initialize_xyz=False, build_method='INTERNAL_COORDINATES')
mdl.restraints.make(aln=aln, restraint_type='stereo', spline_on_site=False)

# Change the Pro-56 restraint from trans to cis:
cispeptide(mdl.restraints, ids1=( 'O:56', 'C:56', 'N:57', 'CA:57'),
                           ids2=('CA:56', 'C:56', 'N:57', 'CA:57'))

mdl.restraints.write(file='1fas.rsr')
mdl.energy()