from __future__ import division
from __future__ import with_statement

import scitbx.math
import scitbx.matrix
import phaser

import operator
import os.path
import itertools
import math

class Blackboard(object):
    """
    Blackboard for storing data
    """
    
    def __init__(self):
            
        self.data = set()
        
        
    def write(self, obj):
        
        self.data.add( obj )
        
    
    def read(self, type):
        
        return [ obj for obj in self.data if isinstance( obj, type ) ]
    
    
    def readiter(self, type):
        
        return ( obj for obj in self.data if isinstance( obj, type ) )
    
    
    def delete(self, obj):
        
        self.data.remove( obj )


# Rotation descriptions
class Rotation(object):
    """
    Rotation in Euler angles
    """
    
    def __init__(self, elements):
        
        self.setting = elements
        
        
    def euler(self):
        
        return scitbx.math.euler_angles_zyz_angles( self.setting )
    
    
    def matrix(self):
        
        return self.setting
    
    
    @classmethod
    def Euler(cls, angles):
        
        return cls( elements = scitbx.math.euler_angles_zyz_matrix( *angles ) )
    
    
    @classmethod
    def Matrix(cls, elements):
        
        return cls( elements = elements )


class Component(object):
    """
    A component
    """
    
    def __init__(self, mw):
        
        self.mw = mw
        
        
    def apply(self, input):
        
        input.addCOMP_PROT_MW_NUM( self.mw, 1 )
        
        
    def superposable(self):
        
        return True
    
    
    def empty(self):
        
        return self.mw == 0
    
    
    def subtractable(self, other):
        
        return self == other
    
    
    def __sub__(self, other):
        
        if not self.subtractable( other = other ):
            raise RuntimeError, "Incompatible components"
        
        return Component( mw = 0 )
    
    
    def __str__(self):
        
        return "Component(mw = %.1f)" % self.mw
        
        
class SequenceData(object):
    """
    A sequence
    """
    
    def __init__(self, sequence, negligible_fraction = 0.1):
        
        self.sequence = sequence
        self.negligible_count = max(
            int( negligible_fraction * len( self.sequence ) ),
            1
            ) 
        
        
    def mw(self, selection):
        
        from phaser.rsam import sequence_weight
        from phaser import mmt
        
        return sequence_weight(
            sequence = self.residue_sequence( selection = selection ),
            mmt = mmt.PROTEIN
            )
        
        
    def residue_sequence(self, selection):
        
        return [
            ch for ( ch, is_selected )
            in zip( self.sequence.sequence, selection.selection )
            if is_selected
            ]
        
        
    def superposable(self, selection):
        
        return self.negligible_count < len( selection )
    
    
    def is_negligible(self, selection):
        
        return len( selection ) <= self.negligible_count
        
        
    def __len__(self):
        
        return len( self.sequence )
        
        
class SequenceComponent(object):
    """
    A component specified by the sequence
    """
    
    def __init__(self, sequence, selection):
        
        self.sequence = sequence
        self.selection = selection
        self.mw = self.sequence.mw( selection = self.selection )
    
    def apply(self, input):
        
        input.addCOMP_PROT_MW_NUM( self.mw, 1 )
        
        
    def superposable(self):
        
        return self.sequence.superposable( selection = self.selection )
    
    
    def empty(self):
        
        return self.sequence.is_negligible( selection = self.selection )
    
    
    def subtractable(self, other):
        
        if not isinstance( other, self.__class__ ) or self.sequence != other.sequence:
            return False
        
        overlap = self.selection.overlap_with( other = other.selection )
        
        if self.sequence.is_negligible( selection = overlap ):
            return False
        
        diff = other.selection.difference_from( other = self.selection )
        return self.sequence.is_negligible( selection = diff )
    
    
    def residue_sequence(self):
        
        return self.sequence.residue_sequence( selection = self.selection )
    
    
    def __sub__(self, other):
        
        if not self.subtractable( other = other ):
            raise RuntimeError, "Incompatible components"
        
        return self.__class__(
            sequence = self.sequence,
            selection = self.selection.difference_from( other = other.selection )
            )
        
        
    def __str__(self):
        
        return "SequenceComponent(mw = %.1f)" % self.mw
        

class Composition(object):
    """
    Composition
    """
    
    def __init__(self, components):
        
        self.components = sorted( components, key = lambda s: s.mw )
        
        
    def weight(self):
        
        return sum( [ c.mw for c in self.components ] )
        
        
    def apply(self, input):
        
        for c in self.components:
            c.apply( input = input )
            
            
    def empty(self):
        
        return all( c.empty() for c in self.components )
    
    
    def contains(self, other):
        
        try:
            res = self - other
            
        except RuntimeError:
            return False
        
        return True
    
    
    def superposable(self):
        
        return any( c.superposable() for c in self.components )
    
    
    def overlap(self, other):
        
        common = []
        copy = self.__class__( components = self.components )
        
        for c in sorted( other.components, key = lambda s: s.mw, reverse = True ):
            try:
                copy -= self.__class__( components = [ c ] )
                
            except RuntimeError:
                continue
            
            common.append( c )
            
        return self.__class__( components = common )
    
    
    def subtract_component(self, component):
        
        for ( index, comp ) in enumerate( self.components ):
            try:
                diff = comp - component
                
            except RuntimeError:
                continue
            
            self.components[ index ] = diff
            break
        
        else:
            raise RuntimeError, "Component not in composition"
    
    
    def copy(self):
        
        return self.__class__( components = self.components )
            
            
    def __eq__(self, other):
        
        try:
            diff = self - other
            
        except RuntimeError:
            return False
        
        return diff.empty()
    
    
    def __ne__(self, other):
        
        return not self == other
    
    
    def __add__(self, other):
        
        return self.__class__( components = self.components + other.components )
    
    
    def __sub__(self, other):
        
        remainder = self.copy()
        
        for c in sorted( other.components, key = lambda s: s.mw, reverse = True ):
            remainder.subtract_component( component = c )
            
        return remainder
    
    
    def __mul__(self, count):
        
        return self.__class__( components = self.components * count )
    
    
class ChainModel(object):
    """
    Universal model format
    """
    
    def __init__(self, model, rmsd):
        
        self.model = model
        self.rmsd = rmsd
        
        
    def transform(self, rt):
        
        atoms = self.model.atoms()
        coords = atoms.extract_xyz()
        atoms.set_xyz( rt * coords )
        
        
    def copy(self):
        
        return self.__class__(
            model = self.model.detached_copy(),
            rmsd = self.rmsd,
            )
        
        
class ChainModelGroup(object):
    """
    Composition and set of ChainModels
    """
    
    def __init__(self, models, composition):
        
        self.models = models
        self.composition = composition
        
        
    def transform(self, rt):
        
        for cm in self.models:
            cm.transform( rt = rt )
            
            
    def copy(self):
        
        return self.__class__(
            models = [ cm.copy() for cm in self.models ],
            composition = self.composition,
            )
            
            
class Ensemble(object):
    """
    Ensemble prototype
    """
    
    def __init__(self):
        
        self._matching_ensemble_info = None
        
    
    def matching_ensemble_info(self):
        
        if self._matching_ensemble_info is None:
            from phaser import matching
            self._matching_ensemble_info = matching.EnsembleInfo.from_hierarchy(
                root = self.root()
                )
            
        return self._matching_ensemble_info
    
    
    def superpose_template(self):
        
        return self
    
    
    def template_equivalent_ensemble_info(self):
        
        return self.matching_ensemble_info()
    
    
    def substituteable_with(self, other):
        
        if not self.composition.contains( other = other.composition ):
            return False
        
        overlap = self.composition.overlap( other = other.composition )
        
        return overlap.superposable()
    
    
    def identifier(self):
        
        return "ensemble%s" % id( self )
    
    
class MultifileEnsemble(Ensemble):
    """
    Ensemble defined with multiple PDB files
    """
    
    def __init__(self, pdbs, rmsds, composition):
        
        assert pdbs
        assert len( pdbs ) == len( rmsds )
        assert all( os.path.isfile( p ) for p in pdbs )
        self.pdbs = pdbs
        self.rmsds = rmsds
        self.composition = composition
        super( MultifileEnsemble, self ).__init__()
        
    
    def apply_ensemble_definition(self, input):
        
        name = self.identifier()
        
        for ( p, rms ) in zip( self.pdbs, self.rmsds ):
            data = open( p ).read()
            input.addENSE_STR_RMS(
                modlid = name,
                content = data,
                name = name,
                rms = rms,
                ptgrp_tolerances = phaser.data_ptgrp(),
                )
    
    
    def score(self):
        
        return -self.rms()
    
    
    def root(self):
        
        import iotbx.pdb
        return iotbx.pdb.input( self.pdbs[0] ).construct_hierarchy()
    
    
    def roots(self):
        
        import iotbx.pdb
        
        for pdb in self.pdbs:
            yield iotbx.pdb.input( pdb ).construct_hierarchy()
    
    
    def chain_model_groups(self):
        
        import iotbx.pdb
        models = []
        
        for ( pdb, rmsd ) in zip( self.pdbs, self.rmsds ):
            for m in iotbx.pdb.input( pdb ).construct_hierarchy().models():
                copy = m.detached_copy()
                copy.id = ""
                models.append( ChainModel( model = copy, rmsd = rmsd ) )
        
        yield ChainModelGroup( models = models, composition = self.composition )
    
    
    def rms(self):
        
        assert self.rmsds
        return sum( self.rmsds ) / len( self.rmsds )
    
    
    def suffix(self):
        
        return "_".join( os.path.basename( f ) for f in self.pdbs )
    
    
    def root_suffix(self):
        
        return os.path.splitext( os.path.basename( self.pdbs[0] ) )[0]
    
    
    def assembly_instructions(self):
        
        return [
            ( self, scitbx.matrix.rt( ( ( 1, 0, 0, 0, 1, 0, 0, 0, 1 ), ( 0, 0, 0 ) ) ) )
            ]
    
    
    def phaser_keyword_format(self):
        
        return "ENSEMBLE %s %s" % (
            self.identifier(),
            " ".join( "PDB %s RMS %.3f" % ( pdb, rms )
                for ( pdb, rms ) in zip( self.pdbs, self.rmsds ) ),
            )
    
    
    def __str__(self):
        
        representative = self.root_suffix()
        
        if len( self.pdbs ) == 1:
            return representative
        
        elif len( self.pdbs ) == 2:
            return "%s + %s" % ( representative, os.path.basename( self.pdbs[1] ) )
        
        else:
            return "%s + %d models" % ( representative, len( self.pdbs ) - 1 )
        
        
class MultimodelEnsemble(Ensemble):
    """
    Ensemble defined with multiple iotbx.pdb.hierarchy.root-objects
    """
    
    def __init__(self, roots, rmsds, composition):
        
        assert roots
        assert len( roots ) == len( rmsds )
        self._roots = roots
        self.rmsds = rmsds
        self.composition = composition
        super( MultimodelEnsemble, self ).__init__()
        
    
    def apply_ensemble_definition(self, input):
        
        name = self.identifier()
        
        for ( r, rms ) in zip( self._roots, self.rmsds ):
            input.addENSE_STR_RMS(
                modlid = name,
                content = r.as_pdb_string(),
                name = name,
                rms = rms,
                ptgrp_tolerances = phaser.data_ptgrp(),
                )
    
    
    def score(self):
        
        return -self.rms()
    
    
    def root(self):
        
        return self._roots[0]
    
    
    def roots(self):
        
        return iter( self._roots )
    
    
    def chain_model_groups(self):
        
        models = []
        
        for ( root, rmsd ) in zip( self._roots, self.rmsds ):
            for m in root.models():
                copy = m.detached_copy()
                copy.id = ""
                models.append( ChainModel( model = copy, rmsd = rmsd ) )
        
        yield ChainModelGroup( models = models, composition = self.composition )
    
    
    def rms(self):
        
        assert self.rmsds
        return sum( self.rmsds ) / len( self.rmsds )
    
    
    def suffix(self):
        
        return "_".join( "root%s" % id( r ) for r in self._roots )
    
    
    def root_suffix(self):
        
        return "root%s" % id( self._roots[0] )
    
    
    def assembly_instructions(self):
        
        return [
            ( self, scitbx.matrix.rt( ( ( 1, 0, 0, 0, 1, 0, 0, 0, 1 ), ( 0, 0, 0 ) ) ) )
            ]
    
    
    def phaser_keyword_format(self):
        
        return "ENSEMBLE %s %s" % (
            self.identifier(),
            " ".join( "MEMADDR %s RMS %.3f" % ( id( root ), rms )
                for ( root, rms ) in zip( self._roots, self.rmsds ) ),
            )
    
    
    def __str__(self):
        
        representative = self.root_suffix()
        
        if len( self._roots ) == 1:
            return representative
        
        elif len( self._roots ) == 2:
            return "%s + %s" % ( representative, "root%s" % id( self._roots[1] ) )
        
        else:
            return "%s + %d models" % ( representative, len( self._roots ) - 1 )
        
        
class SculptorEnsemble(MultifileEnsemble):
    """
    A MultifileEnsemble created by sculptor
    """
    
    def __init__(self, pdbs, rmsds, composition, template):
        
        self.template = template
        super( SculptorEnsemble, self ).__init__(
            pdbs = pdbs,
            rmsds = rmsds,
            composition = composition
            )
    
    
    def superpose_template(self):
        
        return self.template
    
    
    def template_equivalent_ensemble_info(self):
        
        return self.template.ensembles[0].matching_ensemble_info()
        
        
class AssemblyEnsemble(Ensemble):
    """
    Ensemble defined with an assembly
    """
    
    def __init__(self, assembly):
        
        self.assembly = assembly
        self.composition = assembly.composition
        super( AssemblyEnsemble, self ).__init__()
        
    
    def apply_ensemble_definition(self, input):
        
        input.addENSE_STR_RMS(
            modlid = self.identifier(),
            content = self.root().as_pdb_string(),
            name = self.root_suffix(),
            rms = self.rms(),
            ptgrp_tolerances = phaser.data_ptgrp(),
            )
    
    
    def score(self):
        
        return -self.rms()
    
    
    def root(self):
        
        return self.assembly.root()
    
    
    def roots(self):
        
        yield self.root()
    
    
    def chain_model_groups(self):
        
        return self.assembly.chain_models()
    
    
    def rms(self):
        
        return self.assembly.rms
    
    
    def suffix(self):
        
        return self.root_suffix()
    
    
    def root_suffix(self):
        
        return "multimer%s" % id( self )
    
    
    def assembly_instructions(self):
        
        return self.assembly.description
    
    
    def phaser_keyword_format(self):
        
        return "ENSEMBLE %s %s %s" % (
            self.identifier(),
            "ASSEMBLY %s" % self.assembly,
            "RMS %.3f" % self.rms(),
            )
    
    
    def __str__(self):
        
        return "Complex: %s" % self.assembly.info()

    
class RotationPeak(object):
    """
    A rotation peak
    """
    
    def __init__(self, ensemble, rotation, rf, rfz):
        
        self.ensemble = ensemble
        self.rotation = rotation
        self.rf = rf
        self.rfz = rfz
        
    
    def mr_rlist(self):
        
        return phaser.mr_rlist(
            self.ensemble.identifier(),
            self.rotation.euler(),
            self.rf,
            self.rfz
            )
        
        
    def score(self):
        
        return self.rf
    
    
    def __str__(self):
        
        return "ENSE %s EULER (%.3f, %.3f, %.3f) RF=%.3f RFZ=%s" % (
            ( self.ensemble, ) + self.rotation.euler()
            + ( self.rf, "%.2f" % self.rfz if self.rfz is not None else "0.0" )
            )
    
    
class TranslationPeak(object):
    """
    A translation peak 
    """
    
    def __init__(self, ensemble, rotation, translation, tf, tfz, bfactor):
        
        self.ensemble = ensemble
        self.rotation = rotation
        self.translation = translation
        self.tf = tf
        self.tfz = tfz
        self.bfactor = bfactor
    
    
    def mr_ndim(self, fixr = False, fixt = False, fixb = False):
        
        return phaser.mr_ndim(
            MODLID_ = self.ensemble.identifier(),
            R_ = self.rotation.matrix(),
            FRAC_ = True,
            inFRAC_ = True,
            TRA_ = self.translation,
            BFAC_= self.bfactor,
            FIXR_ = fixr,
            FIXT_ = fixt,
            FIXB_ = fixb,
            )
        
        
    def score(self):
        
        return self.tf
    
    
    def as_scitbx_rt_operator(self, cell):
        
        return scitbx.matrix.rt(
            ( self.rotation.matrix(), cell.orthogonalize( self.translation ) )
            )
    
    
    def independent_component_list(self, cell):
        
        independent = []
        operation = self.as_scitbx_rt_operator( cell = cell )
        
        for ( ense, transf ) in self.ensemble.assembly_instructions():
            final = operation * transf
            independent.append(
                self.__class__(
                    ensemble = ense,
                    rotation = Rotation.Matrix( elements = final.r.elems ),
                    translation = cell.fractionalize( final.t.elems ),
                    tf = self.tf,
                    tfz = self.tfz,
                    bfactor = self.bfactor
                    )
                )
            
        return independent
    
    
    def as_peak(self):
        
        return Peak(
            ensemble = self.ensemble,
            rotation = self.rotation,
            translation = self.translation,
            bfactor = self.bfactor,
            )
    
    
    def __str__(self):
        
        return "ENSE %s EULER (%.3f, %.3f, %.3f) FRAC (%.3f, %.3f, %.3f) TF=%.3f TFZ=%.2f" % (
            ( self.ensemble, ) + self.rotation.euler() + self.translation + ( self.tf, self.tfz )
            )    


class Peak(object):
    """
    A peak with a rotation and a translation
    """
    
    def __init__(self, ensemble, rotation, translation, bfactor):
        
        self.ensemble = ensemble
        self.rotation = rotation
        self.translation = translation
        self.bfactor = bfactor
        
        
    def mr_ndim(self, fixr = False, fixt = False, fixb = False):
        
        return phaser.mr_ndim(
            MODLID_ = self.ensemble.identifier(),
            R_ = self.rotation.matrix(),
            FRAC_ = True,
            inFRAC_ = True,
            TRA_ = self.translation,
            BFAC_= self.bfactor,
            FIXR_ = fixr,
            FIXT_ = fixt,
            FIXB_ = fixb,
            )
        
        
    def root(self, cell):
        
        root = self.ensemble.root().deep_copy()
        atoms = root.atoms()
        coords = atoms.extract_xyz()
        atoms.set_xyz( self.as_scitbx_rt_operator( cell = cell ) * coords )
        return root
    
    
    def roots(self, cell):
        
        for r in self.ensemble.roots():
            root = r.deep_copy()
            atoms = root.atoms()
            coords = atoms.extract_xyz()
            atoms.set_xyz( self.as_scitbx_rt_operator( cell = cell ) * coords )
            yield root
    
    
    def chain_model_groups(self, cell):
        
        rt = self.as_scitbx_rt_operator( cell = cell )
        
        for cmgroup in self.ensemble.chain_model_groups():
            cmgroup.transform( rt = rt )
            yield cmgroup
    
    
    def as_ncs_operator(self, cell):
        
        from phaser import ncs
        
        return ncs.Operator(
            rotation = self.rotation.matrix(),
            translation = cell.orthogonalize( self.translation )
            )
        
        
    def as_scitbx_rt_operator(self, cell):
        
        return scitbx.matrix.rt(
            ( self.rotation.matrix(), cell.orthogonalize( self.translation ) )
            )
        
        
    def transformed(self, operation, cell):
        
        final = operation * self.as_scitbx_rt_operator( cell = cell )
        return self.__class__(
            ensemble = self.ensemble,
            rotation = Rotation.Matrix( elements = final.r.elems ),
            translation = cell.fractionalize( final.t ),
            bfactor = self.bfactor
            )
        
        
    def independent_component_list(self, cell):
        
        independent = []
        operation = self.as_scitbx_rt_operator( cell = cell )
        
        for ( ense, transf ) in self.ensemble.assembly_instructions():
            final = operation * transf
            independent.append(
                self.__class__(
                    ensemble = ense,
                    rotation = Rotation.Matrix( elements = final.r.elems ),
                    translation = cell.fractionalize( final.t.elems ),
                    bfactor = self.bfactor
                    )
                )
            
        return independent
    
    
    def ensembles_sorted_order(self):
        
        return sorted( [ p.ensemble for p in self.peaks ] )
        
        
    def matching_mrpeak(
        self,
        crystal,
        dmin,
        cached = False,
        multiplier = 1,
        template_equivalence = False
        ):
        
        from phaser import matching
        
        if cached:
            detail_orientation = matching.cached_rotation_differences
            detail_position = matching.cached_position_differences
            
        else:
            detail_orientation = matching.simple_rotation_differences
            detail_position = matching.simple_position_differences
            
        if template_equivalence:
            ensemble = self.ensemble.template_equivalent_ensemble_info()
            
        else:
            ensemble = self.ensemble.matching_ensemble_info()
            
        return matching.Labelled.phaser_mr_solution(
            rotation = self.rotation.matrix(),
            translation = self.translation,
            crystal = crystal,
            ensemble = ensemble,
            dmin = dmin,
            multiplier = multiplier,
            detail_position = detail_position,
            detail_orientation = detail_orientation,
            )
    
    
    def __str__(self):
        
        return "ENSE %s EULER (%.3f, %.3f, %.3f) FRAC (%.5f, %.5f, %.5f) BFAC %s" % (
            ( self.ensemble, ) + self.rotation.euler() + self.translation
            + ( "%.3f" % self.bfactor if self.bfactor is not None else "None", )
            )
            
            
class Structure(object):
    """
    A collection of peaks
    """
    
    def __init__(
        self,
        peaks,
        significant = False,
        rfz = 0.0,
        tfz = 0.0,
        tfze = 0.0,
        predecessor = None,
        probability = None,
        ):
        
        self.peaks = peaks
        self.predecessor = predecessor
        self.significant = significant
        self.rfz = rfz
        self.tfz = tfz
        self.tfze = tfze
        self.probability = probability
        
        
    def ensembles(self):
        
        return [ p.ensemble for p in self.peaks ]
    
    
    def composition(self):
        
        return reduce(
            operator.add,
            [ e.composition for e in self.ensembles() ],
            Composition( components = [] )
            )
        
        
    def mr_set(
        self,
        rotations = [],
        translations = [],
        sflags = ( False, False, False ),
        oflags = ( False, False, False ),
        ):
        
        flags = [ sflags ] * len( self.peaks ) + [ oflags ] * len( translations )
         
        return phaser.mr_set(
            self.identifier(),
            [ p.mr_ndim( fixr = r, fixt = t, fixb = b ) for ( p, ( r, t, b ) )
                in zip( self.peaks + translations, flags ) ],
            [ r.mr_rlist() for r in rotations ],
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            [],
            phaser.MapCoefs(),
            phaser.map_str_float1D(),
            )
        
        
    def identifier(self):
        
        return "structure%s" % id( self )
    
    
    def old_root(self, cell):
        
        import iotbx.pdb
        
        model = iotbx.pdb.hierarchy.model()
        chain_id = ord( "A" )
        
        for p in self.peaks:
            r = p.root( cell = cell )
            assert r.models()
            
            for c in r.models()[0].chains():
                copy = c.detached_copy()
                copy.id = chr( chain_id )
                chain_id += 1
                model.append_chain( copy )
            
        root = iotbx.pdb.hierarchy.root()
        root.append_model( model )
        
        return root
    
    
    def chain_model_groups(self, cell):
        
        for peak in self.peaks:
            for cmg in peak.chain_model_groups( cell = cell ):
                yield cmg
                
                
    def atomic_constituents(self, cell):
        
        for peak in self.peaks:
            for const in peak.independent_component_list( cell = cell ):
                yield const
    
    
    def matching_mr_solution(
        self,
        crystal,
        dmin,
        multiplier = 1,
        cached = False,
        template_equivalence = False,
        ):
        
        return [
            p.matching_mrpeak(
                crystal = crystal,
                dmin = dmin,
                cached = cached,
                multiplier = multiplier,
                template_equivalence = template_equivalence,
                )
            for p in self.peaks
            ]
    
    
    def history(self):
        
        if self.predecessor is None:
            return [ self ]
        
        
        else:
            return [ self ] + self.predecessor.history()
    
    
    def __str__(self):
        
        if self.peaks:
            centre = "\n".join( str( p ) for p in self.peaks )
        
        else:
            centre = "<EMPTY>"
            
        return centre
    
    
class XrayData(object):
    """
    X-ray data
    """
    
    def __init__(self, miller, fp, sigfp, cell, resolution):
        
        self.miller = miller
        self.fp = fp
        self.sigfp = sigfp
        self.cell = cell
        self.resolution = resolution
        
        
    def data_refl(self):
        
        return phaser.data_refl( self.miller, self.fp, self.sigfp )
    
    
class PhaserData(object):
    """
    Maximal collection of data items that may be needed by Phaser 
    """
    
    def __init__(self, miller, fp, sigfp, composition, cell, hall, resolution):
        
        self.miller = miller
        self.fp = fp
        self.sigfp = sigfp
        self.composition = composition
        self.cell = cell
        self.hall = hall
        self.resolution = resolution
        
    
    def apply(self, input):
         
        input.setREFL( self.miller, self.fp, self.sigfp )
        input.setTNCS_USE( False )
        input.setRESO_HIGH( self.resolution )
        self.composition.apply( input = input )
        self.apply_symmetry( input = input )
        
        
    def apply_symmetry(self, input):
        
        input.setCELL6( self.cell )
        input.setSPAC_HALL( self.hall )
        
        
    def uctbx_unit_cell(self):
        
        from cctbx import uctbx
        return uctbx.unit_cell( parameters = self.cell )
    
    
    def sgtbx_space_group(self):
        
        from cctbx import sgtbx
        return sgtbx.space_group( self.hall )
    
    
    def crystal_symmetry(self):
        
        from cctbx import crystal
        return crystal.symmetry(
            unit_cell = self.uctbx_unit_cell(),
            space_group = self.sgtbx_space_group(),
            )
            
    
class Problem(object):
    """
    Represents the mr problem
    """
    
    def __init__(self, hklin, labin, crystal_symmetry, resolution, no_tncs):
        
        self.xray_data = None
        self.space_group_hall = None
        self.hklin = hklin
        self.labin = labin
        self.crystal_symmetry = crystal_symmetry
        self.resolution = resolution
        self.no_tncs = no_tncs
        self.root = Structure( peaks = [] )
        self.data = Blackboard()
        self.pickled = None
        
        
    def load(self, logger):
        
        if self.pickled is not None:
            import pickle
            import gzip
            ( self.xray_data, self.space_group_hall ) = pickle.load(
                gzip.open( self.pickled )
                )
        
        else:
            from phaser.pipeline import brunett
            ( self.xray_data, self.space_group_hall ) = brunett.read_any_format_data(
                hklin = self.hklin,
                labin = self.labin,
                crystal_symmetry = self.crystal_symmetry,
                resolution = self.resolution,
                no_tncs = getattr( self, "no_tncs", False ),
                logger = logger,
                )
        
        
    def enpickle(self, root):
        
        import pickle
        import gzip
        
        with gzip.open( "%s_data.pkl.gz" % root, "wb" ) as fout:
            pickle.dump( ( self.xray_data, self.space_group_hall ), fout )
            self.pickled = fout.name
            
            
    def crystal(self):
        
        from cctbx import crystal, uctbx, sgtbx
        
        if not self.xray_data:
            raise RuntimeError, "Dataset not loaded"
        
        assert self.space_group_hall is not None
        
        return crystal.symmetry(
            unit_cell = uctbx.unit_cell( parameters = self.xray_data.cell ),
            space_group = sgtbx.space_group( self.space_group_hall )
            )
        
    
    def __getstate__(self):
        
        result = self.__dict__.copy()
        result[ "xray_data" ] = None
        result[ "space_group_hall" ] = None
        return result
        
        
    def __setstate__(self, dict):
        
        self.__dict__ = dict


class Case(object):
    """
    Holds all MRSets for a particular problem (data+symmetry combination)
    """
    
    def __init__(self, composition, space_group_hall, problem):
        
        self.composition = composition
        self._space_group_hall = space_group_hall
        self.problem = problem
        self.partials = set( [ self.problem.root ] )
        self.score_for = { self.problem.root: 0 }
        self.data = Blackboard()
        self._matching_crystal_info = None
        
        
    def calculation_data(self):
        
        return PhaserData(
            miller = self.problem.xray_data.miller,
            fp = self.problem.xray_data.fp,
            sigfp = self.problem.xray_data.sigfp,
            composition = self.composition,
            cell = self.problem.xray_data.cell,
            hall = self.space_group_hall(),
            resolution = self.problem.xray_data.resolution
            )
    
    
    def insert(self, structure, score):
        
        self.partials.add( structure )
        self.score_for[ structure ] = score
        
        
    def uctbx_unit_cell(self):
        
        from cctbx import uctbx
        return uctbx.unit_cell( parameters = self.problem.xray_data.cell )
    
    
    def space_group(self):
        
        from cctbx import sgtbx
        return sgtbx.space_group( self.space_group_hall() )
    
    
    def space_group_type(self):
        
        return self.space_group().type()
    
    
    def space_group_info(self):
        
        return self.space_group().info()
    
    
    def space_group_hall(self):
        
        return self._space_group_hall
    
    
    def space_group_number(self):
        
        return self.space_group_type().number()
    
    
    def space_group_symbol(self):
        
        return self.space_group_type().universal_hermann_mauguin_symbol()
    
    
    def space_group_suffix_symbol(self):
        
        return "".join(
            [ c for c in self.space_group_type().lookup_symbol() if not c.isspace() ]
            )
        
        
    def matching_crystal_info(self):
        
        if not self._matching_crystal_info:
            from phaser import matching
            self._matching_crystal_info = matching.CrystalInfo(
                space_group = self.space_group(),
                cell = self.uctbx_unit_cell()
                )
            
        return self._matching_crystal_info
    
    
    def create_matching_mr_peak(
        self,
        peak,
        multiplier = 1.0,
        cached = False,
        template_equivalence = False
        ):
        
        assert self.problem.xray_data is not None
        return peak.matching_mrpeak(
            crystal = self.matching_crystal_info(),
            dmin = self.problem.xray_data.resolution,
            cached = cached,
            multiplier = multiplier,
            template_equivalence = template_equivalence,
            )
        
        
    def create_matching_origin_search_mr_peak(
        self,
        peak,
        multiplier = 1.0,
        cached = False,
        template_equivalence = False
        ):
        
        from phaser import matching
        return matching.OriginSearchEntity(
            entity = self.create_matching_mr_peak(
                peak = peak,
                multiplier = multiplier,
                cached = cached,
                template_equivalence = template_equivalence,
                )
            )
        
        
    def create_matching_mr_structure(
        self,
        structure,
        multiplier = 1.0,
        cached = False,
        template_equivalence = False
        ):
        
        return [
            self.create_matching_mr_peak(
                peak = p,
                multiplier = multiplier,
                cached = cached,
                template_equivalence = template_equivalence
                )
            for p in structure.peaks
            ]
    
    
    def create_matching_origin_search_structure(
        self,
        structure,
        multiplier = 1.0,
        cached = False,
        template_equivalence = False
        ):
        
        return [
            self.create_matching_origin_search_mr_peak(
                peak = p,
                multiplier = multiplier,
                cached = cached,
                template_equivalence = template_equivalence
                )
            for p in structure.peaks
            ]
        

class PostProcessingData(object):
    """
    A marker class so that these could easily be deleted
    """        


class HistoricRotation(PostProcessingData):
    """
    A rotation that has been found in a previous round but not tested
    """
    
    def __init__(self, structure, peak):
        
        self.structure = structure
        self.peak = peak
        
        
class GenuineRotation(HistoricRotation):
    """
    A rotation found by the rotation function
    """
    
    
class ExtractedRotation(HistoricRotation):
    """
    A rotation extracted from a related peak (because of origin problems)
    """
    
    
class PotentialPeakFixedOriginEquivalence(object):
    """
    Equivalence relation for potential peaks with fixed origin
    """
    
    def __init__(self, case):
        
        self.case = case
        
        
    def __call__(self, left, right):
        
        for ( lstr, lpeak ) in left.associations():
            ltmf = self.case.create_matching_mr_peak( peak = lpeak )
            
            for ( rstr, rpeak ) in right.associations():
                if lstr != rstr:
                    continue
                
                rtmf = self.case.create_matching_mr_peak( peak = rpeak )
                
                if any( ltmf.operations_between( other = rtmf ) ):
                    return True
                
        return False
    
    
class PotentialPeakVariableOriginEquivalence(object):
    """
    Equivalence relation for potential peaks with variable origin
    """
    
    def __init__(self, case):
        
        self.case = case
        
        
    def __call__(self, left, right):
        
        for ( lstr, lpeak ) in left.associations():
            ltmf = self.case.create_matching_origin_search_mr_peak( peak = lpeak )
            
            for ( rstr, rpeak ) in right.associations():
                if lstr != rstr:
                    continue
                
                rtmf = self.case.create_matching_origin_search_mr_peak( peak = rpeak )
                
                if any( ltmf.potential_origins_between( other = rtmf ) ):
                    return True
                
        return False
        
        
class PotentialPeak(PostProcessingData):
    """
    A peak that is implicated without search
    """
    
    
class Amalgamation(PotentialPeak):
    """
    A different solution from the same starting point
    """
    
    entry_code = "A"
    
    def __init__(self, structures, origins):
        
        assert len( structures ) == 2
        self.structures = structures
        assert origins
        self.origins = origins
        
        
    def is_related(self, partial, ensemble):
        
        return (
            partial in self.structures
            and self.other( current = partial ).peaks[-1].ensemble == ensemble
            )
        
        
    def associations(self):
        
        return [
            ( s, self.other( current = s ).peaks[-1] ) for s in self.structures
            ]
        
        
    def is_defined_origin(self):
        
        return len( self.origins ) == 1
    
    
    def other(self, current):
        
        if current not in self.structures:
            raise ValueError, "%s is not member of %s" % ( current, self.structures )
        
        remaining = self.structures.difference( [ current ] )
        assert len( remaining ) == 1
        return list( remaining )[0]
    
    
    def get_handler_for(self, partial, index):
        
        return FillHandler(
            structure = partial,
            peak = self.other( current = partial ).peaks[-1],
            index = index,
            provider = self
            )

    
    def calculation(self, data, partial, peak, context):
        
        from phaser.pipeline import calculation
        return calculation.MultiOriginEvaluation(
            data = data,
            partial = partial,
            peak = peak,
            origins = self.origins
            )
        
        
    def process(self, result):
        
        ( is_pak, peak ) = result
        stats = ( "tf = %.3f, tfz = %.2f" % ( peak.tf, peak.tfz ) ) if is_pak else ""
        return FillResult( is_pak = is_pak, peak = peak, statistics = stats )
        
        
    def caption(self):
        
        return "Amalgamation peak (%d origins):" % len( self.origins )
    
    
    def involved(self):
        
        return "\n\n".join(
            "Structure %s:\n%s" % ( i, s ) for ( i, s )
            in enumerate( self.structures, start = 1 )
            )
    
        
    def __str__(self):
        
        return "%s\n\n%s" % ( self.caption(), self.involved() )
    
    
class ObjectCompletion(PotentialPeak):
    """
    A solution that would add a molecule according to some assembly description
    identified before
    """
    
    entry_code = "R"
    
    def __init__(self, structure, peak, strategy):
        
        self.structure = structure
        self.peak = peak
        self.strategy = strategy
        
        
    def is_related(self, partial, ensemble):
        
        return partial == self.structure and ensemble == self.peak.ensemble
    
    
    def associations(self):
        
        return [ ( self.structure, self.peak ) ]
    
    
    def is_defined_origin(self):
        
        return True
        
        
    def get_handler_for(self, partial, index):
        
        assert partial == self.structure
        return FillHandler(
            structure = partial,
            peak = self.peak,
            index = index,
            provider = self
            )

    
    def calculation(self, data, partial, peak, context):
        
        return self.strategy.calculation(
            data = data,
            partial = partial,
            peak = peak,
            context = context,
            )
        
        
    def process(self, result):
        
        return self.strategy.process( result = result )
    
        
    def caption(self):
        
        return "Assembly completion peak:"
    
    
    def __str__(self):
        
        return "%s\n\nStructure:\n%s\nPeak:%s" % ( self.caption(), self.structure, self.peak )
    
    
class FillHandler(object):
    """
    A unified interface for the different PotentialPeaks
    """
    
    def __init__(self, structure, peak, index, provider):
        
        self.structure = structure
        self.peak = peak
        self.index = index
        self.provider = provider
        
        
    def calculation(self, data, context):
        
        return self.provider.calculation(
            data = data,
            partial = self.structure,
            peak = self.peak,
            context = context
            )
    
    def process(self, result):
        
        return self.provider.process( result = result.data )
    
    
    def stage(self, parent, result):
        
        from phaser.pipeline import stage
        return stage.Filled(
            ensemble = parent,
            peak = result.peak,
            index = self.index,
            entry_code = self.provider.entry_code
            )
        
        
    def __str__(self):
        
        return self.provider.caption()
    

class FillResult(object):
    """
    A unified interface for results obtained from PotentialPeak evaluations
    """
    
    def __init__(self, is_pak, peak, statistics):
        
        self.is_pak = is_pak
        self.peak = peak
        self.statistics = statistics
    
    
class FillFailureRecord(object):
    """
    A generated solution that results in a crash. This is recorded so that it
    does not get regenerated every time
    """
    
    def __init__(self, structure, peak):
        
        self.structure = structure
        self.peak = peak
        
        
class ModelCollection(object):
    """
    A collection of models that need to be superposed before it can be used
    """
    
    COUNT = itertools.count()
    
    def __init__(self, pdbs, rmsds, trim, composition):
        
        assert pdbs
        assert len( pdbs ) == len( rmsds )
        assert all( os.path.isfile( p ) for p in pdbs )
        self.pdbs = pdbs
        self.rmsds = rmsds
        self.trim = trim
        self.composition = composition
        self.index = self.COUNT.next() + 1
        self.ensemble = None
        
        
    def score(self):
        
        return -self.rms()
    
    
    def pdb_objects(self):
        
        from phaser import tbx_utils
        return tbx_utils.read_in_files(
            file_names = self.pdbs,
            object_type = tbx_utils.PDBObject,
            )
        
        
    def rms(self):
        
        return min( self.rmsds )
        
    
    def __str__(self):
    
        return "%d superposable PDBs (#%d)" % ( len( self.pdbs ), self.index )
        
        
class Template(object):
    """
    Structural template for a certain component, with an alignment
    """
    
    def __init__(self, pdb_file, alignment, composition):
        
        assert os.path.isfile( pdb_file )
        self.pdb_file = pdb_file
        self.alignment = alignment
        self.composition = composition
        self.ensembles = None
        
        
    def root(self):
        
        import iotbx.pdb
        return iotbx.pdb.input( self.pdb_file ).construct_hierarchy()
    
    
    def pdb_object(self):

        from phaser import tbx_utils
        return tbx_utils.PDBObject.from_file( file_name = self.pdb_file )
    
    
    def suffix(self):
        
        ( head, tail ) = os.path.split( self.pdb_file )
        ( ali_head, ali_tail ) = os.path.split( self.alignment.name )
        return "%s_%s" % (
            os.path.splitext( tail )[0],
            os.path.splitext( ali_tail )[0],
            )
        
        
    def score(self):
        
        return self.identity()
    
    
    def identity(self):
        
        return self.alignment.object.identity_fraction()
    
    
    def __str__(self):
    
        return os.path.basename( self.pdb_file )
    
    
class BadHomologySearchEntry(object):
    """
    A homology search hit that cannot be used
    """
    
    def __init__(self, hss, reason):
        
        self.hss = hss
        self.reason = reason


class HomologySearchHit(object):
    """
    Homology search hit
    """
    
    def __init__(self, index, search, hit, composition):
        
        self.index = index
        self.search = search
        self.hit = hit
        self.composition = composition
        self.fetched = None
        
        
    def score(self):
        
        return -self.index
    
    
    def alignment_object(self):
        
        from phaser import tbx_utils
        ( head, tail ) = os.path.split( self.search )
        return tbx_utils.AlignmentObject(
            alignment = self.hit.alignment,
            name = "%s_hit%d.aln" % ( os.path.splitext( tail )[0], self.index )
            )
        
        
    def __str__(self):
        
        return "%s, hit %d" % ( self.search, self.index )
    
    
class Generator(object):
    """
    A rotation operator stored as ( fold, axis, centre )
    """
    
    TWOPI = 2.0 * math.pi
    
    def __init__(self, fold, axis, centre):
        
        assert 1 < fold
        self.fold = fold
        self.axis = axis
        self.centre = centre
        
        
    @property
    def rotation_angle(self):
        
        return self.rotation_angle_for( power = 1, fold = self.fold )
        
        
    @property
    def r(self):
        
        return self.rotation_matrix_for( axis = self.axis, angle = self.rotation_angle )
        
        
    @property
    def displacement(self):
        
        return -( self.r * self.centre ) + self.centre
    
    
    @property
    def rt(self):
        
        return scitbx.matrix.rt( ( self.r, self.displacement ) )
        
        
    def inverted_axis_setting(self):
        
        return self.__class__(
            fold = self.fold,
            axis = -self.axis,
            centre = self.centre
            )
        
        
    def transformed(self, operation):
        
        axis = scitbx.matrix.col( operation.r * self.axis )
        centre = scitbx.matrix.col( operation * self.centre )
        
        return self.__class__.new( fold = self.fold, axis = axis, centre = centre )
    
    
    def approx_equal_spatially(self, other, max_misalignment, max_spatial_error):
        
        if abs( self.axis.dot( other.axis ) ) < ( 1.0 - max_misalignment ):
            return False
            
        return ( self.centre - other.centre ).length() <= max_spatial_error
    
    
    def common_rotation_fold(self, other):
        
        return self.lcm( left = self.fold, right = other.fold )
    
    
    def nontrivial_series(self):
        
        rt = self.rt
        series = [ rt ]
        
        for index in range( 2, self.fold ):
            series.append( series[-1] * rt )
            
        return series
    
    
    def __str__(self):
        
        return "rotation: %d, axis %s, around %s" % (
            self.fold,
            "(%.3f, %.3f, %.3f)" % self.axis.elems,
            "(%.3f, %.3f, %.3f)" % self.centre.elems,
            )

    
    @classmethod
    def new(cls, fold, axis, centre):
        
        return cls(
            fold = fold,
            axis = axis,
            centre = centre - axis.dot( centre ) * axis
            )
    
    
    @classmethod
    def from_rotation_translation(cls, power, fold, axis, translation, spatial = 0.10):
        
        assert 1 <= power
        assert 0 < spatial
        parallel = axis.dot( translation )
        
        if spatial < abs( parallel ):
            raise ValueError, "Screw component above tolerance"
        
        perp = translation - parallel * axis
        
        r = cls.rotation_matrix_for(
            axis = axis,
            angle = cls.rotation_angle_for( power = power, fold = fold )
            )
        centre = cls.rotation_axis_row_echelon( r = r, perp = perp )
        
        return cls.new( fold = fold, axis = axis, centre = centre )
        
        
    @classmethod
    def from_scitbx_rt_operator(cls, op, max_fold = 100, angular = 0.005, spatial = 0.10):
        
        assert 0 < angular
        aaf = scitbx.math.r3_rotation_axis_and_angle_from_matrix( op.r )
        
        if abs( aaf.angle() ) < angular:
            raise ValueError, "Quasi-identity rotation"
        
        angle = aaf.angle() % cls.TWOPI
        assert 0 <= angle and angle < cls.TWOPI
        
        for fold in range( 1, max_fold + 1 ):
            rotation = fold * angle
            remainder = rotation % cls.TWOPI
            
            if remainder <= angular or ( cls.TWOPI - remainder ) <= angular:
                power = int( round( rotation / cls.TWOPI ) )
                break
            
        else:
            raise ValueError, "Foldedness outside limits"
            
        return cls.from_rotation_translation(
            power = power,
            fold = fold,
            axis = scitbx.matrix.col( aaf.axis ),
            translation = op.t,
            spatial = spatial,
            )
    
    
    @staticmethod
    def rotation_matrix_for(axis, angle):
        
        return scitbx.matrix.sqr(
            scitbx.math.r3_rotation_axis_and_angle_as_matrix( axis, angle )
            )
        
        
    @staticmethod
    def rotation_angle_for(power, fold):
        
        return Generator.TWOPI * power / fold
    
    
    @staticmethod
    def rotation_axis_row_echelon(r, perp):
        
        a = ( scitbx.matrix.identity( 3 ) - r ).as_flex_double_matrix()
        e = scitbx.math.row_echelon_full_pivoting(
            a_work = a,
            b_work = perp.as_flex_double_matrix().as_1d(),
            min_abs_pivot = 1e-12,
            )
        from scitbx.array_family import flex
        return scitbx.matrix.col(
            e.back_substitution(
                free_values = flex.double( e.nullity ),
                epsilon = 1e-12
                )
            )
        
    
    @staticmethod
    def lcm(left, right):
        
        import fractions
        gcd = fractions.gcd( left, right )
        return left * right // gcd
    
        
class PointGroup(object):
    """
    Base class for identification
    
    Also stands in for no symmetry 
    """
    
    IDENTITY = scitbx.matrix.rt(
        ( scitbx.matrix.identity( 3 ), scitbx.matrix.col( ( 0, 0, 0 ) ) )
        )
    
    def nontrivial_series(self):
        
        return []
    
    
    def series(self):
        
        return [ self.IDENTITY ] + self.nontrivial_series()
    
    
    def transformed(self, operation):
        
        return self
    
    
    def generator_classes(self):
        
        return []
    
    
    def generators(self):
        
        return []
    
    
    def approx_has_operation(self, operation, min_cos_angle, max_spatial_disp):
        
        inverse = operation.inverse()
        
        for op in self.nontrivial_series():
            trial = inverse * op
            
            if self.approx_identity_operation(
                operation = trial,
                min_cos_angle = min_cos_angle,
                max_spatial_disp = max_spatial_disp
                ):
                return True
            
        return False
    
    
    def approx_contains(self, other, min_cos_angle, max_spatial_disp):
        
        return all(
            self.approx_has_operation(
                operation = op,
                min_cos_angle = min_cos_angle,
                max_spatial_disp = max_spatial_disp,
                )
            for op in other.nontrivial_series()
            )
        
        
    def approx_equals(self, other, min_cos_angle, max_spatial_disp):
        
        return (
            len( other.series() ) == len( self.series() )
            and self.approx_contains(
                other = other,
                min_cos_angle = min_cos_angle,
                max_spatial_disp = max_spatial_disp,
                )
            )
        
        
    @staticmethod
    def approx_identity_operation(operation, min_cos_angle, max_spatial_disp):
        
        return (
            min_cos_angle <= scitbx.math.r3_rotation_cos_rotation_angle_from_matrix( operation.r )
            and operation.t.length() <= max_spatial_disp
            )
        
        
    def __str__(self):
        
        return "1"
    
    
    @classmethod
    def from_generators(cls, generators):
        
        if len( generators ) != 0:
            raise ValueError, "This point group type requires exactly zero generators"
        
        return cls()

    
class RotationAxis(PointGroup):
    """
    Point group symmetry, generated by rotation axes
    Currently limited to one axis only
    """
    
    def __init__(self, axis):
        
        self.axis = axis
        
        
    def nontrivial_series(self):
        
        return self.axis.nontrivial_series()
    
    
    def transformed(self, operation):
        
        return self.__class__(
            axis = self.axis.transformed( operation = operation )
            )
        
        
    def generator_classes(self):
        
        return [
            GeneratorClass(
                standard = self.axis,
                generators = [ self.axis ],
                operations = [ self.IDENTITY ],
                ),
            ]
    
    
    def generators(self):
        
        return [ self.axis ]
        
        
    def __str__(self):
        
        return str( self.axis.fold )
    
    
    @classmethod
    def from_generators(cls, generators):
        
        if len( generators ) != 1:
            raise ValueError, "This point group type requires exactly one generator"
        
        return cls( axis = generators[0] )
    
    
class EquivalenceRelation(object):
    """
    Partition members based on an equivalence relation into disjunct sets
    """
    
    def __init__(self, relation):
        
        self.relation = relation
        self.values = []
        self.equivalences = set()
        
        
    def add(self, value):
        
        for v in self.values:
            if self.relation( v, value ):
                self.equivalences.add( frozenset( ( v, value ) ) )
                
        self.values.append( value )
            
            
    def equivalents_of(self, value):
        
        return [ self.firstother( collection = relation, member = value )
            for relation in self.equivalences if value in relation ]
        
        
    def equivalence_group_of(self, value):
        
        group = set( [ value ] )
        leads = [ value ]
        
        while leads:
            next_round_leads = [] 
            
            for l in leads:
                equivs = self.equivalents_of( value = l )
                next_round_leads.extend( [ e for e in equivs if e not in group ] )
                group.update( equivs )
            
            leads = next_round_leads
            
        return group
    
    
    def classes(self):
        
        remaining = set( self.values )
        groups = []
        
        while remaining:
            first = remaining.pop()
            equivs = self.equivalence_group_of( value = first )
            groups.append( equivs )
            remaining.difference_update( equivs )
            
        return groups
                
    
    @staticmethod
    def firstother(collection, member):
        
        assert 2 <= len( collection )
        return iter( collection.difference( [ member ] ) ).next()
    
    
class EquivalenceRelationChain(object):
    """
    Create equivalence classes from pairwise relations
    """
    
    def __init__(self):
        
        self.classes = set()
        
        
    def add(self, relation):
        
        overlapping = [ c for c in self.classes if c.intersection( relation ) ]
        
        for cl in overlapping:
            self.classes.remove( cl )
                
        self.classes.add(
            frozenset(
                reduce(
                    operator.or_,
                    [ set( cl ) for cl in overlapping ],
                    set( relation ) 
                    )
                )
            )
    
    
class Anchor(object):
    """
    An anchor peak
    """
    
    def __init__(self, ensemble, peaks):
        
        self.ensemble = ensemble
        self.peaks = peaks
        
        
    def __str__(self):
        
        return "Anchor: %s, %d peak" % ( self.ensemble, len( self.peaks ) )
    
    
class RefineStrategy(object):
    """
    Refine around expected component
    """
    
    def calculation(self, data, partial, peak, context):
        
        from phaser.pipeline import calculation
        return calculation.RefinedEvaluation(
            data = data,
            partial = partial,
            peak = peak,
            protocol = context.b_factor_refinement
            )
        
        
    def process(self, result):
        
        ( is_pak, peak, llg ) = result
        stats = ( "tf = %.3f, tfz = %.2f, llg = %.2f" % ( peak.tf, peak.tfz, llg ) ) if is_pak else "",
        return FillResult( is_pak = is_pak, peak = peak, statistics = stats )
    
    
    def coverage(self):
        
        return 0
    
    
class SearchStrategy(object):
    """
    Restricted search around expected positions
    """
    
    def __init__(self, sweep, extent):
        
        self.sweep = sweep
        self.extent = extent
        
    
    def calculation(self, data, partial, peak, context):
        
        from phaser.pipeline import calculation
        return calculation.RestrictedSearchEvaluation(
            data = data,
            partial = partial,
            peak = peak,
            sweep = self.sweep,
            extent = self.extent,
            )
        
        
    def process(self, result):
        
        ( is_pak, peak ) = result
        stats = ( "tf = %.3f, tfz = %.2f" % ( peak.tf, peak.tfz ) ) if is_pak else "",
        return FillResult( is_pak = is_pak, peak = peak, statistics = stats )
    
    
    def coverage(self):
        
        return 1
    

class Assembly(object):
    """
    Any assembly, including additional symmetry
    
    The symmetry is stored so that any update will update the assembly as well
    """
    
    def __init__(self, definitions, symmetry, strategy):
        
        assert definitions, "Empty assembly"
        self.definitions = definitions
        self.symmetry = symmetry
        self.strategy = strategy
        
        
    @property
    def members(self):
        
        return self.definitions
    
    
    @property
    def ensembles(self):
        
        return [ d[0] for d in self.definitions ]
    
    
    @property
    def composition(self):
        
        return  reduce( operator.add, [ e.composition for e in self.ensembles ] ) * len( self.symmetry.series() )
    
    
    @property
    def operations(self):
        
        return [ d[1] for d in self.definitions ]
    
    
    @property
    def description(self):
        
        return self.definitions + self.symmetry_generated_part() 
        
        
    @property
    def anchors(self):
        
        complete = self.description
        
        for ( index, ( ense, transf ) ) in enumerate( self.definitions ):
            inverse = transf.inverse()
            
            yield Anchor(
                ensemble = ense,
                peaks = [ ( e, inverse * op ) for ( i, ( e, op ) )
                    in enumerate( complete ) if index != i  ]
                )
            
            
    @property
    def rms(self):
        
        weights = [ e.composition.weight() for e in self.ensembles ]
        total = sum( weights )
        assert 0 < total, "Zero weight for assembly"
        return 1.0 / total * sum(
            [ w * e.rms() for ( w, e ) in zip( weights, self.ensembles ) ]
            )
            
            
    def symmetry_generated_part(self):
        
        return reduce(
            operator.add,
            [ [ ( e, symop * op ) for ( e, op ) in self.definitions ]
                for symop in self.symmetry.nontrivial_series() ],
            []
            )
        
        
    def root(self):
        
        import iotbx.pdb
        from phaser import chisel
        
        new_model = iotbx.pdb.hierarchy.model()
        chain_ids = chisel.chain_ids()
        
        for ( ensemble, operation ) in self.definitions:
            root = ensemble.root()
            assert root.models() # process only first model
            
            for op in self.symmetry.series():
                for c in root.models()[0].chains():
                    copy = c.detached_copy()
                    copy.id = chain_ids.next()
                    atoms = copy.atoms()
                    coords = atoms.extract_xyz()
                    atoms.set_xyz( op * operation * coords )
                    new_model.append_chain( copy )
                    
        new_root = iotbx.pdb.hierarchy.root()
        new_root.append_model( new_model )
        return new_root
    
    
    def chain_model_groups(self):
        
        for ( ensemble, operation ) in self.definitions:
            for cmg in ensemble.chain_model_groups():
                for op in self.symmetry.series():
                    copy = cmg.copy()
                    copy.transform( rt = operation )
                    yield copy
    
    
    def info(self):
        
        return "%d components, symmetry: %s" % (
            len( self.definitions ),
            self.symmetry,
            )
            
            
    def __str__(self):
        
        return "Assembly: %s" % self.info()
    
    
    @classmethod
    def Simple(cls, ensemble, symmetry, strategy):
        
        return cls(
            definitions = [ ( ensemble, PointGroup.IDENTITY ) ],
            symmetry = symmetry,
            strategy = strategy,
            )


class GeneratorSetting(object):
    """
    A generator with a basis model
    """
    
    def __init__(self, ensemble, generator):
        
        self.ensemble = ensemble
        self.generator = generator
        
        
    def __str__(self):
        
        return "ensemble: %s, operator: %s" % ( self.ensemble, self.generator )
        
        
class ModelRelation(object):
    """
    Stores information about geometric relations between models
    """
    
    def __init__(self, setting, relating):
        
        self.setting = setting
        self.relating = relating
        

class StructureRelations(object):
    """
    Geometric relations for models within a structure
    """
    
    def __init__(self, structure, relations):
        
        self.structure = structure
        self.relations = relations
        
        
class GeneratorCluster(object):
    """
    Generator cluster
    """
    
    def __init__(self, average):
        
        self.average = average
        self.relations_from = {}
        
        
    def add(self, structure, relation):
        
        self.relations_from.setdefault( structure, [] ).append( relation )
        
        
class GeneratorClass(object):
    """
    Class of equivalent generators in a given symmetry
    """
    
    def __init__(self, standard, generators, operations):
        
        assert len( generators ) == len( operations )
        self.standard = standard
        self.generators = generators
        self.operations = operations


class SymmetryAnalysis(object):
    """
    A combination of generators that give some point group symmetry
    """
    
    def __init__(self, setting, relations_from, formed_by):
        
        self.setting = setting
        self.relations_from = relations_from
        self.formed_by = set( formed_by )
        
    
    def contains(self, other):
        
        return self.formed_by.issuperset( other.formed_by )
    
    
    def __str__(self):
        
        return str( self.setting )
        

class SymmetrySetting(object):
    """
    A symmetry object associated with an ensemble to define the basis
    """
    
    def __init__(self, ensemble, point_group):
        
        self.ensemble = ensemble
        self.point_group = point_group
        
        
    def __str__(self):
        
        return str( self.point_group )
    
    
class SymmetrySetRelationAnalysis(object):
    """
    Analysis of containment relations based on operators
    """
    
    def __init__(self, analysis, contained_by, contains):
        
        self.analysis = analysis
        self.contained_by = contained_by
        self.contains = contains
        
        
    def __str__(self):
        
        return str( self.analysis )


class NewSymmetry(object):
    """
    Marker that a new symmetry has been created from a SymmetrySetRelationAnalysis
    """
    
    def __init__(self, analysis, assemblies_in):
        
        self.analysis = analysis
        self.assemblies_in = assemblies_in
        
        
    @property
    def setting(self):
        
        return self.analysis.analysis.setting
        
        
    def __str__(self):
        
        return str( self.analysis )


class ObsoletedBySuperSymmetry(object):
    """
    Marker that a SymmetrySetting has been removed
    """
    
    def __init__(self, setting, supersymmetry):
        
        self.setting = setting
        self.supersymmetry = supersymmetry
        
        
class SymmetryDowngrade(object):
    """
    Marker that a SymmetrySetting has been removed
    """
    
    def __init__(self, setting, highest):
        
        self.setting = setting
        self.highest = highest
        
        
class UnsupportedSymmetry(object):
    """
    Marker that a SymmetrySetting has been removed
    """
    
    def __init__(self, setting):
        
        self.setting = setting


class SymmetryUpdateRecord(object):
    """
    Marker that a SymmetrySetting has been updated to something else
    """
    
    def __init__(self, setting, previous, assemblies_in):
        
        self.setting = setting
        self.previous = previous
        self.assemblies_in = assemblies_in
        

class NewAssemblyRecord(object):
    """
    Marker that a new Assembly has been created
    """
    
    def __init__(self, assembly):
        
        self.assembly = assembly


class AssemblyRemovalRecord(object):
    """
    Marker that an Assembly has been removed
    """
    
    def __init__(self, assembly):
        
        self.assembly = assembly


class ForSymmetryAnalysis(object):
    """
    Marks that symmetry analysis should be performed on a structure
    """
    
    def __init__(self, structure, index):
        
        self.structure = structure
        self.index = index