mmcif.py

"""@namespace IMP.pmi.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.pmi.mmcif.ProtocolOutput class, and attaching it to an
   IMP.pmi.representation.Representation object, after which any
   generated models and metadata are collected and output as mmCIF.
"""

from __future__ import print_function
import copy
import RMF
import IMP.core
import IMP.algebra
import IMP.atom
import IMP.em
import IMP.isd
import IMP.pmi.tools
import IMP.pmi.alphabets
from IMP.pmi.tools import OrderedDict
import IMP.pmi.output
import IMP.mmcif.data
import re
import ast
import sys
import os
try:
    from pathlib import Path
except ImportError:  # Use bundled pathlib on Python 2 without pathlib
    from IMP._compat_pathlib import Path  # noqa: F401
import weakref
import ihm.location
import ihm.dataset
import ihm.dumper
import ihm.reader
import ihm.metadata
import ihm.startmodel
import ihm.model
import ihm.protocol
import ihm.analysis
import ihm.representation
import ihm.geometry
import ihm.cross_linkers


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.pmi.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, mmcif=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)
        self._seen_ranges = {}

    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 ranges
            if IMP.atom.Residue.get_is_setup(p):
                rind = IMP.atom.Residue(p).get_index()
                rng = 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()
                rng = asym(rinds[0], rinds[-1])
            else:
                raise ValueError("Unsupported particle type %s" % str(p))
            # Join contiguous ranges
            if len(rngs) > 0 and rngs[-1].asym == asym \
               and rngs[-1].seq_id_range[1] == rng.seq_id_range[0] - 1:
                rngs[-1].seq_id_range = (rngs[-1].seq_id_range[0],
                                         rng.seq_id_range[1])
            else:
                rngs.append(rng)
        # If an identical feature already exists, return that
        # todo: python-ihm should handle this automatically for us
        hrngs = tuple(rngs)
        if hrngs in self._seen_ranges:
            return self._seen_ranges[hrngs]
        else:
            feat = ihm.restraint.ResidueFeature(rngs)
            self._seen_ranges[hrngs] = feat
            return feat


class _AllSoftware(object):
    def __init__(self, system):
        self.system = system
        self.modeller_used = self.phyre2_used = False
        self.pmi = ihm.Software(
                name="IMP PMI module",
                version=IMP.pmi.__version__,
                classification="integrative model building",
                description="integrative model building",
                location='https://integrativemodeling.org')
        self.imp = ihm.Software(
                 name="Integrative Modeling Platform (IMP)",
                 version=IMP.__version__,
                 classification="integrative model building",
                 description="integrative model building",
                 location='https://integrativemodeling.org')
        # Only recent versions of python-ihm support adding citations for
        # software
        if hasattr(self.imp, 'citation'):
            if sys.version_info[0] > 2:
                # Don't include UTF8 characters in source; it confuses Python 2
                javi = 'Vel\u00e1zquez-Muriel J'
            else:
                javi = 'Velazquez-Muriel J'
            self.imp.citation = ihm.Citation(
                pmid='22272186',
                title='Putting the pieces together: integrative modeling '
                      'platform software for structure determination of '
                      'macromolecular assemblies',
                journal='PLoS Biol', volume=10, page_range='e1001244',
                year=2012,
                authors=['Russel D', 'Lasker K', 'Webb B', javi, 'Tjioe E',
                         'Schneidman-Duhovny D', 'Peterson B', 'Sali A'],
                doi='10.1371/journal.pbio.1001244')
            self.pmi.citation = ihm.Citation(
                pmid='31396911',
                title='Modeling Biological Complexes Using Integrative '
                      'Modeling Platform.',
                journal='Methods Mol Biol', volume=2022, page_range=(353, 377),
                year=2019,
                authors=['Saltzberg D', 'Greenberg CH', 'Viswanath S',
                         'Chemmama I', 'Webb B', 'Pellarin R', 'Echeverria I',
                         'Sali A'],
                doi='10.1007/978-1-4939-9608-7_15')
        self.system.software.extend([self.pmi, self.imp])

    def set_modeller_used(self, version, date):
        if self.modeller_used:
            return
        self.modeller_used = True
        s = ihm.Software(
            name='MODELLER', classification='comparative modeling',
            description='Comparative modeling by satisfaction '
                        'of spatial restraints, build ' + date,
            location='https://salilab.org/modeller/', version=version)
        self.system.software.append(s)
        if hasattr(s, 'citation'):
            s.citation = ihm.Citation(
                pmid='8254673',
                title='Comparative protein modelling by satisfaction of '
                      'spatial restraints.',
                journal='J Mol Biol', volume=234, page_range=(779, 815),
                year=1993, authors=['Sali A', 'Blundell TL'],
                doi='10.1006/jmbi.1993.1626')

    def set_phyre2_used(self):
        if self.phyre2_used:
            return
        self.phyre2_used = True
        s = 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/')
        if hasattr(s, 'citation'):
            s.citation = ihm.Citation(
                pmid='25950237',
                title='The Phyre2 web portal for protein modeling, '
                      'prediction and analysis.',
                journal='Nat Protoc', volume=10, page_range=(845, 858),
                authors=['Kelley LA', 'Mezulis S', 'Yates CM', 'Wass MN',
                         'Sternberg MJ'],
                year=2015, doi='10.1038/nprot.2015.053')
        self.system.software.append(s)


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     # noqa: E721
                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'}
    _descriptor_map = {'DSS': ihm.cross_linkers.dss,
                       'EDC': ihm.cross_linkers.edc}

    def __init__(self, pmi_restraint):
        self.pmi_restraint = pmi_restraint
        # Use ChemDescriptor from PMI restraint if available, otherwise guess
        # it using the label
        linker = getattr(self.pmi_restraint, 'linker', None)
        label = self.pmi_restraint.label
        self.label = label
        super(_CrossLinkRestraint, self).__init__(
                dataset=self.pmi_restraint.dataset,
                linker=linker or self._get_chem_descriptor(label))

    @classmethod
    def _get_chem_descriptor(cls, label):
        # Map commonly-used subtypes to more standard labels
        label = cls._label_map.get(label, label)
        if label not in cls._descriptor_map:
            # If label is not a standard linker type, make a new chemical
            # descriptor containing just the name. We don't know the chemistry
            # so cannot specify a SMILES or INCHI string for it at this point
            d = ihm.ChemDescriptor(label)
            cls._descriptor_map[label] = d
        return cls._descriptor_map[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.model:
            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'
        self.monte_carlo_temperature = rex.vars['monte_carlo_temperature']
        self.replica_exchange_minimum_temperature = \
            rex.vars['replica_exchange_minimum_temperature']
        self.replica_exchange_maximum_temperature = \
            rex.vars['replica_exchange_maximum_temperature']
        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, ensemble=True)


class _ReplicaExchangeProtocolDumper(ihm.dumper.Dumper):
    """Write IMP-specific information about replica exchange to mmCIF.
       Note that IDs will have already been assigned by python-ihm's
       standard modeling protocol dumper."""
    def dump(self, system, writer):
        with writer.loop("_imp_replica_exchange_protocol",
                         ["protocol_id", "step_id", "monte_carlo_temperature",
                          "replica_exchange_minimum_temperature",
                          "replica_exchange_maximum_temperature"]) as lp:
            for p in system._all_protocols():
                for s in p.steps:
                    if isinstance(s, _ReplicaExchangeProtocolStep):
                        self._dump_step(p, s, lp)

    def _dump_step(self, p, s, lp):
        mintemp = s.replica_exchange_minimum_temperature
        maxtemp = s.replica_exchange_maximum_temperature
        lp.write(protocol_id=p._id, step_id=s._id,
                 monte_carlo_temperature=s.monte_carlo_temperature,
                 replica_exchange_minimum_temperature=mintemp,
                 replica_exchange_maximum_temperature=maxtemp)


class _ReplicaExchangeProtocolHandler(ihm.reader.Handler):
    category = '_imp_replica_exchange_protocol'

    """Read IMP-specific information about replica exchange from mmCIF."""
    def __call__(self, protocol_id, step_id, monte_carlo_temperature,
                 replica_exchange_minimum_temperature,
                 replica_exchange_maximum_temperature):
        p = self.sysr.protocols.get_by_id(protocol_id)
        # todo: match step_id properly (don't assume contiguous)
        s = p.steps[int(step_id)-1]
        # Upgrade from plain ihm Step to IMP subclass
        s.__class__ = _ReplicaExchangeProtocolStep
        s.monte_carlo_temperature = \
            self.get_float(monte_carlo_temperature)
        s.replica_exchange_minimum_temperature = \
            self.get_float(replica_exchange_minimum_temperature)
        s.replica_exchange_maximum_temperature = \
            self.get_float(replica_exchange_maximum_temperature)


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,
                ensemble=True)


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.pmi.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():
            if protname in simo.entities:
                entity_for_chain[chain_id] = simo.entities[protname]
            else:
                # Remove copy number
                pn = protname.split('.')[0]
                entity_for_chain[chain_id] = simo.entities[pn]
            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(str(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.pmi.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
            # We don't use fragment.offset here because fragment numbering
            # is already in seq_id, not PDB numbering
            pmi_offset = models[-1].asym_unit.entity.pmi_offset
            sid_begin = min(fragment.start - pmi_offset,
                            models[-1].asym_unit.seq_id_range[0])
            sid_end = max(fragment.end - 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)
        source = r['entity_source'].get(f.chain)
        if source:
            f.asym_unit.entity.source = source
        pmi_offset = f.asym_unit.entity.pmi_offset
        m = _StartingModel(
                    asym_unit=f.asym_unit.asym.pmi_range(f.start, f.end),
                    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(str(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 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(str(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 (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 fname.exists():
            model.parse_rmsf_file(fname, component)

    def get_localization_density_file(self, fname):
        return (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 fullpath.exists():
            details = "Localization density for %s %s" \
                      % (fname, self.model_group.name)
            local_file = ihm.location.OutputFileLocation(str(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(str(stat_fname)) as fh:
            stats = ast.literal_eval(fh.readline())
            # Correct path
            rmf_file = stat_fname.parent / ("%d.rmf3" % model_num)
            # todo: test with real PMI2 systems
            if rmf_file.exists():
                rh = RMF.open_rmf_file_read_only(str(rmf_file))
                system = state._pmi_object.system
                IMP.rmf.link_hierarchies(rh, [system.hier])
                IMP.rmf.load_frame(fh, RMF.FrameID(0))
            # todo: fill in other data from stat file, e.g. cross-link 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 = (self.post_process.rex._outputdir /
                    ("precision.%d.%d.out" % (self.cluster_num,
                                              self.cluster_num)))
        if not precfile.exists():
            return ihm.unknown
        # Fail if the precision.x.x.out file doesn't match the cluster
        r = re.compile(
            r'All .*/cluster.%d/ average centroid distance ([\d\.]+)'
            % self.cluster_num)
        with open(str(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 _CustomDNAAlphabet(object):
    """Custom DNA alphabet that maps A,C,G,T (rather than DA,DC,DG,DT
       as in python-ihm)"""
    _comps = dict([cc.code_canonical, cc]
                  for cc in ihm.DNAAlphabet._comps.values())


class _EntityMapper(dict):
    """Handle mapping from IMP components (without copy number) 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 _get_alphabet(self, alphabet):
        """Map a PMI alphabet to an IHM alphabet"""
        # todo: we should probably ignore alphabets entirely and use
        # residue types directly (e.g. we shouldn't compare one-letter
        # code sequence to determine if an Entity is unique)
        alphabet_map = {None: ihm.LPeptideAlphabet,
                        IMP.pmi.alphabets.amino_acid: ihm.LPeptideAlphabet,
                        IMP.pmi.alphabets.rna: ihm.RNAAlphabet,
                        IMP.pmi.alphabets.dna: _CustomDNAAlphabet}
        if alphabet in alphabet_map:
            return alphabet_map[alphabet]
        else:
            raise TypeError("Don't know how to handle %s" % alphabet)

    def add(self, component_name, sequence, offset, alphabet):
        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,
                            alphabet=self._get_alphabet(alphabet))
            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.pmi.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.pmi.topology.System object. After this, any
    generated models and metadata are automatically collected in the
    `system` attribute, which is an
    [ihm.System](https://python-ihm.readthedocs.io/en/latest/main.html#ihm.System) object.
    Once the protocol is complete, call finalize() to make sure `system`
    contains everything, then use the
    [python-ihm API](https://python-ihm.readthedocs.io/en/latest/dumper.html#ihm.dumper.write)
    to write out files in mmCIF or BinaryCIF format.

    See also get_handlers(), get_dumpers().
    """  # noqa: E501
    def __init__(self):
        # 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._state_ensemble_offset = 0
        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(name=name)
        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_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, asym_name=None):
        if asym_name is None:
            asym_name = name
        new_comp = name not in self._all_components
        self._all_components[name] = None
        if modeled:
            state.all_modeled_components.append(name)
            if asym_name not in self.asym_units:
                # assign asym once we get sequence
                self.asym_units[asym_name] = None
            if new_comp:
                self.all_modeled_components.append(name)

    def add_component_sequence(self, state, name, seq, asym_name=None,
                               alphabet=None):
        if asym_name is None:
            asym_name = name

        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, alphabet)
        if asym_name in self.asym_units:
            if self.asym_units[asym_name] is None:
                # Set up a new asymmetric unit for this component
                entity = self.entities[name]
                asym = AsymUnit(entity, details=asym_name)
                self.system.asym_units.append(asym)
                self.asym_units[asym_name] = asym
            state.modeled_assembly.append(self.asym_units[asym_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 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()

    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=False)
        # 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)
        step = _ReplicaExchangeProtocolStep(state, rex)
        step.software = self.all_software.pmi
        self.all_protocols.add_step(step, 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:
            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:
                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."""
        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)
        pp.software = self.all_software.pmi
        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 get_dumpers():
    """Get custom python-ihm dumpers for writing PMI to from mmCIF.
       This returns a list of custom dumpers that can be passed as all or
       part of the `dumpers` argument to ihm.dumper.write(). They add
       PMI-specific information to mmCIF or BinaryCIF files written out
       by python-ihm."""
    return [_ReplicaExchangeProtocolDumper]


def get_handlers():
    """Get custom python-ihm handlers for reading PMI data from mmCIF.
       This returns a list of custom handlers that can be passed as all or
       part of the `handlers` argument to ihm.reader.read(). They read
       PMI-specific information from mmCIF or BinaryCIF files read in
       by python-ihm."""
    return [_ReplicaExchangeProtocolHandler]


class GMMParser(ihm.metadata.Parser):
    """Extract metadata from an EM density GMM file."""

    def parse_file(self, filename):
        """Extract metadata from `filename`.
           @return a dict with key `dataset` pointing to the GMM file and
           `number_of_gaussians` to the number of GMMs (or None)"""
        loc = ihm.location.InputFileLocation(
            filename,
            details="Electron microscopy density map, "
                    "represented as a Gaussian Mixture Model (GMM)")
        # A 3DEM restraint's dataset ID uniquely defines the mmCIF
        # restraint, so we need to allow duplicates
        loc._allow_duplicates = True
        d = ihm.dataset.EMDensityDataset(loc)
        ret = {'dataset': d, 'number_of_gaussians': None}

        with open(filename) as fh:
            for line in fh:
                if line.startswith('# data_fn: '):
                    p = ihm.metadata.MRCParser()
                    fn = line[11:].rstrip('\r\n')
                    dataset = p.parse_file(os.path.join(
                                     os.path.dirname(filename), fn))['dataset']
                    ret['dataset'].parents.append(dataset)
                elif line.startswith('# ncenters: '):
                    ret['number_of_gaussians'] = int(line[12:])
        return ret