import scitbx.matrix
from scitbx.array_family import flex

import math
        
class EnsembleInfo(object):
    """
    Information about ensemble extent, centre, symmetry
    """
    
    def __init__(self, root):
        
        import cctbx.eltbx.tiny_pse
        atoms = root.atoms()
        coords = atoms.extract_xyz()
        elems = [ a.determine_chemical_element_simple() for a in atoms ]
        weights = atoms.extract_occ() * flex.double(
            [ cctbx.eltbx.tiny_pse.table( e if e else " C", True ).weight()
                for e in elems ]
            )
        self.principal = PrincipalComponents( coords = coords, weights = weights )
        self.centre = self.principal.centre
        ( txs, tys, tzs ) = ( self.principal.r * coords + self.principal.t ).parts()
        self.extent = 1.0 / 3.0 * flex.sum(
            0.5 * flex.double(
                [ max( txs ) - min( txs ), max( tys ) - min( tys ), max( tzs ) - min( tzs ) ]
                )
            )
        from phaser import ncs
        ensym = ncs.EnsembleSymmetry()
        ensym.set_root( root = root )
        self.pgops = [ scitbx.matrix.rt( ( op.get_rotation(), op.get_translation() ) )
            for op in ensym.get_point_group().get_operators() ]
        
        
class MapEnsembleInfo(object):
    """
    Information about ensemble defined with electron density
    """
    
    def __init__(self, centre, extent):
        
        self.centre = centre
        self.extent = sum( extent ) / len( extent )
        
        from phaser import ncs
        self.pgops = [
            scitbx.matrix.rt(
                ( ncs.Operator.unity.get_rotation(), ncs.Operator.unity.get_translation() )
                )
            ]
        
        
class PrincipalComponents(object):
    """
    Principal components of inertia
    """
    
    def __init__(self, coords, weights):
        
        total = flex.sum( weights )
        self.centre = 1.0 / total * scitbx.matrix.col( ( weights * coords ).sum() )
        ( xs, ys, zs ) = ( coords - self.centre ).parts()
        xx = flex.sum( xs * xs * weights )
        xy = flex.sum( xs * ys * weights )
        xz = flex.sum( xs * zs * weights )
        yy = flex.sum( ys * ys * weights )
        yz = flex.sum( ys * zs * weights )
        zz = flex.sum( zs * zs * weights )
        
        from scitbx.linalg import eigensystem
        eig = eigensystem.real_symmetric(
            ( yy + zz, xx + zz, xx + yy, -xy, -xz, -yz )
            )
        eiv = eig.vectors()
        inertia = scitbx.matrix.sqr(
            tuple( eiv[6:9] ) + tuple( eiv[3:6] ) + tuple( eiv[0:3] )
            )
        
        if inertia.determinant() < 0:
            inertia = -inertia
            
        self.r = inertia.elems
        self.t = -inertia * self.centre
        
        
class CrystalInfo(object):
    """
    Symmetry information
    """
    
    def __init__(self, space_group, cell, dmin):
        
        self.cell = cell
        self.space_group = space_group
        self.to_frac = scitbx.matrix.sqr( cell.fractionalization_matrix() )
        self.to_orth = scitbx.matrix.sqr( cell.orthogonalization_matrix() )
        self.sgops = space_group.all_ops()
        self.frac_rotations = [
            scitbx.matrix.sqr( op.r().as_double() ) for op in self.sgops
            ]
        self.orth_rotations = [
            self.to_orth * m * self.to_frac for m in self.frac_rotations
            ]
        self.dmin = dmin
        self.cb_op = self.space_group.type().cb_op()
        self.to_reference_setting = scitbx.matrix.sqr(
            self.cb_op.c().r().as_double()
            )
        self.from_reference_setting = scitbx.matrix.sqr(
            self.cb_op.c_inv().r().as_double()
            )
        
        
class PeakInfo(object):
    """
    Information about a placed ensemble 
    """
    
    def __init__(self, transformation, ensemble, cell):
        
        self.transformation = transformation
        self.ensemble = ensemble
        self.centre_orth = transformation * ensemble.centre
        self.centre_frac = scitbx.matrix.col(
            cell.fractionalize( self.centre_orth )
            )
        
        
    def translated(self, shift, cell):
        
        shiftop = scitbx.matrix.rt(
            ( ( 1, 0, 0, 0, 1, 0, 0, 0, 1 ), cell.orthogonalize( shift ) )
            ) 
        
        return self.__class__(
            transformation = shiftop * self.transformation,
            ensemble = self.ensemble,
            cell = cell
            )
        
        
    @classmethod
    def from_phaser_solution(cls, rotation, translation, ensemble, cell):
        
        return cls(
            transformation = scitbx.matrix.rt(
                ( rotation, cell.orthogonalize( translation ) )
                ),
            ensemble = ensemble,
            cell = cell
            )
        
        
class MatchTolerance(object):
    """
    Tolerances for equivalence
    """
    
    MT_FOR = {}
    
    def __init__(self, rotation, translation):
        
        self.rotation = rotation
        self.translation = translation
        
        
    def __mul__(self, other):
        
        if not isinstance( other, ( int, float ) ):
            raise NotImplemented
        
        return self.__class__(
            rotation = self.rotation * other,
            translation = self.translation * other,
            )
        
        
    def __rmul__(self, other):
        
        return self.__mul__( other = other )
        
        
    @classmethod
    def from_ensemble(cls, crystal, ensemble):
        
        key = ( crystal, ensemble )
        
        if key not in cls.MT_FOR:
            cls.MT_FOR[ key ] = cls(
                rotation = 2.0 * math.atan( crystal.dmin / ( 4.0 * ensemble.extent ) ),
                translation = crystal.dmin / 5.0,
                )
            
        return cls.MT_FOR[ key ]
    
    
def no_origin_shift():
    """
    No origin shifts between the two structures
    """
    
    yield ( 0, 0, 0 )
        
        
def plausible_origin_shifts(left, right, crystal, multiplier):
    """
    Origin shifts consistent with space group symmetry that link at least one
    molecule to another one
    """
    
    from cctbx import sgtbx
    semis = sgtbx.structure_seminvariants( crystal.space_group )
    flags = sgtbx.search_symmetry_flags( True, 0, True )
    ssym = sgtbx.search_symmetry( flags, crystal.space_group.type(), semis )
    euclidian = CrystalInfo(
        space_group = ssym.subgroup(),
        cell = crystal.cell,
        dmin = crystal.dmin
        )
    
    # Space group reference setting
    reference_space_group = sgtbx.space_group(
        crystal.space_group.change_basis( crystal.cb_op )
        )
    reference_ssym = sgtbx.search_symmetry(
        flags,
        reference_space_group.type(),
        sgtbx.structure_seminvariants( reference_space_group )
        )
    assert reference_ssym.continuous_shifts_are_principal()
    shift_flags = reference_ssym.continuous_shift_flags()
    origins = flex.vec3_double()
    
    for l in left:
        for r in right:
            match = symmetry_operators_between(
                left = l,
                right = r,
                shift_flags = shift_flags,
                crystal = euclidian,
                multiplier = multiplier
                )
            
            if not match:
                continue
            
            ( sgop, pgop, shift ) = match
            transformed = scitbx.matrix.col( sgop * ( 0, 0, 0 ) )
            yield ( transformed + shift ).elems

        
def is_symmetry_equivalent(
    left,
    right,
    crystal,
    multiplier = 1.0,
    origins = no_origin_shift()
    ):
    """
    Determines whether all peaks in left have an equivalent in right
    """
    
    # Quit early if not same composition
    if sorted( l.ensemble for l in left ) != sorted( r.ensemble for r in right ):
        return False
    
    # Do matching
    for ori in origins:
        pairs = overlap_iterator(
            left = left,
            right = [
                r.translated( shift = ori, cell = crystal.cell ) for r in right
                ],
            crystal = crystal,
            multiplier = multiplier
            )
        
        if all( l is not None and r is not None for ( l, r ) in pairs ):
            return True
    
    return False


def overlap(left, right, crystal, multiplier = 1.0):
    """
    Pairs peaks from left with their equivalents in right
    """
    
    pairs = overlap_iterator(
        left = left,
        right = right,
        crystal = crystal,
        multiplier = multiplier
        )
    return [ ( l, r ) for ( l, r ) in pairs if l is not None and r is not None ]


def overlap_iterator(left, right, crystal, multiplier):
    
    copy = set( right )
    
    for l in left:
        matched = None
        
        for r in copy: 
            if symmetry_operators_between(
                left = l,
                right = r,
                shift_flags = ( False, False, False ),
                crystal = crystal,
                multiplier = multiplier
                ):
                matched = r
                break
        
        if matched:
            copy.remove( matched )
            
        yield ( l, matched )
        
    for r in copy:
        yield ( None, r )
        
        
def symmetry_operators_between(left, right, shift_flags, crystal, multiplier):
    """
    Calculates whether left is a symmetry equivalent of right
    (considering crystal symmetry and model point group symmetry)
    """
    
    if left.ensemble == right.ensemble:
        ensemble = left.ensemble
        tolerance = multiplier * MatchTolerance.from_ensemble(
            crystal = crystal,
            ensemble = ensemble
            )
        
        for ( sgop, orth_rotation, delta ) in orthogonal_distances_between(
            left = left.centre_frac,
            right = right.centre_frac,
            crystal = crystal
            ):
            delta_ref = crystal.to_reference_setting * delta
            shift_ref = scitbx.matrix.col(
                [ -d if f else 0 for ( f, d ) in zip( shift_flags, delta_ref.elems ) ]
                )
            delta_ref += shift_ref
            shift = crystal.from_reference_setting * shift_ref
            diff = scitbx.matrix.col(
                crystal.cell.orthogonalize( crystal.from_reference_setting * delta_ref )
                )
            
            if tolerance.translation < diff.length():
                continue
            
            for ( pgop, difference ) in rotation_differences_between(
                left = left.transformation.r,
                right = orth_rotation * right.transformation.r,
                ensemble = ensemble
                ):
                if tolerance.rotation < difference:
                    continue
                
                return ( sgop, pgop, shift )
        
    return None


def orthogonal_distances_between(left, right, crystal):
    """
    Iterator returning orthogonal distance between left and all symmetry
    equivalents of right
    """
    
    for ( sgop, orth_rotation ) in zip( crystal.sgops, crystal.orth_rotations ):
        transformed = scitbx.matrix.col( sgop * right )
        delta = [ fractional_modulus( f = v ) for v in ( transformed - left ).elems ]
        yield ( sgop, orth_rotation, scitbx.matrix.col( delta ) )
        
        
def fractional_modulus(f):
    
    r = f % 1.0
    return r -1.0 if 0.5 < r else r

            
def rotation_differences_between(left, right, ensemble):
    """
    Iterator returning rotation angle differences between left and all
    point group equivalent orientations for right
    """
        
    transpose = left.transpose()
    
    for pgop in ensemble.pgops:
        yield (
            pgop,
            rotation_matrix_angle( rotmat = pgop.r * right * transpose )
            )


def rotation_matrix_angle(rotmat):
    """
    Rotation angle of a rotation matrix
    """
    
    cosdel = ( rotmat[0] + rotmat[4] + rotmat[8] - 1.0 ) / 2.0
    return math.acos( min( max( cosdel, -1.0 ), 1.0 ) )