from __future__ import division

from phaser import rsam
from phaser import residue_substitution

from iotbx.pdb.hierarchy import root, model, chain, residue_group, atom_group, atom
from scitbx.array_family import flex
import scitbx.math

import operator
import math


class ForwardIndexer(object):
    """
    Class to lookup a descriptor object from the actual object
    """
    
    def __init__(self, root):
        
        self.arrays = [
            dict( [ ( Entity.Root( obj = root ), 0 ) ] ),
            dict(
                [ ( Entity.Model( obj = e ), i ) for ( i, e )
                    in enumerate( root.models() ) ]
                ),
            dict(
                [ ( Entity.Chain( obj = e ), i ) for ( i, e )
                    in enumerate( root.chains() ) ]
                ),
            dict(
                [ ( Entity.ResidueGroup( obj = e ), i ) for ( i, e )
                    in enumerate( root.residue_groups() ) ]
                ),
            dict(
                [ ( Entity.AtomGroup( obj = e ), i )  for ( i, e )
                    in enumerate( root.atom_groups() ) ]
                ),
            dict(
                [ ( Entity.Atom( obj = e ), i ) for ( i, e )
                    in enumerate( root.atoms() ) ]
                ),
            ]


class ReverseIndexer(object):
    """
    Class to lookup from the actual object from a deferred one
    """
    
    def __init__(self, root):
        
        self.arrays = [
            [ root ],
            list( root.models() ),
            list( root.chains() ),
            list( root.residue_groups() ),
            list( root.atom_groups() ),
            list( root.atoms() ),
            ]
        
        
class MMDBIndexer(object):
    """
    From mmdb selection to index
    """
    
    CREATION = [
        ( 1, "models" ),
        ( 2, "chains" ),
        ( 3, "residue_groups" ),
        ( 4, "atom_groups" ),
        ( 5, "atoms" ),
        ]
    
    def __init__(self, root):
        
        self.arrays = [
            { (): [ 0 ] },
            {},
            {},
            {},
            {},
            {}
            ]
        
        for ( level, method ) in self.CREATION:
            for ( index, obj ) in enumerate( getattr( root, method )() ):
                key = Entity( obj = obj, level = level ).specifiers()
                self.arrays[ level ].setdefault( key, [] ).append( index )
                
                
def to_coot_s_expression(root):
        
    expression = []
    
    for m in root.models():
        m_expr = []
        
        for c in m.chains():
            rg_expr = []
            
            for rg in c.residue_groups():
                a_expr = [
                    [
                        [ a.name, a.parent().altloc.strip() ],
                        [ a.occ, a.b, a.element ],
                        list( a.xyz ),
                        ]
                    for a in rg.atoms()
                    ]
                rg_expr.append(
                    [
                        rg.resseq_as_int(),
                        rg.icode.strip(),
                        rg.atom_groups()[0].resname,
                        a_expr
                        ]
                    )
            m_expr.append( [ c.id, rg_expr ] )
            
        expression.append( m_expr )
        
    return expression
        
            
class Descriptor(object):
    """
    Descriptor to portably identify an entity unambiguously 
    """
    
    def __init__(self, level, index):
        
        self.level = level
        self.index = index
        
        
    def as_entity(self, indexer):
        
        return Entity(
            obj = indexer.arrays[ self.level ][ self.index ],
            level = self.level
            )
        
        
    def params(self):
        
        return ( self.level, self.index )
    
    
    def __hash__(self):
        
        return hash( self.params() )
    
    
    def __eq__(self, other):
        
        if not isinstance( other, self.__class__ ):
            return False
        
        return self.params() == other.params()
    
    
    def __ne__(self, other):
        
        return not( self == other )
    
    
    @classmethod
    def from_specifiers(cls, level, specifiers, indexer):
        
        return [ cls( level = level, index = i )
            for i in indexer.arrays[ level ][ specifiers ] ]
        

class Entity(object):
    """
    Hashing, naming and selectors for iotbx.pdb.hierarchy.* objects
    """
    
    IOTBX_HIERARCHY = [ root, model, chain, residue_group, atom_group, atom ]
    LEVEL_FOR = dict( [ ( c, i ) for ( i, c ) in enumerate( IOTBX_HIERARCHY ) ] )
    NAMES = [
        ( "", lambda h: () ),
        ( "id = '%s'", lambda h: h.id ),
        ( "id = '%s'", lambda h: h.id ),
        ( "resid = '%s'", lambda h: h.resid() ),
        ( "resname = '%s', altloc = '%s'", lambda h: ( h.resname, h.altloc ) ),
        ( "name = '%s'", lambda h: h.name ),
        ]
    SELECTORS = [
        ( "all", lambda h: () ),
        ( "model '%s'", lambda h: h.id ),
        ( "chain '%s'", lambda h: h.id ),
        ( "resseq '%s' and icode '%s'", lambda h: ( h.resseq, h.icode ) ),
        ( "altid '%s'", lambda h: h.altloc ),
        ( "name '%s'", lambda h: h.name ),
        ]
    SPECIFIERS = [
        lambda h: (),
        lambda h: ( h.id, ),
        lambda h: ( h.id, ),
        lambda h: ( ( h.resseq, h.icode ), ),
        lambda h: ( h.altloc, ),
        lambda h: ( h.name, ),
        ]
    COOT_SPECIFIERS = [
        lambda h: (),
        lambda h: (),
        lambda h: ( h.id, ),
        lambda h: ( h.resseq_as_int(), h.icode.strip() ),
        lambda h: (),
        lambda h: ( h.name, h.parent().altloc.strip() ),
        ]
    
    def __init__(self, obj, level):
        
        assert isinstance( obj, self.IOTBX_HIERARCHY[ level ] )
        self._object = obj
        self._level = level
        
        
    def object(self):
        
        return self._object
    
    
    def level(self):
        
        return self._level
    
    
    def parent(self):
        
        if self._level == 0:
            raise ValueError, "root has no parent"
        
        parent = self._object.parent()
        
        if not parent:
            return None
        
        return self.__class__( obj = parent, level = self._level - 1 )
    
    
    def iotbx_class(self):
        
        return self.IOTBX_HIERARCHY[ self._level ]
    
    
    def as_descriptor(self, indexer):
        
        return Descriptor(
            level = self._level,
            index = indexer.arrays[ self._level ][ self ]
            )
        
        
    def level_identifier(self):
        
        return self.format( data = self.NAMES )
    
    
    def level_selector(self):
        
        return self.format( data = self.SELECTORS )
    
    
    def level_specifier(self):
        
        return self.SPECIFIERS[ self._level ]( self._object )
    
    
    def level_coot_selector(self):
        
        return self.COOT_SPECIFIERS[ self._level ]( self._object )
    
    
    def level_name(self):
        
        return "%s (%s)" % (
            self.IOTBX_HIERARCHY[ self._level ].__name__,
            self.level_identifier()
            )
    
    
    def selector(self):
        
        return " and ".join( [ e.level_selector() for e in self.hierarchy() ] )
    
    
    def specifiers(self):
        
        return reduce(
            operator.add,
            [ e.level_specifier() for e in self.hierarchy() ],
            (),
            )
        
    
    def coot_selector(self):
        
        return reduce(
            operator.add,
            [ e.level_coot_selector() for e in self.hierarchy() ],
            (),
            )
    
    
    def hierarchy(self):
        
        h = [ self ]
        current = self
        
        while 0 < current._level:
            parent = current.parent()
                
            if not parent:
                break
            
            h.append( parent )
            current = parent
            
        return reversed( h )
    
    
    def is_connected(self):
        
        top = self.hierarchy().next()
        
        return top.level() == 0
        
        
    def format(self, data):
        
        ( format, method ) = data[ self._level ]
        return format % method( self._object )
    
    
    def __hash__(self):
        
        return hash( self._object.memory_id() )
    
    
    def __eq__(self, other):
        
        return self._object.memory_id() == other._object.memory_id()
    
    
    def __ne__(self, other):
        
        return not ( self == other )
    
    
    def __str__(self):
        
        if self._level == 0:
            return self.level_name()
        
        h = list( self.hierarchy() )
        
        if h[0]._level == 0:
            h = h[1:]
            
        return ", ".join( [ e.level_name() for e in h ] )
    
    
    @classmethod
    def Create(cls, obj):
        
        if obj.__class__ not in cls.LEVEL_FOR:
            raise ValueError, "Unknown entity class"
        
        return cls( obj = obj, level = cls.LEVEL_FOR[ obj.__class__ ] )
    
    
    @classmethod
    def Root(cls, obj):
        
        return cls( obj = obj, level = cls.LEVEL_FOR[ root ] )
    
    
    @classmethod
    def Model(cls, obj):
        
        return cls( obj = obj, level = cls.LEVEL_FOR[ model ] )
    
    
    @classmethod
    def Chain(cls, obj):
        
        return cls( obj = obj, level = cls.LEVEL_FOR[ chain ] )
    
    
    @classmethod
    def ResidueGroup(cls, obj):
        
        return cls( obj = obj, level = cls.LEVEL_FOR[ residue_group ] )
    
    
    @classmethod
    def AtomGroup(cls, obj):
        
        return cls( obj = obj, level = cls.LEVEL_FOR[ atom_group ] )
    
    
    @classmethod
    def Atom(cls, obj):
        
        return cls( obj = obj, level = cls.LEVEL_FOR[ atom ] )
    
    
class Instruction(object):
    """
    Generic instruction related functionality
    """
    
    def __ne__(self, other):
        
        return not ( self == other )
    
    
class DeleteInstruction(Instruction):
    """
    Deletes a residue group
    """
    
    METHOD_FOR = {
        chain: "remove_residue_group",
        residue_group: "remove_atom_group",
        atom_group: "remove_atom",
        }
    
    def apply_to(self, entity):
         
        parent = entity.parent()
        
        if not parent:
            raise RuntimeError, "%s has no parent" % entity
        
        getattr( parent.object(), self.METHOD_FOR[ parent.iotbx_class() ] )( entity.object() )
        
        
    def log(self, entity):
        
        return "Delete %s" % entity
    
    
    def parseable_form(self):
        
        return ( self.__class__.__name__, () )
    
    
    def is_structural_change(self):
        
        return True
    
    
    def __hash__(self):
        
        return hash( self.__class__ )
    
    
    def __eq__(self, other):
        
        return isinstance( other, self.__class__ )
    
    
    @classmethod
    def from_params(cls, params):
        
        return cls()
        
        
class SetInstruction(Instruction):
    """
    Set a property
    """
    
    def __init__(self, setter):
        
        self.apply_to = setter
        
        
    def log(self, entity):
        
        return "Set '%s' to '%s' on %s" % ( self.apply_to.description + ( entity, ) )
    
    
    def property(self):
        
        return self.apply_to.description[0]
    
    
    def parseable_form(self):
        
        return ( self.__class__.__name__, self.apply_to.description )
    
    
    def is_structural_change(self):
        
        return ( 
            ( self.property() == "b-factor" )
            or ( self.property() == "element" )
            )
        
        
    def __hash__(self):
        
        return hash( ( self.__class__, self.apply_to.description ) )
    
    
    def __eq__(self, other):
        
        if not isinstance( other, self.__class__ ):
            return False
        
        return self.apply_to.description == other.apply_to.description
    
    
    @classmethod
    def Resseq(cls, value):
        
        def setter(entity):
        
            entity.object().resseq = value
        
        setter.description = ( "resseq", value )
        
        return cls( setter = setter )
    
    
    @classmethod
    def Icode(cls, value):
        
        def setter(entity):
        
            entity.object().icode = value
        
        setter.description = ( "icode", value )
        
        return cls( setter = setter )
    
    
    @classmethod
    def Resname(cls, value):
        
        def setter(entity):
        
            entity.object().resname = value
        
        setter.description = ( "resname", value )
        
        return cls( setter = setter )
    
    
    @classmethod
    def Bfactor(cls, value):
        
        def setter(entity):
        
            entity.object().set_b( value )
        
        setter.description = ( "b-factor", value )
        
        return cls( setter = setter )
    
    
    @classmethod
    def Element(cls, value):
        
        def setter(entity):
        
            entity.object().set_element( value )
        
        setter.description = ( "element", value )
        
        return cls( setter = setter )
    
    
    @classmethod
    def Hetero(cls, value):
        
        def setter(entity):
        
            entity.object().set_hetero( value )
        
        setter.description = ( "hetero", value )
        
        return cls( setter = setter )
    
    
    @classmethod
    def Name(cls, value):
        
        def setter(entity):
        
            entity.object().set_name( value )
        
        setter.description = ( "name", value )
        
        return cls( setter = setter )
    
    
    @classmethod
    def from_params(cls, params):
        
        constructor_for = {
            "resseq": cls.Resseq,
            "icode": cls.Icode,
            "resname": cls.Resname,
            "b-factor": cls.Bfactor,
            "element": cls.Element,
            "hetero": cls.Hetero,
            "name": cls.Name,
            }
        
        return constructor_for[ params[0] ]( value = params[1] )
        
        
class AttachInstruction(Instruction):
    """
    Attaches atoms to an atom_group
    """
    
    METHOD_FOR = {
        atom_group: "append_atom",
        }
    
    def __init__(self, data):
        
        self.data = data
        
    
    def apply_to(self, entity):
        
        child = self.data.as_iotbx_object()
        getattr( entity.object(), self.METHOD_FOR[ entity.iotbx_class() ] )( child )
    
    
    def log(self, entity):
        
        return "Attach %s to %s" % ( self.data.identifier(), entity )
    
    
    def parseable_form(self):
        
        return (
            self.__class__.__name__,
            ( self.data.__class__.__name__, self.data.parseable_form() ),
            )
        
        
    def is_structural_change(self):
        
        return True
    
    
    def __hash__(self):
        
        return hash( ( self.__class__, self.data ) )
        
        
    def __eq__(self, other):
        
        if not isinstance( other, self.__class__ ):
            return False
        
        return self.data == other.data
        
        
    @classmethod
    def from_params(cls, params):
        
        return cls( data = getattr( rsam, params[0] )( *params[1] ) )
    
    
class ResetIsotropicInstruction(Instruction):
    """
    Deletes uij from an entity
    """
    
    def apply_to(self, entity):
        
        entity.object().uij_erase()
    
    
    def log(self, entity):
        
        return "Reset %s to isotropic" % entity
    
    
    def parseable_form(self):
        
        return ( self.__class__.__name__, () )
    
    
    def is_structural_change(self):
        
        return True
    
    
    def __hash__(self):
        
        return hash( self.__class__ )
    
    
    def __eq__(self, other):
        
        return isinstance( other, self.__class__ )
    
    
    @classmethod
    def from_params(cls, params):
        
        return cls()
    

# Data types
class ChainSample(object):
    """
    A chain to be modified
    """
    
    def __init__(self, chain, mmt):
        
        self.chain = chain
        self.mmt = mmt
        self.store = {}
        
        
    def atom_radii(self):
        
        key = ( "radii", )
        
        if key not in self.store:
            self.store[ key ] = rsam.VanDerWaalsRadius(
                root = self.chain,
                radius_for_mainchain_atom = dict(
                    zip(
                        self.mmt.mainchain_atom_names,
                        self.mmt.mainchain_atom_radii )
                    ),
                sidechain_data_for = dict(
                    zip(
                        self.mmt.three_letter_codes,
                        self.mmt.sidechain_atom_radii_dicts
                        )
                    ),
                unknown_radius = self.mmt.unknown_atom_radius
                )
        
        return self.store[ key ]
        
        
    def accessible_surface_area(self, radius, precision):
        
        key = ( "asa", radius, precision )
        
        if key not in self.store:
            self.store[ key ] = rsam.AccessibleSurfaceArea(
                root = self.chain,
                vdw_radii = self.atom_radii()
                )
            
        return self.store[ key ]
    
    
    def identity(self):
        
        return None


class AlignedChainSample(ChainSample):
    """
    A chain to be modified
    """
    
    def __init__(self, chain, alignment, index, mmt, min_hssp_length, min_matching_fraction):
        
        aligned = rsam.match_chain_to_alignment(
            chain = chain,
            mtype = mmt,
            sequence = alignment.alignments[ index ],
            gap = alignment.gap,
            min_hssp_length = min_hssp_length,
            min_matching_fraction = min_matching_fraction
            )
        
        assert len( aligned ) == alignment.length()
        super( AlignedChainSample, self ).__init__( chain = chain, mmt = mmt )
        
        self.alignment = alignment
        self.index = index
        self.aligned = aligned
        self.unaligned = rsam.chain_residue_groups_difference(
            chain = self.chain,
            residue_groups = self.aligned
            )
        
    
    def target_sequence(self):
        
        return self.alignment.alignments[0]
    
    
    def model_sequence(self):
        
        return self.alignment.alignments[ self.index ]
    
    
    def pairwise_alignment(self):
        
        import iotbx.bioinformatics
        return iotbx.bioinformatics.alignment(
            names = [ "TARGET", "MODEL" ],
            alignments = [ self.target_sequence(), self.model_sequence() ]
            )
    
    
    def sequence_similarity(self, matrix, unknown):
             
        key = ( "similarity", matrix, unknown )
        
        if key not in self.store:
            if matrix not in self.mmt.similarity_matrix_suite.names():
                raise RuntimeError, "Unknown matrix %s for type %s" % (
                    matrix,
                    self.mmt.name
                    )
            m = self.mmt.similarity_matrix_suite.matrix_for( name = matrix )
            scores = residue_substitution.SequenceSimilarity(
                alignment = self.alignment,
                matrix = m,
                unknown = unknown
                )
            ( scale, base ) = m.normalization()
            self.store[ key ] = [
                ( value - base ) / scale for value in scores.values
                ]
        
        return self.store[ key ]
    
    
    def identity(self):
        
        return self.pairwise_alignment().identity_fraction()
    

class Blueprint(object):
    """
    Object to contain deferred modification data
    """
    
    def __init__(self):
        
        self.data = []
        
        
    def add(self, datum):
        
        self.data.append( datum )
        
        
    def instructions(self, r_indexer):
        
        instructions = []
        
        for ( i, d ) in self.data:
            e = d.as_entity( indexer = r_indexer )
            instructions.append( ( i, e ) )
            
        return instructions
    
    
    def formatted_modifications(self, r_indexer):
        
        modifications = []
        
        for ( i, d ) in self.data:
            specs = d.as_entity( indexer = r_indexer ).specifiers()
            modifications.append( ( i.parseable_form(), d.level, specs ) )
            
        return modifications
    
    
    @classmethod
    def from_instructions(cls, instructions, indexer):
        
        bp = cls()
        
        for ( i, e ) in instructions:
            bp.add( datum = ( i, e.as_descriptor( indexer = indexer ) ) )
            
        return bp
    
                
# Sculptor class to modify chains
class Sculptor(object):
    """
    Coordinates mainchain, sidechain and b-factor processing
    """
    
    def __init__(
        self,
        deletions,
        polishings,
        renumber,
        prunings,
        completions,
        rename,
        bfactors,
        minimum_b
        ):
        
        self.deletions = deletions
        self.polishings = polishings
        self.renumber = renumber
        
        self.prunings = prunings
        self.completions = completions
        self.rename = rename
        
        self.bfactors = bfactors
        self.minimum_b = minimum_b
        
    
    def run(self, sample, logger):
        
        moribund_rgs = self.run_mainchain_design( sample = sample )
        ( resseq_for, icode_for ) = self.renumber.calculate( sample = sample )
        ( moribund_atoms, resname_for, hetero_resets, name_for, element_for, generated_for ) = (
            self.run_sidechain_design( sample = sample )
            )
        ( bfactor_for, isotropic_resets, b_shift ) = self.run_bfactor_design( sample = sample )
        
        instructions = []
        
        for residue_group in sample.chain.residue_groups(): 
            rg = Entity.ResidueGroup( obj = residue_group )
            
            if rg in resseq_for:
                instructions.append(
                    ( SetInstruction.Resseq( value = resseq_for[ rg ] ), rg )
                    )
            
            if rg in icode_for: 
                instructions.append(
                    ( SetInstruction.Icode( value = icode_for[ rg ] ), rg ) 
                    )
            
            if rg in moribund_rgs:
                instructions.append( ( DeleteInstruction(), rg ) )
                
            for atom_group in residue_group.atom_groups():
                ag = Entity.AtomGroup( obj = atom_group )
                
                if ag in resname_for: 
                    instructions.append(
                        ( SetInstruction.Resname( value = resname_for[ ag ] ), ag )
                        )
                    
                if ag in generated_for:
                    for d in generated_for[ ag ]:
                        d.b += b_shift
                        instructions.append(
                            ( AttachInstruction( data = d ), ag )
                            )
                    
                for atom in atom_group.atoms():
                    a = Entity.Atom( obj = atom )
                    
                    if a in moribund_atoms:
                        instructions.append( ( DeleteInstruction(), a ) )
                    
                    if a in name_for: 
                        instructions.append(
                            ( SetInstruction.Name( value = name_for[ a ] ), a ),
                            )
                        
                    if a in element_for: 
                        instructions.append(
                            ( SetInstruction.Element( value = element_for[ a ] ), a )
                            )
                        
                    if a in bfactor_for:
                        instructions.append(
                            ( SetInstruction.Bfactor( value = bfactor_for[ a ] ), a )
                            )
                        
                    if a in isotropic_resets:
                        instructions.append( ( ResetIsotropicInstruction(), a ) )
                        
                    if a in hetero_resets:
                        instructions.append(
                            ( SetInstruction.Hetero( value = False ), a )
                            )
        
        return instructions
    
    
    def run_deletion_algorithms(self, sample):
        
        if not self.deletions:
            return set()
        
        # Run algorithms
        result = self.collate_alignment_frame_algorithm_results(
            algorithms = self.deletions,
            sample = sample,
            reducer = any
            )
        
        moribund = set(
            [ Entity.ResidueGroup( obj = rg ) for ( rg, mb )
                in zip( sample.aligned, result )
                if rg and mb ]
            )
        
        # Add unaligned residues
        moribund.update(
            [ Entity.ResidueGroup( obj = rg ) for rg in sample.unaligned ]
            )
        return moribund
        
        
    def run_mainchain_design(self, sample):
        
        moribund = self.run_deletion_algorithms( sample = sample )
        
        # Polishing algorithms modify moribund residues sequentially
        for algorithm in self.polishings:
            ( delete, restore ) = algorithm.calculate(
                sample = sample,
                moribund = moribund
                )
             
            moribund.difference_update( restore )    
            moribund.update( delete )
        
        return moribund
    
    
    def run_pruning_algorithms(self, sample):
        
        if not self.prunings:
            return {}
        
        result = self.collate_alignment_frame_algorithm_results(
            algorithms = self.prunings,
            sample = sample,
            reducer = min
            )
        
        level_for = dict(
            [ ( Entity.ResidueGroup( obj = rg ), l ) for ( rg, l )
                in zip( sample.aligned, result )
                if rg and l ]
            )
        unaligned = max(
            sample.mmt.topology_cache.get_mainchain_distance_for(
                sidechain = sample.mmt.unknown_sidechain_topology
                ) \
                .from_atom.values()
                )
        level_for.update(
            [ ( Entity.ResidueGroup( obj = rg ), unaligned ) for rg in sample.unaligned ]
            )
        
        return level_for
        
    
    def run_sidechain_design(self, sample):
        
        sidechain_for = dict(
            zip( sample.mmt.three_letter_codes, sample.mmt.sidechain_topologies )
            )
        unknown = sample.mmt.unknown_sidechain_topology
        ( resname_for, hetero_resets ) = self.rename.calculate( sample = sample )
        cache = sample.mmt.topology_cache
        max_level_for = self.run_pruning_algorithms( sample = sample )
        moribund = set()
        name_for = {}
        element_for = {}
        generated_for = {}
        mainchain_atoms = set( sample.mmt.mainchain_atom_names )
        generators = [ g for g in sample.mmt.sidechain_completion
            if g.NAME in self.completions ]
        
        for rg in sample.chain.residue_groups():
            d_rg = Entity.ResidueGroup( obj = rg )
            ag_altcoded = {}
            do_pruning = d_rg in max_level_for
            level = max_level_for.get( d_rg, None )
            resname = None
            
            for ag in rg.atom_groups():
                d_ag = Entity.AtomGroup( obj = ag )
                sc = sidechain_for.get( ag.resname, unknown )
                ca_dist = cache.get_mainchain_distance_for( sidechain = sc )
                
                do_mapping = d_ag in resname_for
                
                if do_mapping:
                    am = cache.get_mapping_between(
                        source = sc,
                        target = sidechain_for.get( resname_for[ d_ag ], unknown )
                        )
                else:
                    am = {}
                
                atom_named = {}
            
                for atom in ag.atoms():
                    d_atom = Entity.Atom( obj = atom )
                    
                    # Skip for mainchain atoms
                    if atom.name in mainchain_atoms:
                        atom_named[ atom.name ] = atom
                        continue
                    
                    # Pruning
                    if do_pruning and level < ca_dist.from_atom.get( atom.name, level + 1 ):
                        moribund.add( d_atom )
                        
                    # Remapping - no match
                    elif do_mapping and atom.name not in am.mapping:
                        moribund.add( d_atom )
                            
                    # Remapping - different name
                    elif do_mapping and atom.name != am.mapping[ atom.name ]:
                        tname = am.mapping[ atom.name ]
                        name_for[ d_atom ] = tname
                        element = tname.strip()[0]
                        
                        if atom.element.strip() != element:
                            element_for[ d_atom ] = element
                             
                        atom_named[ tname ] = atom
                            
                    # Remapping - same name or no remapping
                    else:
                        atom_named[ atom.name ] = atom
            
                ag_altcoded[ ag.altloc ] = (
                    d_ag,
                    atom_named,
                    resname_for.get( d_ag, ag.resname )
                    )
            
            # Propagate common atoms
            if "" in ag_altcoded and len( ag_altcoded ) != 1:
                ( descr, common_named, rn ) = ag_altcoded[ "" ]
                del ag_altcoded[ "" ]
                
                for ( desc, atom_named, rn ) in ag_altcoded.values():
                    atom_named.update( common_named )
            
            # Sidechain completion
            for ( descr, atom_named, resname ) in ag_altcoded.values():
                generated_for[ descr ] = self.do_sidechain_completion(
                    generators = generators,
                    atom_named = atom_named,
                    max_level = level,
                    resname = resname
                    )
            
        return ( moribund, resname_for, hetero_resets, name_for, element_for, generated_for )
    
    
    def run_bfactor_design(self, sample):
        
        individual = [ p.calculate( sample = sample ) for p in self.bfactors ]
        length = len( sample.chain.atoms() )
        assert all( length == len( i ) for i in individual )
        
        cumulative = [ sum( ind ) for ind in zip( *individual ) ]
        
        if cumulative:
            lowest = min( cumulative )
                
        else:
            lowest = self.minimum_b
            
        b_shift = max( self.minimum_b - lowest, 0 )
        cumulative = [ c + b_shift for c in cumulative ]
            
        bfactor_for = {}
        isotropic_resets = set()
        
        for ( at, b ) in zip( sample.chain.atoms(), cumulative ):
            descr = Entity.Atom( obj = at )
            
            if 0.01 < abs( at.b - b ):
                bfactor_for[ descr ] = b
                
            if at.uij_is_defined():
                isotropic_resets.add( descr )
        
        return ( bfactor_for, isotropic_resets, b_shift )

    
    def processable(self, sample_type):
        
        return all(
            [
                all( issubclass( sample_type, a.TYPE ) for a in self.deletions ),
                all( issubclass( sample_type, a.TYPE ) for a in self.polishings ),
                issubclass( sample_type, self.renumber.TYPE ),
                all( issubclass( sample_type, a.TYPE ) for a in self.prunings ),
                issubclass( sample_type, self.rename.TYPE ),
                all( issubclass( sample_type, a.TYPE ) for a in self.bfactors ),
                ]
            )
        
        
    @staticmethod
    def collate_alignment_frame_algorithm_results(algorithms, sample, reducer):
        
        length = len( sample.aligned )
        
        individuals = [ a.calculate( sample = sample ) for a in algorithms ]
        assert all( len( i ) == length for i in individuals )
        return [ reducer( equi ) for equi in zip( *individuals ) ]
    
    
    @staticmethod
    def do_sidechain_completion(generators, atom_named, max_level, resname):
        
        generated = []
        
        for gen in generators:
            result = gen(
                atom_named = atom_named,
                resname = resname,
                max_level = max_level
                )
            generated.extend( result.generated )
            atom_named.update(
                dict(
                    [
                        ( a.name, a ) for a in
                        [ d.as_iotbx_object() for d in result.generated ]
                        ]
                    )
                )
        
        return generated
    
    
class Discarder(object):
    """
    Sculptor-like class that deletes every residue except the ones named
    """
    
    def __init__(self, keep):
        
        self.keep = set( keep )
    
    
    def processable(self, sample_type):
        
        return issubclass( sample_type, ChainSample )
    
    
    def run(self, sample, logger):
        
        instructions = []
        
        for rg in sample.chain.residue_groups():
            for ag in rg.atom_groups():
                if ag.resname.strip() not in self.keep:
                    instructions.append(
                        ( DeleteInstruction(), Entity.AtomGroup( obj = ag ) )
                        )
                    
        return instructions


# Some common functionality        
class LinearAveragedSequenceSimilarity(object):
    """
    Sequence similarity with linear averaging
    """
    
    def __init__(self, matrix, unknown, window, weights):
        
        assert 2 * window + 1 == len( weights )
        norm = flex.sum( weights )
        assert 0 < norm
        self.matrix = matrix
        self.unknown = unknown
        self.weights = weights / norm
        self.window = window
        
    
    def scores_for(self, sample):
        
        matrix = sample.mmt.similarity_matrix_suite.matrix_for( name = self.matrix )
        ( scale, base ) = matrix.normalization()
        
        localized = rsam.localize(
            sequence = sample.sequence_similarity(
                matrix = self.matrix,
                unknown = self.unknown
                ),
            width = self.window,
            padding = ( matrix.gap_penalty - base ) / scale,
            )
        
        return [ flex.sum( flex.double( subseq ) * self.weights )
            for subseq in localized ]
    
    
class SpaceAveragedSequenceSimilarity(object):
    """
    Sequence similarity averaged over space
    """
    
    def __init__(self, matrix, unknown, maxdist):
        
        assert 0 <= maxdist 
        self.matrix = matrix
        self.unknown = unknown
        self.maxdist = maxdist
        
        
    def scores_for(self, sample):
        
        positions = self.calculate_mean_positions( sample = sample )
        scores = sample.sequence_similarity(
            matrix = self.matrix,
            unknown = self.unknown
            )
        assert len( positions ) == len( scores )
        
        ( points, values ) = self.calculate_ss_map(
            positions = positions,
            scores = scores
            )
        assert len( points ) == len( values )
        
        averaged = []
        
        for ( position, score ) in zip( positions, scores ):
            if not position:
                averaged.append( score )
                continue
            
            averaged.append(
                self.calculate_average(
                    position = position,
                    points = points,
                    scores = values
                    )
                )
            
        return averaged
    
    
    def calculate_mean_positions(self, sample):
        
        mainchain = set( sample.mmt.mainchain_atom_names )
        selector = lambda a: a.name in mainchain
        positions = []
        
        for rg in sample.aligned:
            if not rg:
                positions.append( None )
                continue
                
            accepted = flex.vec3_double(
                [ a.xyz for a in rg.atoms() if selector( a ) ]
                )
            
            if not accepted:
                positions.append( None )
            
            else:
                positions.append( accepted.mean() )
        
        return positions
    
    
    def calculate_ss_map(self, positions, scores):
        
        points = flex.vec3_double()
        values = flex.double()
        
        for ( p, s ) in zip( positions, scores ):
            if not p:
                continue
            
            points.append( p )
            values.append( s )
            
        return ( points, values )
    
    
    def calculate_average(self, position, points, scores):
        
        distance_sqs = rsam.distance_sqs_from(
            centre = position,
            coords = points
            )
        selection = distance_sqs <= ( self.maxdist ** 2 )
        assert 0 < len( selection.iselection() )
        weights = self.calculate_weights(
            distances = flex.sqrt( distance_sqs.select( selection ) )
            )
        norm = flex.sum( weights )
        return flex.sum( scores.select( selection ) * weights ) / norm
        
        
    def calculate_weights(self, distances):
        
        return flex.exp( -distances / ( self.maxdist / 2.0 ) ) 


# Mainchain deletion algorithms
class MDAGap(object):
    """
    Include residue if present in the search sequence
    """
    
    TYPE = AlignedChainSample
    
    def calculate(self, sample):
        
        return [ t == sample.alignment.gap for t in sample.target_sequence() ]
        
        
class MDASimilarity(object):
    """
    Include residue if the sequence similarity is higher than a threshold
    """
    
    TYPE = AlignedChainSample
    
    def __init__(self, ss_calc, threshold):
        
        self.ss_calc = ss_calc
        self.threshold = threshold
        
    
    def calculate(self, sample):
        
        return [
            s < self.threshold for s in self.ss_calc.scores_for( sample = sample )
            ]
        
        
class MDAThresholdAdjustSimilarity(object):
    """
    Include residue if the sequence similarity is higher than a threshold
    """
    
    TYPE = AlignedChainSample
    
    def __init__(self, ss_calc, offset):
        
        self.ss_calc = ss_calc
        self.offset = offset
        
    
    def calculate(self, sample):
        
        gap_count = MDAGap().calculate( sample = sample ).count( True )
        model_count = len(
            [ p for p in zip( sample.target_sequence(), sample.model_sequence() )
                if None not in p ]
            )
        target = gap_count - self.offset * model_count
        scores = self.ss_calc.scores_for( sample = sample )
        threshold = self.get_threshold( scores = scores, target = target )
        
        return [ s < threshold for s in scores ]

    
    def get_threshold(self, scores, target):
        
        assert scores
        ordered = sorted( set( scores ) )
        previous = ( ordered[0], 0 )
        
        for threshold in ordered[1:]:
            current = ( threshold, len( [ s for s in scores if s < threshold ] ) )
            
            if target < current[1]:
                break
            
            previous = current
            
        else:
            current = ( ordered[-1] + 0.01, len( scores ) )
        
        if target - previous[1] <= current[1] - target:
            return previous[0]
        
        else:
            return current[0]
        
        
class MDARemoveLong(object):
    """
    Remove residues only if the deletion is long 
    """
    
    TYPE = AlignedChainSample
    
    def __init__(self, min_length):
        
        self.min_length = min_length
    
    
    def calculate(self, sample):
        
        deletion = []
        
        for contiguous in self.split_alignment( sample = sample ):
            length = len( contiguous )
            verdict = [ False ]
            
            if sample.alignment.gap in contiguous: 
                if self.min_length <= length:
                    verdict = [ True ]
            
            deletion.extend( verdict * length )
            
        return deletion
    
    
    def split_alignment(self, sample):
        
        gap = sample.alignment.gap
        return rsam.split(
            sequence = sample.target_sequence(),
            consecutivity = lambda left, right: (
                right == gap if left == gap else right != gap
                )
            )
        

# Mainchain polishing algorithms
class MPARemoveShort(object):
    """
    Removes residues that would map to a short segment on the model structure
    """
    
    TYPE = ChainSample
    
    def __init__(self, min_length):
        
        assert 0 < min_length
        self.min_length = min_length
        
        
    def calculate(self, sample, moribund):
        
        present = [ rg for rg in sample.chain.residue_groups()
            if Entity.ResidueGroup( obj = rg ) not in moribund ]
        consec = sample.mmt.consecutivity
        delete = []
        
        for contiguous in rsam.split( sequence = present, consecutivity = consec ):
            if len( contiguous ) < self.min_length:
                delete.extend(
                    [ Entity.ResidueGroup( obj = rg ) for rg in contiguous ]
                    )
            
        return ( delete, [] )
    
    
class MPAKeepRegular(object):
    """
    Reinstates consecutive residues up to max_length if they are in regular
    secondary structure 
    """
    
    TYPE = ChainSample
    
    def __init__(self, max_length):
        
        assert 0 < max_length
        self.max_length = max_length
        
        
    def calculate(self, sample, moribund):
        
        sec_str = self.get_secondary_structure(
            mmt = sample.mmt,
            chain = sample.chain
            )
        
        consec = sample.mmt.consecutivity
        reinstate = []
        
        for sse in sec_str.elements:
            mb = [ rg for rg in sse if Entity.ResidueGroup( obj = rg ) in moribund ]
            
            for segment in rsam.split( sequence = mb, consecutivity = consec ):
                if len( segment ) <= self.max_length:
                    reinstate.extend(
                        [ Entity.ResidueGroup( obj = rg ) for rg in segment ]
                        )
                    
        return ( [], reinstate )
    
    
    def get_secondary_structure(self, mmt, chain):
        
        if not mmt.secondary_structure_analyser:
            raise RuntimeError, "No secondary structure analyser available for %s" % mmt.name
        
        try:
            sec_str = mmt.secondary_structure_analyser( chain = chain ) 
        
        except ValueError, e:
            raise RuntimeError, "Error in secondary structure analysis: %s" % e
        
        return sec_str


# Mainchain renumbering algorithms
class MRAOriginal(object):
    """
    Algorithm that does not do anything
    """
    
    TYPE = ChainSample
    
    def calculate(self, sample):
        
        return ( {}, {} )
    
    
class MRASequence(object):
    """
    Renumbers residue_groups according to sequence
    """
    
    TYPE = ChainSample
    
    def __init__(self, enquire, start):
        
        self.enquire = enquire
        self.start = start
    
    
    def calculate(self, sample):
         
        sequence = self.enquire( sample )
        gap = sample.alignment.gap
        
        resseq = self.start - 1
        icode = " "
        
        instructions = []
        resseq_for = {}
        icode_for = {}
        
        for ( rescode, rg ) in zip( sequence, sample.aligned ):
            if rescode != gap:
                resseq += 1
                icode = " "
            
            else:
                if icode == " " or icode == "Z":
                    icode = "A"
                
                else:
                    icode = chr( ord( icode ) + 1 )
                    
            if not rg:
                continue
            
            ent = Entity.ResidueGroup( obj = rg )
            
            if rg.resseq_as_int() != resseq:
                resseq_for[ ent ] = resseq
                
            if rg.icode != icode:
                icode_for[ ent ] = icode
                
        return ( resseq_for, icode_for )
    
    
    @classmethod
    def Target(cls, start):
        
        return cls(
            enquire = AlignedChainSample.target_sequence,
            start = start
            )
    
    
    @classmethod
    def Model(cls, start):
        
        return cls(
            enquire = AlignedChainSample.model_sequence,
            start = start
            )
        

# Renaming algorithms
class SROriginal(object):
    """
    Algorithm that does not do anything
    """
    
    TYPE = ChainSample
    
    def calculate(self, sample):
        
        return ( {}, set() )
    
    
class SRSequence(object):
    """
    Rename according to target sequence
    """
    
    TYPE = AlignedChainSample
    
    def calculate(self, sample):
        
        three_letter_for = dict(
            zip( sample.mmt.one_letter_codes, sample.mmt.three_letter_codes )
            )
        resname_for = {}
        hetero_resets = set()
        
        # Rename aligned residues according to sequence
        for ( rescode, rg ) in zip( sample.target_sequence(), sample.aligned ):
            if not rg:
                continue
            
            if rescode == sample.alignment.gap:
                resname = "GAP"
            
            else:
                resname = three_letter_for.get(
                    rescode,
                    sample.mmt.unknown_three_letter
                    )
            
            ( rn_for, het_res ) = self.atom_group_name_settings(
                rg = rg,
                resname = resname
                )
            resname_for.update( rn_for )
            hetero_resets.update( het_res )
                    
        # Rename unaligned residues as unknowns
        for rg in sample.unaligned:
            ( rn_for, het_res ) = self.atom_group_name_settings(
                rg = rg,
                resname = sample.mmt.unknown_three_letter
                )
            resname_for.update( rn_for )
            hetero_resets.update( het_res )
            
        return ( resname_for, hetero_resets )
            
            
    @staticmethod
    def atom_group_name_settings(rg, resname):
        
        resname_for = {}
        hetero_resets = set()
        
        for ag in rg.atom_groups():
            if ag.resname != resname:
                resname_for[ Entity.AtomGroup( obj = ag ) ] = resname
                
            for a in ag.atoms():
                if a.hetero:
                    hetero_resets.add( Entity.Atom( obj = a ) )
                
        return ( resname_for, hetero_resets )
      
      
class SRSingleLetterTarget(object):
    """
    Rename according to target sequence, but equivalence is tested on single-
    letter residue codes
    """
    
    TYPE = AlignedChainSample
    
    def calculate(self, sample):
        
        three_letter_for = dict(
            zip( sample.mmt.one_letter_codes, sample.mmt.three_letter_codes )
            )
        one_letter_for = dict(
            zip( sample.mmt.three_letter_codes, sample.mmt.one_letter_codes )
            )
        resname_for = {}
        hetero_resets = set()
        
        # Rename aligned residues according to sequence
        for ( rescode, rg ) in zip( sample.target_sequence(), sample.aligned ):
            if not rg:
                continue
            
            ( rn_for, het_res ) = self.atom_group_name_settings(
                rg = rg,
                rescode = rescode,
                gap_char = sample.alignment.gap,
                one_letter_for = one_letter_for,
                unknown_one_letter = sample.mmt.unknown_one_letter,
                three_letter_for = three_letter_for,
                unknown_three_letter = sample.mmt.unknown_three_letter,
                )
            resname_for.update( rn_for )
            hetero_resets.update( het_res )
                    
        # Rename unaligned residues as unknowns
        for rg in sample.unaligned:
            ( rn_for, het_res ) = self.atom_group_name_settings(
                rg = rg,
                rescode = sample.mmt.unknown_one_letter,
                gap_char = sample.alignment.gap,
                one_letter_for = one_letter_for,
                unknown_one_letter = sample.mmt.unknown_one_letter,
                three_letter_for = three_letter_for,
                unknown_three_letter = sample.mmt.unknown_three_letter,
                )
            resname_for.update( rn_for )
            hetero_resets.update( het_res )
            
        return ( resname_for, hetero_resets )
            
            
    @staticmethod
    def atom_group_name_settings(
        rg,
        rescode,
        gap_char,
        one_letter_for,
        unknown_one_letter,
        three_letter_for,
        unknown_three_letter,
        ):
        
        resname_for = {}
        hetero_resets = set()
        
        for ag in rg.atom_groups():
            d_ag = Entity.AtomGroup( obj = ag )
            
            if rescode == gap_char:
                resname_for[ d_ag ] = "GAP"
                
            elif one_letter_for.get( ag.resname ) != rescode:
                resname_for[ d_ag ] = three_letter_for.get(
                    rescode,
                    unknown_three_letter,
                    )
                
            for a in ag.atoms():
                if a.hetero:
                    hetero_resets.add( Entity.Atom( obj = a ) )
                
        return ( resname_for, hetero_resets )
    

# Sidechain pruning algorithms
class SidechainPruningAlgorithm(object):
    """
    Some generic features
    """
    
    @classmethod
    def residue_max_bond_distance(cls, sidechain_for, cache, rescode):
        
        if rescode in sidechain_for:
            sc = sidechain_for[ rescode ]
            dist = cache.get_mainchain_distance_for( sidechain = sc ).from_atom.values()
            
            if dist:
                return max( dist )
            
            else:
                return 0
            
        else:
            return 99 # Unknown sidechain
    
    
    @classmethod
    def residue_min_bond_distance(cls, sidechain_for, cache, rescode):
        
        if rescode in sidechain_for:
            sc = sidechain_for[ rescode ]
            dist = cache.get_mainchain_distance_for( sidechain = sc ).from_atom.values()
            
            if dist:
                return min( dist )
            
            else:
                return 0
            
        else:
            return 1 # Unknown sidechain
    

class SPASchwarzenbacher(SidechainPruningAlgorithm):
    """
    Truncates sidechain if residues types do not agree
    """
    
    TYPE = AlignedChainSample
    
    def __init__(self, level):
        
        assert 0 <= level
        self.level = level
        
        
    def calculate(self, sample):
        
        sidechain_for = dict(
            zip( sample.mmt.one_letter_codes, sample.mmt.sidechain_topologies )
            )
        cache = sample.mmt.topology_cache
        
        return [
            self.pruning_level_for(
                target = t,
                is_identical = ( t == m ),
                sidechain_for = sidechain_for,
                cache = cache
                )
            for ( t, m ) in zip( sample.target_sequence(), sample.model_sequence() )
            ]
        
        
    def pruning_level_for(self, target, is_identical, sidechain_for, cache):
        
        if is_identical:
            return self.residue_max_bond_distance(
                sidechain_for = sidechain_for,
                cache = cache,
                rescode = target
                )
        
        else:
            return self.level
        
        
class SPASimilarity(SidechainPruningAlgorithm):
    """
    Sequence similarity defined sidechain truncation:
        Full length if upper_limit < conservation,
        Full truncation if conservation < lower_limit,
        Predefined truncation if upper_limit < conservation < lower_limit
    """
    
    TYPE = AlignedChainSample
    
    def __init__(self, ss_calc, lower, upper, level):
        
        assert 0 <= level and lower <= upper
        
        self.ss_calc = ss_calc
        self.lower = lower
        self.upper = upper
        self.level = level
        
        
    def calculate(self, sample):
        
        sidechain_for = dict(
            zip( sample.mmt.one_letter_codes, sample.mmt.sidechain_topologies )
            )
        cache = sample.mmt.topology_cache
        
        return [
            self.pruning_level_for(
                target = t,
                similarity = s,
                sidechain_for = sidechain_for,
                cache = cache
                )
            for ( t, s ) in zip(
                sample.target_sequence(),
                self.ss_calc.scores_for( sample = sample )
                ) ]
    
    
    def pruning_level_for(self, target, similarity, sidechain_for, cache):
        
        if self.upper <= similarity:
            return self.residue_max_bond_distance(
                sidechain_for = sidechain_for,
                cache = cache,
                rescode = target
                )
            
        elif self.lower <= similarity:
            return self.level

        else:
            return self.residue_min_bond_distance(
                sidechain_for = sidechain_for,
                cache = cache,
                rescode = target
                )

class Scaled(object):
    """
    Algorithms transforming their primary result
    """

    def __init__(self, factor):
        
        self.factor = factor
        
            
    def transform(self, primary):
        
        return [ self.factor * value for value in primary ]


class BPAOriginalBfactor(Scaled):
    """
    Original bfactor with linear transformation
    """
    
    TYPE = ChainSample
    
    def calculate(self, sample):
        
        return self.transform( primary = [ a.b for a in sample.chain.atoms() ] )
    
    
class BPAAccessibleSurfaceArea(Scaled):
    """
    Accessible surface area with transformation
    """
    
    TYPE = ChainSample
    
    def __init__(self, probe, precision, factor):
        
        assert 0 < probe
        assert 0 < precision
        self.probe = probe
        self.precision = precision
        super( BPAAccessibleSurfaceArea, self ).__init__( factor = factor )
        
        
    def calculate(self, sample):
        
        asa = sample.accessible_surface_area(
            radius = self.probe,
            precision = self.precision
            )
        
        return self.transform( primary = list( asa.values ) )
    
    
class BPASequenceSimilarity(Scaled):
    """
    Sequence similarity score with transformation
    """
    
    TYPE = AlignedChainSample
    
    def __init__(self, ss_calc, factor):
        
        super( BPASequenceSimilarity, self ).__init__( factor = factor )
        self.ss_calc = ss_calc
        
    
    def calculate(self, sample):
        
        scores = self.ss_calc.scores_for( sample = sample )
        assert len( scores ) == len( sample.aligned )
        
        score_for = dict(
            [ ( Entity.ResidueGroup( obj = rg ), score )
                for ( rg, score ) in zip( sample.aligned, scores ) if rg
                ]
            )
        
        for rg in sample.unaligned:
            descr = Entity.ResidueGroup( obj = rg )
            assert descr not in score_for
            score_for[ descr ] = 0 # set value for unaligned residues
        
        return self.transform(
            primary = [ score_for[ Entity.ResidueGroup( a.parent().parent() ) ]
                for a in sample.chain.atoms() ]
            )


def chain_ids():
    
    import itertools
    import string
    allchars = string.ascii_uppercase + string.ascii_lowercase + string.digits
    it = itertools.chain(
        allchars,
        itertools.product( string.ascii_uppercase, allchars ),
        itertools.product( string.ascii_lowercase, allchars ),
        itertools.product( string.digits, allchars ),
        )
    
    for comb in it:
        yield "".join( comb )