/mmcif.py

"""@namespace IMP.pmi1.mmcif
   @brief Support for the mmCIF file format.

   IMP has basic support for writing out files in mmCIF format, for
   deposition in [PDB-Dev](https://pdb-dev.wwpdb.org/).
   mmCIF files are currently generated by creating an
   IMP.pmi1.mmcif.ProtocolOutput class, and attaching it to an
   IMP.pmi1.representation.Representation object, after which any
   generated models and metadata are collected and output as mmCIF.
"""

from __future__ import print_function
import copy
import IMP.core
import IMP.algebra
import IMP.atom
import IMP.em
import IMP.isd
import IMP.pmi1.representation
import IMP.pmi1.tools
from IMP.pmi1.tools import OrderedDict
import IMP.pmi1.output
import IMP.pmi1.metadata
import IMP.mmcif.data
import re
import ast
import sys
import os
import textwrap
import weakref
import operator
import itertools
import ihm.location
import ihm.dataset
import ihm.dumper
import ihm.metadata
import ihm.startmodel
import ihm.model
import ihm.protocol
import ihm.analysis
import ihm.representation
import ihm.geometry

def _assign_id(obj, seen_objs, obj_by_id):
    """Assign a unique ID to obj, and track all ids in obj_by_id."""
    if obj not in seen_objs:
        if not hasattr(obj, 'id'):
            obj_by_id.append(obj)
            obj.id = len(obj_by_id)
        seen_objs[obj] = obj.id
    else:
        obj.id = seen_objs[obj]


def _get_by_residue(p):
    """Determine whether the given particle represents a specific residue
       or a more coarse-grained object."""
    return IMP.atom.Residue.get_is_setup(p) or IMP.atom.Atom.get_is_setup(p)


class _ComponentMapper(object):
    """Map a Particle to a component name"""
    def __init__(self, prot):
        self.o = IMP.pmi1.output.Output()
        self.prot = prot
        self.name = 'cif-output'
        self.o.dictionary_pdbs[self.name] = self.prot
        self.o._init_dictchain(self.name, self.prot, multichar_chain=True)

    def __getitem__(self, p):
        protname, is_a_bead = self.o.get_prot_name_from_particle(self.name, p)
        return protname


class _AsymMapper(object):
    """Map a Particle to an asym_unit"""
    def __init__(self, simo, prot):
        self.simo = simo
        self._cm = _ComponentMapper(prot)

    def __getitem__(self, p):
        protname = self._cm[p]
        return self.simo.asym_units[protname]

    def get_feature(self, ps):
        """Get an ihm.restraint.Feature that covers the given particles"""
        # todo: handle things other than residues
        rngs = []
        for p in ps:
            asym = self[p]
            # todo: handle overlapping or adjoining ranges
            if IMP.atom.Residue.get_is_setup(p):
                rind = IMP.atom.Residue(p).get_index()
                rngs.append(asym(rind, rind))
            elif IMP.atom.Fragment.get_is_setup(p):
                # PMI Fragments always contain contiguous residues
                rinds = IMP.atom.Fragment(p).get_residue_indexes()
                rngs.append(asym(rinds[0], rinds[-1]))
            else:
                raise ValueError("Unsupported particle type %s" % str(p))
        return ihm.restraint.ResidueFeature(rngs)


class _AllSoftware(object):
    def __init__(self, system):
        self.system = system
        self.modeller_used = self.phyre2_used = False
        self.system.software.extend([
           ihm.Software(
                 name="Integrative Modeling Platform (IMP)",
                 version=IMP.__version__,
                 classification="integrative model building",
                 description="integrative model building",
                 location='https://integrativemodeling.org'),
           ihm.Software(
                name="IMP PMI module",
                version=IMP.pmi1.__version__,
                classification="integrative model building",
                description="integrative model building",
                location='https://integrativemodeling.org')])

    def set_modeller_used(self, version, date):
        if self.modeller_used:
            return
        self.modeller_used = True
        self.system.software.append(ihm.Software(
                    name='MODELLER', classification='comparative modeling',
                    description='Comparative modeling by satisfaction '
                                'of spatial restraints, build ' + date,
                    location='https://salilab.org/modeller/',
                    version=version))

    def set_phyre2_used(self):
        if self.phyre2_used:
            return
        self.phyre2_used = True
        self.system.software.append(ihm.Software(
                   name='Phyre2', classification='protein homology modeling',
                   description='Protein Homology/analogY Recognition '
                               'Engine V 2.0',
                   version='2.0',
                   location='http://www.sbg.bio.ic.ac.uk/~phyre2/'))

def _get_fragment_is_rigid(fragment):
    """Determine whether a fragment is modeled rigidly"""
    leaves = IMP.atom.get_leaves(fragment.hier)
    # Assume all leaves are set up as rigid/flexible in the same way
    return IMP.core.RigidMember.get_is_setup(leaves[0])

class _PDBFragment(ihm.representation.ResidueSegment):
    """Record details about part of a PDB file used as input
       for a component."""
    def __init__(self, state, component, start, end, pdb_offset,
                 pdbname, chain, hier, asym_unit):
        # start, end are PMI residue indexes (not IHM)
        super(_PDBFragment, self).__init__(
                   asym_unit=asym_unit.pmi_range(start, end),
                   rigid=None, primitive='sphere')
        self.component, self.start, self.end, self.offset, self.pdbname \
              = component, start, end, pdb_offset, pdbname
        self.state, self.chain, self.hier = state, chain, hier
        sel = IMP.atom.NonWaterNonHydrogenPDBSelector() \
              & IMP.atom.ChainPDBSelector(chain)
        self.starting_hier = IMP.atom.read_pdb(pdbname, state.model, sel)

    rigid = property(lambda self: _get_fragment_is_rigid(self),
                     lambda self, val: None)

    def combine(self, other):
        pass


class _BeadsFragment(ihm.representation.FeatureSegment):
    """Record details about beads used to represent part of a component."""
    chain = None
    def __init__(self, state, component, start, end, count, hier, asym_unit):
        super(_BeadsFragment, self).__init__(asym_unit=asym_unit(start, end),
                    rigid=None, primitive='sphere', count=count)
        self.state, self.component, self.hier \
              = state, component, hier

    rigid = property(lambda self: _get_fragment_is_rigid(self),
                     lambda self, val: None)

    def combine(self, other):
        # todo: don't combine if one fragment is rigid and the other flexible
        if type(other) == type(self) and \
          other.asym_unit.seq_id_range[0] == self.asym_unit.seq_id_range[1] + 1:
            self.asym_unit.seq_id_range = (self.asym_unit.seq_id_range[0],
                                           other.asym_unit.seq_id_range[1])
            self.count += other.count
            return True


class _AllModelRepresentations(object):
    def __init__(self, simo):
        self.simo = simo
        # dict of fragments, ordered by representation, then component name,
        # then state
        self.fragments = OrderedDict()
        self._all_representations = {}

    def copy_component(self, state, name, original, asym_unit):
        """Copy all representation for `original` in `state` to `name`"""
        def copy_frag(f):
            newf = copy.copy(f)
            newf.asym_unit = asym_unit(*f.asym_unit.seq_id_range)
            return newf
        for rep in self.fragments:
            if original in self.fragments[rep]:
                if name not in self.fragments[rep]:
                    self.fragments[rep][name] = OrderedDict()
                self.fragments[rep][name][state] = [copy_frag(f)
                            for f in self.fragments[rep][original][state]]
                # Assume representation for a component is the same in all
                # states, so only write out the first one
                first_state = list(self.fragments[rep][name].keys())[0]
                if state is first_state:
                    representation = self._all_representations[rep]
                    representation.extend(self.fragments[rep][name][state])

    def add_fragment(self, state, representation, fragment):
        """Add a model fragment."""
        comp = fragment.component
        id_rep = id(representation)
        self._all_representations[id_rep] = representation
        if id_rep not in self.fragments:
            self.fragments[id_rep] = OrderedDict()
        if comp not in self.fragments[id_rep]:
            self.fragments[id_rep][comp] = OrderedDict()
        if state not in self.fragments[id_rep][comp]:
            self.fragments[id_rep][comp][state] = []
        fragments = self.fragments[id_rep][comp][state]
        if len(fragments) == 0 or not fragments[-1].combine(fragment):
            fragments.append(fragment)
            # Assume representation for a component is the same in all states,
            # so only write out the first one
            first_state = list(self.fragments[id_rep][comp].keys())[0]
            if state is first_state:
                representation.append(fragment)


class _AllDatasets(object):
    """Track all datasets generated by PMI and add them to the ihm.System"""
    def __init__(self, system):
        self.system = system
        self._datasets = []
        self._datasets_by_state = {}
        self._restraints_by_state = {}

    def get_all_group(self, state):
        """Get a DatasetGroup encompassing all datasets so far in this state"""
        # Note that if a restraint dataset is replaced after this is called,
        # the group will still contain the old dataset - mark dataset as read
        # only?
        g = ihm.dataset.DatasetGroup(self._datasets_by_state.get(state, [])
                + [r.dataset for r in self._restraints_by_state.get(state, [])
                   if r.dataset])
        return g

    def add(self, state, dataset):
        """Add a new dataset."""
        self._datasets.append(dataset)
        if state not in self._datasets_by_state:
            self._datasets_by_state[state] = []
        self._datasets_by_state[state].append(dataset)
        # Make sure that the dataset ends up in the mmCIF file even if nothing
        # refers to it
        self.system.orphan_datasets.append(dataset)
        return dataset

    def add_restraint(self, state, restraint):
        """Add the dataset for a restraint"""
        if state not in self._restraints_by_state:
            self._restraints_by_state[state] = []
        self._restraints_by_state[state].append(restraint)


class _CrossLinkRestraint(ihm.restraint.CrossLinkRestraint):
    """Restrain to a set of cross links"""

    assembly = None
    _label_map = {'wtDSS':'DSS', 'scDSS':'DSS', 'scEDC':'EDC'}

    def __init__(self, pmi_restraint):
        self.pmi_restraint = pmi_restraint
        label = self.pmi_restraint.label
        # Map commonly-used subtypes to more standard labels
        label = self._label_map.get(label, label)
        self.label = label
        super(_CrossLinkRestraint, self).__init__(
                dataset=self.pmi_restraint.dataset,
                linker_type=label)

    def _set_psi_sigma(self, model):
        old_values = []
        # Do nothing if given a different state than the restraint applies to
        if model.m != self.pmi_restraint.m:
            return
        for resolution in self.pmi_restraint.sigma_dictionary:
            statname = 'ISDCrossLinkMS_Sigma_%s_%s' % (resolution, self.label)
            if model.stats and statname in model.stats:
                sigma = float(model.stats[statname])
                p = self.pmi_restraint.sigma_dictionary[resolution][0]
                old_values.append((p, p.get_scale()))
                p.set_scale(sigma)
        for psiindex in self.pmi_restraint.psi_dictionary:
            statname = 'ISDCrossLinkMS_Psi_%s_%s' % (psiindex, self.label)
            if model.stats and statname in model.stats:
                psi = float(model.stats[statname])
                p = self.pmi_restraint.psi_dictionary[psiindex][0]
                old_values.append((p, p.get_scale()))
                p.set_scale(psi)
        # If the same particle is set multiple times, make sure we get the
        # original value last
        return list(reversed(old_values))

    def add_fits_from_model_statfile(self, model):
        # Set psi/sigma for all particles
        old_values = self._set_psi_sigma(model)
        # If no stats were collected, we can't show the fit
        if not old_values:
            return
        for xl in self.cross_links:
            # Get current psi/sigma particle value for each XL
            xl.fits[model] = ihm.restraint.CrossLinkFit(
                               psi=xl.psi, sigma1=xl.sigma1, sigma2=xl.sigma2)
        # Restore original psi/sigma
        for p, scale in old_values:
            p.set_scale(scale)

    # Have our dataset point to that in the original PMI restraint
    def __set_dataset(self, val):
        self.pmi_restraint.dataset = val
    dataset = property(lambda self: self.pmi_restraint.dataset,
                       __set_dataset)


def get_asym_mapper_for_state(simo, state, asym_map):
    asym = asym_map.get(state, None)
    if asym is None:
        asym = _AsymMapper(simo, state.prot)
        asym_map[state] = asym
    return asym

class _PMICrossLink(object):
    # Query PMI particles to get psi and sigma values at time of use, not
    # time of construction (since they may be set after creating the restraint)
    psi = property(lambda self: self.psi_p.get_scale(),
                   lambda self, val: None)
    sigma1 = property(lambda self: self.sigma1_p.get_scale(),
                      lambda self, val: None)
    sigma2 = property(lambda self: self.sigma2_p.get_scale(),
                      lambda self, val: None)


class _ResidueCrossLink(ihm.restraint.ResidueCrossLink, _PMICrossLink):
    pass


class _FeatureCrossLink(ihm.restraint.FeatureCrossLink, _PMICrossLink):
    pass


class _EM2DRestraint(ihm.restraint.EM2DRestraint):
    def __init__(self, state, pmi_restraint, image_number, resolution,
                 pixel_size, image_resolution, projection_number,
                 micrographs_number):
        self.pmi_restraint, self.image_number = pmi_restraint, image_number
        super(_EM2DRestraint, self).__init__(
                dataset=pmi_restraint.datasets[image_number],
                assembly=state.modeled_assembly,
                segment=False, number_raw_micrographs=micrographs_number,
                pixel_size_width=pixel_size, pixel_size_height=pixel_size,
                image_resolution=image_resolution,
                number_of_projections=projection_number)

    # Have our dataset point to that in the original PMI restraint
    def __get_dataset(self):
        return self.pmi_restraint.datasets[self.image_number]
    def __set_dataset(self, val):
        self.pmi_restraint.datasets[self.image_number] = val
    dataset = property(__get_dataset, __set_dataset)

    def add_fits_from_model_statfile(self, model):
        ccc = self._get_cross_correlation(model)
        transform = self._get_transformation(model)
        rot = transform.get_rotation()
        rm = [[e for e in rot.get_rotation_matrix_row(i)] for i in range(3)]
        self.fits[model] = ihm.restraint.EM2DRestraintFit(
                               cross_correlation_coefficient=ccc,
                               rot_matrix=rm,
                               tr_vector=transform.get_translation())

    def _get_transformation(self, model):
        """Get the transformation that places the model on the image"""
        stats = model.em2d_stats or model.stats
        prefix = 'ElectronMicroscopy2D_%s_Image%d' % (self.pmi_restraint.label,
                                                      self.image_number + 1)
        r = [float(stats[prefix + '_Rotation%d' % i]) for i in range(4)]
        t = [float(stats[prefix + '_Translation%d' % i])
             for i in range(3)]
        # If the model coordinates are transformed, need to back transform
        # them first
        inv = model.transform.get_inverse()
        return IMP.algebra.Transformation3D(IMP.algebra.Rotation3D(*r),
                                            IMP.algebra.Vector3D(*t)) * inv

    def _get_cross_correlation(self, model):
        """Get the cross correlation coefficient between the model projection
           and the image"""
        stats = model.em2d_stats or model.stats
        return float(stats['ElectronMicroscopy2D_%s_Image%d_CCC'
                                 % (self.pmi_restraint.label,
                                    self.image_number + 1)])


class _EM3DRestraint(ihm.restraint.EM3DRestraint):

    def __init__(self, simo, state, pmi_restraint, target_ps, densities):
        self.pmi_restraint = pmi_restraint
        super(_EM3DRestraint, self).__init__(
                dataset=pmi_restraint.dataset,
                assembly=self._get_assembly(densities, simo, state),
                fitting_method='Gaussian mixture models',
                number_of_gaussians=len(target_ps))

    # Have our dataset point to that in the original PMI restraint
    def __set_dataset(self, val):
        self.pmi_restraint.dataset = val
    dataset = property(lambda self: self.pmi_restraint.dataset,
                       __set_dataset)

    def _get_assembly(self, densities, simo, state):
        """Get the Assembly that this restraint acts on"""
        cm = _ComponentMapper(state.prot)
        components = {}
        for d in densities:
            components[cm[d]] = None # None == all residues in this component
        a = simo._get_subassembly(components,
                              name="EM subassembly",
                              description="All components that fit the EM map")
        return a

    def add_fits_from_model_statfile(self, model):
        ccc = self._get_cross_correlation(model)
        self.fits[model] = ihm.restraint.EM3DRestraintFit(
                                           cross_correlation_coefficient=ccc)

    def _get_cross_correlation(self, model):
        """Get the cross correlation coefficient between the model
           and the map"""
        if model.stats is not None:
            return float(model.stats['GaussianEMRestraint_%s_CCC'
                                     % self.pmi_restraint.label])


class _GeometricRestraint(ihm.restraint.GeometricRestraint):

    def __init__(self, simo, state, pmi_restraint, geometric_object,
                 feature, distance, sigma):
        self.pmi_restraint = pmi_restraint
        super(_GeometricRestraint, self).__init__(
                dataset=pmi_restraint.dataset,
                geometric_object=geometric_object, feature=feature,
                distance=distance, harmonic_force_constant = 1. / sigma,
                restrain_all=True)

    # Have our dataset point to that in the original PMI restraint
    def __set_dataset(self, val):
        self.pmi_restraint.dataset = val
    dataset = property(lambda self: self.pmi_restraint.dataset,
                       __set_dataset)


class _ReplicaExchangeProtocolStep(ihm.protocol.Step):
    def __init__(self, state, rex):
        if rex.monte_carlo_sample_objects is not None:
            method = 'Replica exchange monte carlo'
        else:
            method = 'Replica exchange molecular dynamics'
        super(_ReplicaExchangeProtocolStep, self).__init__(
                assembly=state.modeled_assembly,
                dataset_group=None, # filled in by add_step()
                method=method, name='Sampling',
                num_models_begin=None, # filled in by add_step()
                num_models_end=rex.vars["number_of_frames"],
                multi_scale=True, multi_state=False, ordered=False)


class _SimpleProtocolStep(ihm.protocol.Step):
    def __init__(self, state, num_models_end, method):
        super(_SimpleProtocolStep, self).__init__(
                assembly=state.modeled_assembly,
                dataset_group=None, # filled in by add_step()
                method=method, name='Sampling',
                num_models_begin=None, # filled in by add_step()
                num_models_end=num_models_end,
                multi_scale=True, multi_state=False, ordered=False)


class _Chain(object):
    """Represent a single chain in a Model"""
    def __init__(self, pmi_chain_id, asym_unit):
        self.pmi_chain_id, self.asym_unit = pmi_chain_id, asym_unit
        self.spheres = []
        self.atoms = []

    def add(self, xyz, atom_type, residue_type, residue_index,
            all_indexes, radius):
        if atom_type is None:
            self.spheres.append((xyz, residue_type, residue_index,
                                 all_indexes, radius))
        else:
            self.atoms.append((xyz, atom_type, residue_type, residue_index,
                               all_indexes, radius))
    orig_comp = property(lambda self: self.comp)

class _TransformedChain(object):
    """Represent a chain that is a transformed version of another"""
    def __init__(self, orig_chain, asym_unit, transform):
        self.orig_chain, self.asym_unit = orig_chain, asym_unit
        self.transform = transform

    def __get_spheres(self):
        for (xyz, residue_type, residue_index, all_indexes,
             radius) in self.orig_chain.spheres:
            yield (self.transform * xyz, residue_type, residue_index,
                   all_indexes, radius)
    spheres = property(__get_spheres)

    def __get_atoms(self):
        for (xyz, atom_type, residue_type, residue_index, all_indexes,
             radius) in self.orig_chain.atoms:
            yield (self.transform * xyz, atom_type, residue_type, residue_index,
                   all_indexes, radius)
    atoms = property(__get_atoms)

    entity = property(lambda self: self.orig_chain.entity)
    orig_comp = property(lambda self: self.orig_chain.comp)

class _Excluder(object):
    def __init__(self, component, simo):
        self._seqranges = simo._exclude_coords.get(component, [])

    def is_excluded(self, indexes):
        """Return True iff the given sequence range is excluded."""
        for seqrange in self._seqranges:
            if indexes[0] >= seqrange[0] and indexes[-1] <= seqrange[1]:
                return True


class _Model(ihm.model.Model):
    def __init__(self, prot, simo, protocol, assembly, representation):
        super(_Model, self).__init__(assembly=assembly, protocol=protocol,
                                     representation=representation)
        self.simo = weakref.proxy(simo)
        # Transformation from IMP coordinates into mmCIF coordinate system.
        # Normally we pass through coordinates unchanged, but in some cases
        # we may want to translate them (e.g. Nup84, where the deposited PDB
        # files are all centered; we want the mmCIF files to match)
        self.transform = IMP.algebra.get_identity_transformation_3d()
        self.em2d_stats = None
        self.stats = None
        # True iff restraints act on this model
        self._is_restrained = True
        o = self.output = IMP.pmi1.output.Output(atomistic=True)
        name = 'cif-output'
        self.m = prot.get_model()
        o.dictionary_pdbs[name] = prot
        o._init_dictchain(name, prot, multichar_chain=True)
        (particle_infos_for_pdb,
         self.geometric_center) = o.get_particle_infos_for_pdb_writing(name)
        self.geometric_center = IMP.algebra.Vector3D(*self.geometric_center)
        self._make_spheres_atoms(particle_infos_for_pdb, o, name, simo)
        self.rmsf = {}

    def all_chains(self, simo):
        """Yield all chains, including transformed ones"""
        chain_for_comp = {}
        for c in self.chains:
            yield c
            chain_for_comp[c.comp] = c
        for tc in simo._transformed_components:
            orig_chain = chain_for_comp.get(tc.original, None)
            if orig_chain:
                asym = simo.asym_units[tc.name]
                c = _TransformedChain(orig_chain, asym, tc.transform)
                c.comp = tc.name
                yield c

    def _make_spheres_atoms(self, particle_infos_for_pdb, o, name, simo):
        entity_for_chain = {}
        comp_for_chain = {}
        correct_asym = {}
        for protname, chain_id in o.dictchain[name].items():
            entity_for_chain[chain_id] = simo.entities[protname]
            comp_for_chain[chain_id] = protname
            # When doing multi-state modeling, the chain ID returned here
            # (assigned sequentially) might not be correct (states may have
            # gaps in the chain IDs). Map it to the correct asym unit.
            correct_asym[chain_id] = simo.asym_units[protname]

        # Gather by chain ID (should be sorted by chain ID already)
        self.chains = []
        chain = None
        excluder = None

        for (xyz, atom_type, residue_type, chain_id, residue_index,
             all_indexes, radius) in particle_infos_for_pdb:
            if chain is None or chain.pmi_chain_id != chain_id:
                chain = _Chain(chain_id, correct_asym[chain_id])
                chain.entity = entity_for_chain[chain_id]
                chain.comp = comp_for_chain[chain_id]
                self.chains.append(chain)
                excluder = _Excluder(chain.comp, simo)
            if not excluder.is_excluded(all_indexes if all_indexes
                                        else [residue_index]):
                chain.add(xyz, atom_type, residue_type, residue_index,
                          all_indexes, radius)

    def parse_rmsf_file(self, fname, component):
        self.rmsf[component] = rmsf = {}
        with open(fname) as fh:
            for line in fh:
                resnum, blocknum, val = line.split()
                rmsf[int(resnum)] = (int(blocknum), float(val))

    def get_rmsf(self, component, indexes):
        """Get the RMSF value for the given residue indexes."""
        if not self.rmsf:
            return None
        rmsf = self.rmsf[component]
        blocknums = dict.fromkeys(rmsf[ind][0] for ind in indexes)
        if len(blocknums) != 1:
            raise ValueError("Residue indexes %s aren't all in the same block"
                             % str(indexes))
        return rmsf[indexes[0]][1]

    def get_atoms(self):
        for chain in self.all_chains(self.simo):
            pmi_offset = chain.asym_unit.entity.pmi_offset
            for atom in chain.atoms:
                (xyz, atom_type, residue_type, residue_index,
                 all_indexes, radius) = atom
                pt = self.transform * xyz
                yield ihm.model.Atom(asym_unit=chain.asym_unit,
                        seq_id=residue_index - pmi_offset,
                        atom_id=atom_type.get_string(),
                        type_symbol=None, # todo: get element
                        x=pt[0], y=pt[1], z=pt[2])

    def get_spheres(self):
        for chain in self.all_chains(self.simo):
            pmi_offset = chain.asym_unit.entity.pmi_offset
            for sphere in chain.spheres:
                (xyz, residue_type, residue_index,
                 all_indexes, radius) = sphere
                if all_indexes is None:
                    all_indexes = (residue_index,)
                pt = self.transform * xyz
                yield ihm.model.Sphere(asym_unit=chain.asym_unit,
                        seq_id_range=(all_indexes[0] - pmi_offset,
                                      all_indexes[-1] - pmi_offset),
                        x=pt[0], y=pt[1], z=pt[2], radius=radius,
                        rmsf=self.get_rmsf(chain.orig_comp, all_indexes))


class _AllProtocols(object):
    def __init__(self, simo):
        self.simo = weakref.proxy(simo)
        # Lists of Protocols by state
        self.protocols = OrderedDict()

    def add_protocol(self, state):
        """Add a new Protocol"""
        if state not in self.protocols:
            self.protocols[state] = []
        p = ihm.protocol.Protocol()
        self.simo.system.orphan_protocols.append(p)
        self.protocols[state].append(p)

    def add_step(self, step, state):
        """Add a ProtocolStep to the last Protocol of the given State"""
        if state not in self.protocols:
            self.add_protocol(state)
        protocol = self.get_last_protocol(state)
        if len(protocol.steps) == 0:
            step.num_models_begin = 0
        else:
            step.num_models_begin = protocol.steps[-1].num_models_end
        protocol.steps.append(step)
        step.id = len(protocol.steps)
        # Assume that protocol uses all currently-defined datasets
        step.dataset_group = self.simo.all_datasets.get_all_group(state)

    def add_postproc(self, step, state):
        """Add a postprocessing step to the last protocol"""
        protocol = self.get_last_protocol(state)
        if len(protocol.analyses) == 0:
            protocol.analyses.append(ihm.analysis.Analysis())
        protocol.analyses[-1].steps.append(step)

    def get_last_protocol(self, state):
        """Return the most recently-added _Protocol"""
        return self.protocols[state][-1]


class _AllStartingModels(object):
    def __init__(self, simo):
        self.simo = simo
        # dict of starting models (entire PDB files), collected from fragments,
        # ordered by component name and state
        self.models = OrderedDict()
        self.output = IMP.pmi1.output.Output()

    def add_pdb_fragment(self, fragment):
        """Add a starting model PDB fragment."""
        comp = fragment.component
        state = fragment.state
        if comp not in self.models:
            self.models[comp] = OrderedDict()
        if state not in self.models[comp]:
            self.models[comp][state] = []
        models = self.models[comp][state]
        if len(models) == 0 \
           or models[-1].fragments[0].pdbname != fragment.pdbname:
            model = self._add_model(fragment)
            models.append(model)
        else:
            # Break circular ref between fragment and model
            models[-1].fragments.append(weakref.proxy(fragment))
            # Update residue range to cover all fragments
            pmi_offset = models[-1].asym_unit.entity.pmi_offset
            sid_begin = min(fragment.start + fragment.offset - pmi_offset,
                            models[-1].asym_unit.seq_id_range[0])
            sid_end = max(fragment.end + fragment.offset - pmi_offset,
                          models[-1].asym_unit.seq_id_range[1])
            models[-1].asym_unit = fragment.asym_unit.asym(sid_begin, sid_end)
        fragment.starting_model = models[-1]

    def _add_model(self, f):
        parser = ihm.metadata.PDBParser()
        r = parser.parse_file(f.pdbname)

        self.simo._add_dataset(r['dataset'])
        # We only want the templates that model the starting model chain
        templates = r['templates'].get(f.chain, [])
        for t in templates:
            if t.alignment_file:
                self.simo.system.locations.append(t.alignment_file)
            if t.dataset:
                self.simo._add_dataset(t.dataset)
        pmi_offset = f.asym_unit.entity.pmi_offset
        m = _StartingModel(
                    asym_unit=f.asym_unit.asym.pmi_range(f.start + f.offset,
                                                         f.end + f.offset),
                    dataset=r['dataset'], asym_id=f.chain,
                    templates=templates, offset=f.offset + pmi_offset,
                    metadata=r['metadata'],
                    software=r['software'][0] if r['software'] else None,
                    script_file=r['script'])
        m.fragments = [weakref.proxy(f)]
        return m


class _StartingModel(ihm.startmodel.StartingModel):
    def get_seq_dif(self):
        return self._seq_dif

    def get_atoms(self):
        pmi_offset = self.asym_unit.entity.pmi_offset
        mh = IMP.mmcif.data._StartingModelAtomHandler(self.templates,
                                                      self.asym_unit)
        for f in self.fragments:
            sel = IMP.atom.Selection(f.starting_hier,
                            residue_indexes=list(range(f.start - f.offset,
                                                       f.end - f.offset + 1)))
            for a in mh.get_ihm_atoms(sel.get_selected_particles(),
                                      f.offset - pmi_offset):
                yield a
        self._seq_dif = mh._seq_dif


class _ReplicaExchangeAnalysisPostProcess(ihm.analysis.ClusterStep):
    """Post processing using AnalysisReplicaExchange0 macro"""

    def __init__(self, rex, num_models_begin):
        self.rex = rex
        num_models_end = 0
        for fname in self.get_all_stat_files():
            with open(fname) as fh:
                num_models_end += len(fh.readlines())
        super(_ReplicaExchangeAnalysisPostProcess, self).__init__(
                feature='RMSD', num_models_begin=num_models_begin,
                num_models_end=num_models_end)

    def get_stat_file(self, cluster_num):
        return os.path.join(self.rex._outputdir, "cluster.%d" % cluster_num,
                            'stat.out')

    def get_all_stat_files(self):
        for i in range(self.rex._number_of_clusters):
            yield self.get_stat_file(i)


class _ReplicaExchangeAnalysisEnsemble(ihm.model.Ensemble):
    """Ensemble generated using AnalysisReplicaExchange0 macro"""

    num_models_deposited = None # Override base class property

    def __init__(self, pp, cluster_num, model_group, num_deposit):
        with open(pp.get_stat_file(cluster_num)) as fh:
            num_models = len(fh.readlines())
        super(_ReplicaExchangeAnalysisEnsemble, self).__init__(
                num_models=num_models,
                model_group=model_group, post_process=pp,
                clustering_feature=pp.feature,
                name=model_group.name)
        self.cluster_num = cluster_num
        self.num_models_deposited = num_deposit

    def get_rmsf_file(self, component):
        return os.path.join(self.post_process.rex._outputdir,
                            'cluster.%d' % self.cluster_num,
                            'rmsf.%s.dat' % component)

    def load_rmsf(self, model, component):
        fname = self.get_rmsf_file(component)
        if os.path.exists(fname):
            model.parse_rmsf_file(fname, component)

    def get_localization_density_file(self, fname):
        return os.path.join(self.post_process.rex._outputdir,
                            'cluster.%d' % self.cluster_num,
                            '%s.mrc' % fname)

    def load_localization_density(self, state, fname, select_tuple, asym_units):
        fullpath = self.get_localization_density_file(fname)
        if os.path.exists(fullpath):
            details = "Localization density for %s %s" \
                      % (fname, self.model_group.name)
            local_file = ihm.location.OutputFileLocation(fullpath,
                                                         details=details)
            for s in select_tuple:
                if isinstance(s, tuple) and len(s) == 3:
                    asym = asym_units[s[2]].pmi_range(s[0], s[1])
                else:
                    asym = asym_units[s]
                den = ihm.model.LocalizationDensity(file=local_file,
                                                    asym_unit=asym)
                self.densities.append(den)

    def load_all_models(self, simo, state):
        stat_fname = self.post_process.get_stat_file(self.cluster_num)
        model_num = 0
        with open(stat_fname) as fh:
            stats = ast.literal_eval(fh.readline())
            # Correct path
            rmf_file = os.path.join(os.path.dirname(stat_fname),
                                    "%d.rmf3" % model_num)
            for c in state.all_modeled_components:
                # todo: this only works with PMI 1
                state._pmi_object.set_coordinates_from_rmf(c, rmf_file, 0,
                                                       force_rigid_update=True)
            # todo: fill in other data from stat file, e.g. crosslink phi/psi
            yield stats
            model_num += 1
            if model_num >= self.num_models_deposited:
                return

    # todo: also support dRMS precision
    def _get_precision(self):
        precfile = os.path.join(self.post_process.rex._outputdir,
                                "precision.%d.%d.out" % (self.cluster_num,
                                                         self.cluster_num))
        if not os.path.exists(precfile):
            return ihm.unknown
        # Fail if the precision.x.x.out file doesn't match the cluster
        r = re.compile('All .*/cluster.%d/ average centroid distance ([\d\.]+)'
                       % self.cluster_num)
        with open(precfile) as fh:
            for line in fh:
                m = r.match(line)
                if m:
                    return float(m.group(1))

    precision = property(lambda self: self._get_precision(),
                         lambda self, val: None)

class _SimpleEnsemble(ihm.model.Ensemble):
    """Simple manually-created ensemble"""

    num_models_deposited = None # Override base class property

    def __init__(self, pp, model_group, num_models, drmsd,
                 num_models_deposited, ensemble_file):
        super(_SimpleEnsemble, self).__init__(
                model_group=model_group, post_process=pp, num_models=num_models,
                file=ensemble_file, precision=drmsd, name=model_group.name,
                clustering_feature='dRMSD')
        self.num_models_deposited = num_models_deposited

    def load_localization_density(self, state, component, asym, local_file):
        den = ihm.model.LocalizationDensity(file=local_file, asym_unit=asym)
        self.densities.append(den)


class _EntityMapper(dict):
    """Handle mapping from IMP components to CIF entities.
       Multiple components may map to the same entity if they share sequence."""
    def __init__(self, system):
        super(_EntityMapper, self).__init__()
        self._sequence_dict = {}
        self._entities = []
        self.system = system

    def add(self, component_name, sequence, offset):
        def entity_seq(sequence):
            # Map X to UNK
            if 'X' in sequence:
                return ['UNK' if s == 'X' else s for s in sequence]
            else:
                return sequence
        if sequence not in self._sequence_dict:
            # Use the name of the first component, stripped of any copy number,
            # as the description of the entity
            d = component_name.split("@")[0].split(".")[0]
            entity = Entity(entity_seq(sequence), description=d,
                            pmi_offset=offset)
            self.system.entities.append(entity)
            self._sequence_dict[sequence] = entity
        self[component_name] = self._sequence_dict[sequence]


class _TransformedComponent(object):
    def __init__(self, name, original, transform):
        self.name, self.original, self.transform = name, original, transform


class _SimpleRef(object):
    """Class with similar interface to weakref.ref, but keeps a strong ref"""
    def __init__(self, ref):
        self.ref = ref
    def __call__(self):
        return self.ref


class _State(ihm.model.State):
    """Representation of a single state in the system."""

    def __init__(self, pmi_object, po):
        # Point to the PMI object for this state. Use a weak reference
        # since the state object typically points to us too, so we need
        # to break the reference cycle. In PMI1 this will be a
        # Representation object; in PMI2 it is the PMI2 State object itself.
        self._pmi_object = weakref.proxy(pmi_object)
        if hasattr(pmi_object, 'state'):
            # Need a strong ref to pmi_object.state to prevent it from being
            # cleaned up when doing PMI1 multistate modeling
            self._pmi_state = _SimpleRef(pmi_object.state)
        else:
            self._pmi_state = weakref.ref(pmi_object)
        # Preserve PMI state name
        old_name = self.name
        super(_State, self).__init__(experiment_type='Fraction of bulk')
        self.name = old_name

        # The assembly of all components modeled by IMP in this state.
        # This may be smaller than the complete assembly.
        self.modeled_assembly = ihm.Assembly(
                        name="Modeled assembly",
                        description=self.get_postfixed_name(
                                         "All components modeled by IMP"))
        po.system.orphan_assemblies.append(self.modeled_assembly)

        self.all_modeled_components = []

    def __hash__(self):
        return hash(self._pmi_state())
    def __eq__(self, other):
        return self._pmi_state() == other._pmi_state()

    def add_model_group(self, group):
        self.append(group)
        return group

    def get_prefixed_name(self, name):
        """Prefix the given name with the state name, if available."""
        if self.short_name:
            return self.short_name + ' ' + name
        else:
            # Can't use capitalize() since that un-capitalizes everything
            # but the first letter
            return name[0].upper() + name[1:] if name else ''

    def get_postfixed_name(self, name):
        """Postfix the given name with the state name, if available."""
        if self.short_name:
            return "%s in state %s" % (name, self.short_name)
        else:
            return name

    short_name = property(lambda self: self._pmi_state().short_name)
    long_name = property(lambda self: self._pmi_state().long_name)
    def __get_name(self):
        return self._pmi_state().long_name
    def __set_name(self, val):
        self._pmi_state().long_name = val
    name = property(__get_name, __set_name)


class Entity(ihm.Entity):
    """A single entity in the system.
       This functions identically to the base ihm.Entity class, but it
       allows identifying residues by either the IHM numbering scheme
       (seq_id, which is always contiguous starting at 1) or by the PMI
       scheme (which is similar, but may not start at 1 if the sequence in
       the FASTA file has one or more N-terminal gaps"""
    def __init__(self, sequence, pmi_offset, *args, **keys):
        # Offset between PMI numbering and IHM; <pmi_#> = <ihm_#> + pmi_offset
        # (pmi_offset is also the number of N-terminal gaps in the FASTA file)
        self.pmi_offset = pmi_offset
        super(Entity, self).__init__(sequence, *args, **keys)

    def pmi_residue(self, res_id):
        """Return a single IHM residue indexed using PMI numbering"""
        return self.residue(res_id - self.pmi_offset)

    def pmi_range(self, res_id_begin, res_id_end):
        """Return a range of IHM residues indexed using PMI numbering"""
        off = self.pmi_offset
        return self(res_id_begin - off, res_id_end - off)


class AsymUnit(ihm.AsymUnit):
    """A single asymmetric unit in the system.
       This functions identically to the base ihm.AsymUnit class, but it
       allows identifying residues by either the IHM numbering scheme
       (seq_id, which is always contiguous starting at 1) or by the PMI
       scheme (which is similar, but may not start at 1 if the sequence in
       the FASTA file has one or more N-terminal gaps"""

    def __init__(self, entity, *args, **keys):
        super(AsymUnit, self).__init__(
                entity, auth_seq_id_map=entity.pmi_offset, *args, **keys)

    def pmi_residue(self, res_id):
        """Return a single IHM residue indexed using PMI numbering"""
        return self.residue(res_id - self.entity.pmi_offset)

    def pmi_range(self, res_id_begin, res_id_end):
        """Return a range of IHM residues indexed using PMI numbering"""
        off = self.entity.pmi_offset
        return self(res_id_begin - off, res_id_end - off)


class ProtocolOutput(IMP.pmi1.output.ProtocolOutput):
    """Class to encode a modeling protocol as mmCIF.

    IMP has basic support for writing out files in mmCIF format, for
    deposition in [PDB-Dev](https://pdb-dev.wwpdb.org/).
    After creating an instance of this class, attach it to an
    IMP.pmi1.representation.Representation object. After this, any
    generated models and metadata are automatically collected and
    output as mmCIF.
    """
    def __init__(self, fh):
        # Ultimately, collect data in an ihm.System object
        self.system = ihm.System()
        self._state_group = ihm.model.StateGroup()
        self.system.state_groups.append(self._state_group)

        self.fh = fh
        self._state_ensemble_offset = 0
        self._each_metadata = [] # list of metadata for each representation
        self._file_datasets = []
        self._main_script = os.path.abspath(sys.argv[0])

        # Point to the main modeling script
        loc = ihm.location.WorkflowFileLocation(path=self._main_script,
                               details="The main integrative modeling script")
        self.system.locations.append(loc)

        self._states = {}
        self.__asym_states = {}
        self._working_directory = os.getcwd()
        self.default_representation = self.create_representation(
                                                    "Default representation")
        self.entities = _EntityMapper(self.system)
        # Mapping from component names to ihm.AsymUnit
        self.asym_units = {}
        self._all_components = {}
        self.all_modeled_components = []
        self._transformed_components = []
        self.sequence_dict = {}

        # Common geometric objects used in PMI systems
        self._xy_plane = ihm.geometry.XYPlane()
        self._xz_plane = ihm.geometry.XZPlane()
        self._z_axis = ihm.geometry.ZAxis()
        self._center_origin = ihm.geometry.Center(0,0,0)
        self._identity_transform = ihm.geometry.Transformation.identity()

        # Coordinates to exclude
        self._exclude_coords = {}

        self.all_representations = _AllModelRepresentations(self)
        self.all_protocols = _AllProtocols(self)
        self.all_datasets = _AllDatasets(self.system)
        self.all_starting_models = _AllStartingModels(self)

        self.all_software = _AllSoftware(self.system)

    def create_representation(self, name):
        """Create a new Representation and return it. This can be
           passed to add_model(), add_bead_element() or add_pdb_element()."""
        r = ihm.representation.Representation()
        self.system.orphan_representations.append(r)
        return r

    def exclude_coordinates(self, component, seqrange):
        """Don't record coordinates for the given domain.
           Coordinates for the given domain (specified by a component name
           and a 2-element tuple giving the start and end residue numbers)
           will be excluded from the mmCIF file. This can be used to exclude
           parts of the structure that weren't well resolved in modeling.
           Any bead or residue that lies wholly within this range will be
           excluded. Multiple ranges for a given component can be excluded
           by calling this method multiple times."""
        if component not in self._exclude_coords:
            self._exclude_coords[component] = []
        self._exclude_coords[component].append(seqrange)

    def _is_excluded(self, component, start, end):
        """Return True iff this chunk of sequence should be excluded"""
        for seqrange in self._exclude_coords.get(component, ()):
            if start >= seqrange[0] and end <= seqrange[1]:
                return True

    def _add_state(self, state):
        """Create a new state and return a pointer to it."""
        self._state_ensemble_offset = len(self.system.ensembles)
        s = _State(state, self)
        self._state_group.append(s)
        self._last_state = s
        return s

    def get_file_dataset(self, fname):
        for d in self._file_datasets:
            fd = d.get(os.path.abspath(fname), None)
            if fd:
                return fd

    def _get_chain_for_component(self, name, output):
        """Get the chain ID for a component, if any."""
        # todo: handle multiple copies
        if name in self.asym_units:
            return self.asym_units[name]._id
        else:
            # A non-modeled component doesn't have a chain ID
            return None

    def _get_assembly_comps(self, assembly):
        """Get the names of the components in the given assembly"""
        # component name = asym_unit.details
        comps = {}
        for ca in assembly:
            comps[ca.details] = None
        return comps

    def create_transformed_component(self, state, name, original, transform):
        """Make a new component that's a transformed copy of another.
           All representation for the existing component is copied to the
           new one."""
        assembly_comps = self._get_assembly_comps(state.modeled_assembly)
        if name in assembly_comps:
            raise ValueError("Component %s already exists" % name)
        elif original not in assembly_comps:
            raise ValueError("Original component %s does not exist" % original)
        self.create_component(state, name, True)
        self.add_component_sequence(state, name, self.sequence_dict[original])
        self._transformed_components.append(_TransformedComponent(
                                            name, original, transform))
        self.all_representations.copy_component(state, name, original,
                                                self.asym_units[name])

    def create_component(self, state, name, modeled):
        new_comp = name not in self._all_components
        self._all_components[name] = None
        if modeled:
            state.all_modeled_components.append(name)
            if new_comp:
                self.asym_units[name] = None # assign asym once we get sequence
                self.all_modeled_components.append(name)

    def add_component_sequence(self, state, name, seq):
        def get_offset(seq):
            # Count length of gaps at start of sequence
            for i in range(len(seq)):
                if seq[i] != '-':
                    return seq[i:], i
        seq, offset = get_offset(seq)
        if name in self.sequence_dict:
            if self.sequence_dict[name] != seq:
                raise ValueError("Sequence mismatch for component %s" % name)
        else:
            self.sequence_dict[name] = seq
            self.entities.add(name, seq, offset)
        if name in self.asym_units:
            if self.asym_units[name] is None:
                # Set up a new asymmetric unit for this component
                entity = self.entities[name]
                asym = AsymUnit(entity, details=name)
                self.system.asym_units.append(asym)
                self.asym_units[name] = asym
            state.modeled_assembly.append(self.asym_units[name])

    def _add_restraint_model_fits(self):
        """Add fits to restraints for all known models"""
        for group, m in self.system._all_models():
            if m._is_restrained:
                for r in self.system.restraints:
                    if hasattr(r, 'add_fits_from_model_statfile'):
                        r.add_fits_from_model_statfile(m)

    def _add_em2d_raw_micrographs(self):
        """If the deprecated metadata.EMMicrographsDataset class was used,
           transfer its data to the EM2D restraint."""
        for r in self.system.restraints:
            if isinstance(r, _EM2DRestraint):
                for d in r.dataset.parents:
                    if isinstance(d, IMP.pmi1.metadata.EMMicrographsDataset):
                        r.number_raw_micrographs = d.number
                        if isinstance(d.location,
                                      IMP.pmi1.metadata.FileLocation):
                            d.location.content_type = \
                                   ihm.location.InputFileLocation.content_type

    def flush(self):
        self.finalize()
        ihm.dumper.write(self.fh, [self.system])

    def finalize(self):
        """Do any final processing on the class hierarchy.
           After calling this method, the `system` member (an instance
           of `ihm.System`) completely reproduces the PMI modeling, and
           can be written out to an mmCIF file with `ihm.dumper.write`,
           and/or modified using the ihm API."""
        self._add_restraint_model_fits()

        self._add_em2d_raw_micrographs()

        # Add metadata to ihm.System
        self.system.software.extend(m for m in self._metadata
                                    if isinstance(m, ihm.Software))
        self.system.citations.extend(m for m in self._metadata
                                     if isinstance(m, ihm.Citation))
        self.system.locations.extend(m for m in self._metadata
                                   if isinstance(m, ihm.location.FileLocation))

        # Point all locations to repos, if applicable
        all_repos = [m for m in self._metadata
                     if isinstance(m, ihm.location.Repository)]
        self.system.update_locations_in_repositories(all_repos)

    def add_pdb_element(self, state, name, start, end, offset, pdbname,
                        chain, hier, representation=None):
        if self._is_excluded(name, start, end):
            return
        if representation is None:
            representation = self.default_representation
        asym = self.asym_units[name]
        p = _PDBFragment(state, name, start, end, offset, pdbname, chain,
                         hier, asym)
        self.all_representations.add_fragment(state, representation, p)
        self.all_starting_models.add_pdb_fragment(p)

    def add_bead_element(self, state, name, start, end, num, hier,
                         representation=None):
        if self._is_excluded(name, start, end):
            return
        if representation is None:
            representation = self.default_representation
        asym = self.asym_units[name]
        pmi_offset = asym.entity.pmi_offset
        b = _BeadsFragment(state, name, start - pmi_offset, end - pmi_offset,
                           num, hier, asym)
        self.all_representations.add_fragment(state, representation, b)

    def get_cross_link_group(self, pmi_restraint):
        r = _CrossLinkRestraint(pmi_restraint)
        self.system.restraints.append(r)
        self._add_restraint_dataset(r) # so that all-dataset group works
        return r

    def add_experimental_cross_link(self, r1, c1, r2, c2, rsr):
        if c1 not in self._all_components or c2 not in self._all_components:
            # Crosslink refers to a component we didn't model
            # As a quick hack, just ignore it.
            # todo: need to add an entity for this anyway (so will need the
            # sequence, and add to struct_assembly)
            return None
        e1 = self.entities[c1]
        e2 = self.entities[c2]
        xl = ihm.restraint.ExperimentalCrossLink(residue1=e1.pmi_residue(r1),
                                                 residue2=e2.pmi_residue(r2))
        rsr.experimental_cross_links.append([xl])
        return xl

    def add_cross_link(self, state, ex_xl, p1, p2, length, sigma1_p, sigma2_p,
                       psi_p, rsr):
        # Map p1/p2 to asym units
        asym = get_asym_mapper_for_state(self, state, self.__asym_states)
        d = ihm.restraint.UpperBoundDistanceRestraint(length)

        if _get_by_residue(p1) and _get_by_residue(p2):
            cls = _ResidueCrossLink
        else:
            cls = _FeatureCrossLink
        xl = cls(ex_xl, asym1=asym[p1], asym2=asym[p2], distance=d,
                 restrain_all=True)
        # Needed to look up fits later
        xl.psi_p, xl.sigma1_p, xl.sigma2_p = psi_p, sigma1_p, sigma2_p
        rsr.cross_links.append(xl)

    def add_replica_exchange(self, state, rex):
        # todo: allow for metadata to say how many replicas were used in the
        # actual experiment, and how many independent runs were carried out
        # (use these as multipliers to get the correct total number of
        # output models)
        self.all_protocols.add_step(_ReplicaExchangeProtocolStep(state, rex),
                                    state)

    def _add_simple_dynamics(self, num_models_end, method):
        # Always assumed that we're dealing with the last state
        state = self._last_state
        self.all_protocols.add_step(_SimpleProtocolStep(state, num_models_end,
                                                        method), state)

    def _add_protocol(self):
        # Always assumed that we're dealing with the last state
        state = self._last_state
        self.all_protocols.add_protocol(state)

    def _add_dataset(self, dataset):
        return self.all_datasets.add(self._last_state, dataset)

    def _add_restraint_dataset(self, restraint):
        return self.all_datasets.add_restraint(self._last_state, restraint)

    def _add_simple_postprocessing(self, num_models_begin, num_models_end):
        # Always assumed that we're dealing with the last state
        state = self._last_state
        pp = ihm.analysis.ClusterStep('RMSD', num_models_begin, num_models_end)
        self.all_protocols.add_postproc(pp, state)
        return pp

    def _add_no_postprocessing(self, num_models):
        # Always assumed that we're dealing with the last state
        state = self._last_state
        pp = ihm.analysis.EmptyStep()
        pp.num_models_begin = pp.num_models_end = num_models
        self.all_protocols.add_postproc(pp, state)
        return pp

    def _add_simple_ensemble(self, pp, name, num_models, drmsd,
                             num_models_deposited, localization_densities,
                             ensemble_file):
        """Add an ensemble generated by ad hoc methods (not using PMI).
           This is currently only used by the Nup84 system."""
        # Always assumed that we're dealing with the last state
        state = self._last_state
        group = ihm.model.ModelGroup(name=state.get_postfixed_name(name))
        state.add_model_group(group)
        if ensemble_file:
            # Deprecated Location class does not state input vs output
            if isinstance(ensemble_file, IMP.pmi1.metadata.FileLocation):
                ensemble_file.content_type = \
                        ihm.location.OutputFileLocation.content_type
            self.system.locations.append(ensemble_file)
        e = _SimpleEnsemble(pp, group, num_models, drmsd, num_models_deposited,
                            ensemble_file)
        self.system.ensembles.append(e)
        for c in state.all_modeled_components:
            den = localization_densities.get(c, None)
            if den:
                # Deprecated Location class does not state input vs output
                if isinstance(den, IMP.pmi1.metadata.FileLocation):
                    den.content_type = \
                            ihm.location.OutputFileLocation.content_type
                e.load_localization_density(state, c, self.asym_units[c], den)
        return e

    def set_ensemble_file(self, i, location):
        """Point a previously-created ensemble to an 'all-models' file.
           This could be a trajectory such as DCD, an RMF, or a multimodel
           PDB file."""
        # Deprecated Location class does not state input vs output
        if isinstance(location, IMP.pmi1.metadata.FileLocation):
            location.content_type = ihm.location.OutputFileLocation.content_type
        self.system.locations.append(location)
        # Ensure that we point to an ensemble related to the current state
        ind = i + self._state_ensemble_offset
        self.system.ensembles[ind].file = location

    def add_replica_exchange_analysis(self, state, rex, density_custom_ranges):
        # todo: add prefilter as an additional postprocess step (complication:
        # we don't know how many models it filtered)
        # todo: postpone rmsf/density localization extraction until after
        # relevant methods are called (currently it is done from the
        # constructor)
        protocol = self.all_protocols.get_last_protocol(state)
        num_models = protocol.steps[-1].num_models_end
        pp = _ReplicaExchangeAnalysisPostProcess(rex, num_models)
        self.all_protocols.add_postproc(pp, state)
        for i in range(rex._number_of_clusters):
            group = ihm.model.ModelGroup(name=state.get_prefixed_name(
                                                      'cluster %d' % (i + 1)))
            state.add_model_group(group)
            # todo: make # of models to deposit configurable somewhere
            e = _ReplicaExchangeAnalysisEnsemble(pp, i, group, 1)
            self.system.ensembles.append(e)
            # Add localization density info if available
            for fname, stuple in sorted(density_custom_ranges.items()):
                e.load_localization_density(state, fname, stuple,
                                            self.asym_units)
            for stats in e.load_all_models(self, state):
                m = self.add_model(group)
                # Since we currently only deposit 1 model, it is the
                # best scoring one
                m.name = 'Best scoring model'
                m.stats = stats
                # Add RMSF info if available
                for c in state.all_modeled_components:
                    e.load_rmsf(m, c)

    def _get_subassembly(self, comps, name, description):
        """Get an Assembly consisting of the given components.
           `compdict` is a dictionary of the components to add, where keys
           are the component names and values are the sequence ranges (or
           None to use all residues in the component)."""
        asyms = []
        for comp, seqrng in comps.items():
            a = self.asym_units[comp]
            asyms.append(a if seqrng is None else a(*seqrng))

        a = ihm.Assembly(asyms, name=name, description=description)
        return a

    def _add_foxs_restraint(self, model, comp, seqrange, dataset, rg, chi,
                            details):
        """Add a basic FoXS fit. This is largely intended for use from the
           NPC application."""
        assembly = self._get_subassembly({comp:seqrange},
                              name="SAXS subassembly",
                              description="All components that fit SAXS data")
        r = ihm.restraint.SASRestraint(dataset, assembly, segment=False,
                fitting_method='FoXS', fitting_atom_type='Heavy atoms',
                multi_state=False, radius_of_gyration=rg, details=details)
        r.fits[model] = ihm.restraint.SASRestraintFit(chi_value=chi)
        self.system.restraints.append(r)
        self._add_restraint_dataset(r) # so that all-dataset group works

    def add_em2d_restraint(self, state, r, i, resolution, pixel_size,
                           image_resolution, projection_number,
                           micrographs_number):
        r = _EM2DRestraint(state, r, i, resolution, pixel_size,
                           image_resolution, projection_number,
                           micrographs_number)
        self.system.restraints.append(r)
        self._add_restraint_dataset(r) # so that all-dataset group works

    def add_em3d_restraint(self, state, target_ps, densities, pmi_restraint):
        # todo: need to set allow_duplicates on this dataset?
        r = _EM3DRestraint(self, state, pmi_restraint, target_ps, densities)
        self.system.restraints.append(r)
        self._add_restraint_dataset(r) # so that all-dataset group works

    def add_zaxial_restraint(self, state, ps, lower_bound, upper_bound,
                             sigma, pmi_restraint):
        self._add_geometric_restraint(state, ps, lower_bound, upper_bound,
                                      sigma, pmi_restraint, self._xy_plane)

    def add_yaxial_restraint(self, state, ps, lower_bound, upper_bound,
                             sigma, pmi_restraint):
        self._add_geometric_restraint(state, ps, lower_bound, upper_bound,
                                      sigma, pmi_restraint, self._xz_plane)

    def add_xyradial_restraint(self, state, ps, lower_bound, upper_bound,
                               sigma, pmi_restraint):
        self._add_geometric_restraint(state, ps, lower_bound, upper_bound,
                                      sigma, pmi_restraint, self._z_axis)

    def _add_geometric_restraint(self, state, ps, lower_bound, upper_bound,
                                 sigma, pmi_restraint, geom):
        asym = get_asym_mapper_for_state(self, state, self.__asym_states)
        r = _GeometricRestraint(self, state, pmi_restraint, geom,
                             asym.get_feature(ps),
                             ihm.restraint.LowerUpperBoundDistanceRestraint(
                                                    lower_bound, upper_bound),
                             sigma)
        self.system.restraints.append(r)
        self._add_restraint_dataset(r) # so that all-dataset group works

    def _get_membrane(self, tor_R, tor_r, tor_th):
        """Get an object representing a half-torus membrane"""
        if not hasattr(self, '_seen_membranes'):
            self._seen_membranes = {}
        # If a membrane already exists with all dimensions within 0.01 of
        # this one, reuse it
        membrane_id = tuple(int(x * 100.) for x in (tor_R, tor_r, tor_th))
        if membrane_id not in self._seen_membranes:
            m = ihm.geometry.HalfTorus(center=self._center_origin,
                    transformation=self._identity_transform,
                    major_radius=tor_R, minor_radius=tor_r, thickness=tor_th,
                    inner=True, name='Membrane')
            self._seen_membranes[membrane_id] = m
        return self._seen_membranes[membrane_id]

    def add_membrane_surface_location_restraint(
            self, state, ps, tor_R, tor_r, tor_th, sigma, pmi_restraint):
        self._add_membrane_restraint(state, ps, tor_R, tor_r, tor_th, sigma,
                  pmi_restraint, ihm.restraint.UpperBoundDistanceRestraint(0.))

    def add_membrane_exclusion_restraint(
            self, state, ps, tor_R, tor_r, tor_th, sigma, pmi_restraint):
        self._add_membrane_restraint(state, ps, tor_R, tor_r, tor_th, sigma,
                  pmi_restraint, ihm.restraint.LowerBoundDistanceRestraint(0.))

    def _add_membrane_restraint(self, state, ps, tor_R, tor_r, tor_th,
                                sigma, pmi_restraint, rsr):
        asym = get_asym_mapper_for_state(self, state, self.__asym_states)
        r = _GeometricRestraint(self, state, pmi_restraint,
                             self._get_membrane(tor_R, tor_r, tor_th),
                             asym.get_feature(ps), rsr, sigma)
        self.system.restraints.append(r)
        self._add_restraint_dataset(r) # so that all-dataset group works

    def add_model(self, group, assembly=None, representation=None):
        state = self._last_state
        if representation is None:
            representation = self.default_representation
        protocol = self.all_protocols.get_last_protocol(state)
        m = _Model(state.prot, self, protocol,
                   assembly if assembly else state.modeled_assembly,
                   representation)
        group.append(m)
        return m

    def _update_locations(self, filelocs):
        """Update FileLocation to point to a parent repository, if any"""
        all_repos = [m for m in self._metadata
                     if isinstance(m, ihm.location.Repository)]
        for fileloc in filelocs:
            ihm.location.Repository._update_in_repos(fileloc, all_repos)

    _metadata = property(lambda self:
                         itertools.chain.from_iterable(self._each_metadata))