import scitbx.matrix

class Superposition(object):
    """
    Superposition data
    """
    
    UNIT = None
    
    def __init__(self, transformation, rmsd):
        
        self.transformation = transformation
        self.rmsd = rmsd
        
        
    def inverse(self):
        
        return self.__class__(
            transformation = self.transformation.inverse(),
            rmsd = self.rmsd
            )
        
    @classmethod
    def unit(cls):
        
        if cls.UNIT is None:
            cls.UNIT = cls(
                transformation = scitbx.matrix.rt(
                    ( ( 1, 0, 0, 0, 1, 0, 0, 0, 1 ), ( 0, 0, 0 ) )
                    ),
                rmsd = 0.0
                )
            
        return cls.UNIT

        
class Proxy(object):
    """
    Superposition using a certain protocol
    """
    
    def __init__(self, method):
        
        self.method = method
        self.cache = {}
        self.failures = set()
        
        
    def superpose(self, reference, moving, stream):
        
        stream.write( "Superposition:\n" )
        stream.write( "    Reference: %s\n" % reference )
        stream.write( "    Moving: %s\n\n" % moving )
        
        refid = reference.superpose_template()
        movid = moving.superpose_template()
        
        if refid == movid:
            stream.write( "Identical (template equivalent) molecules\n" )
            superposition = Superposition.unit()
            
        elif ( refid, movid ) in self.failures:
            stream.write( "Failure found in cache\n" )
            raise RuntimeError, "Superposition failure"
            
        elif ( refid, movid ) in self.cache:
            stream.write( "Result found in cache\n" )
            superposition = self.cache[ ( refid, movid ) ]
        
        elif ( movid, refid ) in self.cache:
            stream.write( "Result found in cache (inverted)\n" )
            superposition = self.cache[ ( movid, refid ) ].inverse()
            self.cache[ ( refid, movid ) ] = superposition
            
        else:
            stream.write( "Running %s...\n" % self.method.NAME )
            
            try:
                superposition = self.method( reference = reference, moving = moving )
                
            except RuntimeError, e:
                self.failures.add( ( refid, movid ) )
                self.failures.add( ( movid, refid ) )
                raise RuntimeError, e
            
            self.cache[ ( refid, movid ) ] = superposition
            
        stream.write( "Rmsd: %.3f\n" % superposition.rmsd )
        
        stream.write( "Rotation:\n" )
        rotm = superposition.transformation.r.elems
        stream.write( "    %7.4f %7.4f %7.4f\n" % rotm[0:3] )
        stream.write( "    %7.4f %7.4f %7.4f\n" % rotm[3:6] )
        stream.write( "    %7.4f %7.4f %7.4f\n" % rotm[6:9] )
        stream.write(
            "Translation: %.4f %.4f %.4f\n" % superposition.transformation.t.elems
            )
        
        return superposition
    
    
    def substitute(self, peak, ensemble, cell, stream):
       
        stream.write( "Substituting ensemble in peak:\n" )
        stream.write( "    %s\n" % peak )
        stream.write( "with ensemble:\n" )
        stream.write( "    %s\n" % ensemble )
        
        superposition = self.superpose(
            reference = peak.ensemble,
            moving = ensemble,
            stream = stream
            )
        superposed = peak.as_scitbx_rt_operator( cell = cell ) * superposition.transformation
        from phaser.pipeline import mr_object
        substituted = mr_object.Peak(
            ensemble = ensemble,
            rotation = mr_object.Rotation.Matrix( elements = superposed.r.elems ),
            translation = cell.fractionalize( superposed.t.elems ),
            bfactor = peak.bfactor
            )
        
        stream.write( "Substituted peak:\n" )
        stream.write( "    %s\n" % substituted )
        
        return substituted
    
    
def simple_ssm_superposition(reference, moving):
            
    try:
        import ccp4io_adaptbx
        
    except ImportError, e:
        raise RuntimeError, "SSM failure: Cannot import ccp4io_adaptbx: %s" % e
    
    try:
        ssm = ccp4io_adaptbx.SecondaryStructureMatching(
            moving = moving.root(),
            reference = reference.root()
            )
        
    except RuntimeError, e:
        raise RuntimeError, "SSM superposition failed: %s" % e
    
    tm = ssm.get_matrix()
    
    return Superposition(
        transformation = scitbx.matrix.rt(
            ( ( tm[0:3] + tm[4:7] + tm[8:11] ), ( tm[3], tm[7], tm[11] ) )
            ),
        rmsd = ssm.ssm.rmsd
        )

simple_ssm_superposition.NAME = "SSM"