import scitbx.matrix
from scitbx.math import r3_rotation_cos_rotation_angle_from_matrix
from cctbx import sgtbx

import math
import operator

# Crystal
class SymOp(object):
    """
    Data from a symmetry operator
    """

    def __init__(self, sgop, crystal):

        self.sgop = sgop
        self.frac_rotation = scitbx.matrix.sqr( sgop.r().as_double() )
        self.orth_rotation = crystal.to_orth * self.frac_rotation * crystal.to_frac
        
        
class ChangeOfBasisOperator(object):
    """
    Object to store all data for a coordinate transformation
    """
    
    def __init__(self, operator, inverted, original, transformed):
        
        self.operator = operator
        self.inverted = inverted
        self.original = original
        self.transformed = transformed
        
        
    @property
    def operator_r(self):
        
        return scitbx.matrix.sqr( self.operator.c().r().as_double() ) 
    
    
    @property
    def inverted_r(self):
        
        return scitbx.matrix.sqr( self.operator.c_inv().r().as_double() )
        
    
    def transform_position(self, vector):
        
        return self.operator( vector )
    
    
    def transform_orientation(self, matrix):
        
        return self.transformed.to_orth * self.operator_r * self.original.to_frac * matrix
    
    
    def transform_r_frac(self, matrix):
        
        return self.operator_r * matrix * self.inverted_r
    
    
    def transform_r_orth(self, matrix):
        
        return ( self.transformed.to_orth
            * self.transform_r_frac(
                matrix = self.original.to_frac * matrix * self.original.to_orth,
                )
            * self.transformed.to_frac
            )
    
    
    def transform_rt_mx(self, rt_mx):
        
        return self.operator.apply( rt_mx )
    
    
    def inverse(self):
        
        return self.__class__(
            operator = self.inverted,
            inverted = self.operator,
            original = self.transformed,
            transformed = self.original,
            )
    
    
    @classmethod
    def new(cls, cbop, crystal):
        
        return cls(
            operator = cbop,
            inverted = cbop.inverse(),
            original = crystal,
            transformed = crystal.with_changed_basis( cbop = cbop ),
            )


class IdentityChangeOfBasisOperator(object):
    """
    An optimization to not do any transformation when not necessary
    """
    
    UNIT_R = scitbx.matrix.identity( 3 )
    UNIT_CBOP = sgtbx.change_of_basis_op()
    
    def __init__(self, crystal):
        
        self.original = crystal
        
    
    @property
    def transformed(self):
        
        return self.original
    
    
    @property
    def operator(self):
        
        return self.UNIT_CBOP
    
    
    @property
    def inverted(self):
        
        return self.UNIT_CBOP
    
        
    @property
    def operator_r(self):
        
        return self.UNIT_R
    
    
    @property
    def inverted_r(self):
        
        return self.UNIT_R
        
    
    def transform_position(self, vector):
        
        return vector
    
    
    def transform_orientation(self, matrix):
        
        return matrix
    
    
    def transform_r_frac(self, matrix):
        
        return matrix
    
    
    def transform_r_orth(self, matrix):
        
        return matrix
    
    
    def transform_rt_mx(self, rt_mx):
        
        return rt_mx
    
    
    def inverse(self):
        
        return self


class CrystalInfo(object):
    """
    Symmetry information
    """

    def __init__(self, space_group, cell):

        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.symops = [
            SymOp( sgop = op, crystal = self ) for op in space_group
            ]
        self.origin_search = None
        self.reference_setting_info = None
        self.cbop_ref = None
        self.is_ref_setting = None


    def space_group_info(self):

        return self.space_group.info()


    def space_group_type(self):

        return self.space_group.type()


    def change_of_basis_to_reference(self):

        if self.cbop_ref is None:
            if self.is_reference_setting():
                self.cbop_ref = IdentityChangeOfBasisOperator( crystal = self )
                
            else:
                self.cbop_ref = ChangeOfBasisOperator.new(
                    cbop = self.space_group_type().cb_op(),
                    crystal = self,
                    )

        return self.cbop_ref


    def with_changed_basis(self, cbop):

        return self.__class__(
            space_group = self.space_group.change_basis( cbop ),
            cell = self.cell.change_basis( cbop ),
            )

    def origin_search_info(self):

        if self.origin_search is None:
            self.origin_search = OriginSearchInfo( crystal = self )

        return self.origin_search


    def reference_setting(self):

        if self.reference_setting_info is None:
            self.reference_setting_info = self.change_of_basis_to_reference().transformed

        return self.reference_setting_info


    def is_reference_setting(self):

        if self.is_ref_setting is None:
            self.is_ref_setting = self.space_group_info().is_reference_setting()

        return self.is_ref_setting


class OriginSearchInfo(object):
    """
    Data for origin search
    """

    def __init__(self, crystal):

        refset = crystal.reference_setting()

        self.sgi = refset.space_group_info()
        search_symmetry = sgtbx.search_symmetry(
            flags=sgtbx.search_symmetry_flags(
                use_space_group_symmetry = False,
                use_space_group_ltr = -1,
                use_seminvariants = True,
                use_normalizer_k2l = False,
                use_normalizer_l2n = False
                ),
            space_group_type = self.sgi.type(),
            seminvariant = self.sgi.structure_seminvariants()
            )
        self.ops = [
            SymOp( sgop = op, crystal = crystal ) for op in search_symmetry.subgroup()
            ]
        assert search_symmetry.continuous_shifts_are_principal()
        self.continuous_shift_flags = search_symmetry.continuous_shift_flags()

        if any( self.continuous_shift_flags ):
            self.projector = subtract_continuous_allowed_origin_shifts

        else:
            self.projector = unchanged
            
    
    def project(self, delta):
      
      return self.projector( sgi = self.sgi, delta = delta )
            
            
def subtract_continuous_allowed_origin_shifts(sgi, delta):

    return sgi.subtract_continuous_allowed_origin_shifts( delta )
  
  
def unchanged(sgi, delta):

    return delta
  
  
# Backwards compatibility
project_without_shifts = subtract_continuous_allowed_origin_shifts
project_with_shifts = unchanged


# Object description
class EnsembleInfo(object):
    """
    Information about ensemble extent, centre, symmetry
    """

    def __init__(self, centre, extent, pgops):

        self.centre = centre
        self.extent = extent
        self.pgops = pgops


    @classmethod
    def from_hierarchy(cls, root):

        from scitbx.array_family import flex
        principal = PrincipalComponents.from_hierarchy( root = root )
        extent = sum( principal.extent ) / 3.0

        from phaser import ncs
        ensym = ncs.EnsembleSymmetry()
        ensym.set_root( root = root )
        pgops = [ scitbx.matrix.rt( ( op.get_rotation(), op.get_translation() ) )
            for op in ensym.get_point_group().get_operators() ]

        return cls(
            centre = principal.centre,
            extent = extent,
            pgops = pgops,
            )


    @classmethod
    def from_map(cls, centre, extent):

        return cls(
            centre = centre,
            extent = sum( extent ) / 3.0,
            pgops = [
                scitbx.matrix.rt(
                    ( ( 1, 0, 0, 0, 1, 0, 0, 0, 1 ), ( 0, 0, 0 ) )
                    ),
                ]
            )


class PrincipalComponents(object):
    """
    Principal components of inertia
    """

    def __init__(self, coords, weights):

        from scitbx.array_family import flex

        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

        ( txs, tys, tzs ) = ( self.r * coords + self.t ).parts()
        self.extent = (
            ( max( txs ) - min( txs ) ) / 2.0,
            ( max( tys ) - min( tys ) ) / 2.0,
            ( max( tzs ) - min( tzs ) ) / 2.0,
            )


    @classmethod
    def from_hierarchy(cls, root):

        import cctbx.eltbx.tiny_pse
        from scitbx.array_family import flex

        atoms = root.models()[0].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 ]
            )
        return cls( coords = coords, weights = weights )



# Match quantities - basic types
class simple_position_differences(object):
    """
    Simple (i.e. not cached) calculation of positional differences
    """

    def __init__(self, centre, crystal):

        self.centre = centre
        self.symops = crystal.symops


    def __call__(self, other):

        for symop in self.symops:
            transformed = scitbx.matrix.col( symop.sgop * other.centre )
            yield ( symop, mod_short_coords( coords = self.centre - transformed ) )


class cached_position_differences(object):
    """
    Symmetry equivalent positions are precalculated
    """

    def __init__(self, centre, crystal):

        self.symmeqs = [
            scitbx.matrix.col( symop.sgop.inverse() * centre ) for symop in crystal.symops
            ]
        self.symops = crystal.symops


    def __call__(self, other):

        for ( symop, posi ) in zip( self.symops, self.symmeqs ):
            yield ( symop, mod_short_coords( coords = posi - other.centre ) )


class Position(object):
    """
    Information about the position of an object
    'centre' is stored as fractional, but 'tolerance_sq' is in orthogonal
    coordinates
    """

    def __init__(self, centre, crystal, tolerance_sq, detail = simple_position_differences):

        self.centre = centre
        self.crystal = crystal
        self.fractional_differences_between = detail(
            centre = self.centre,
            crystal = self.crystal,
            )
        self.tolerance_sq = tolerance_sq
        
        
    def orthogonal_differences_between(self, other):
      
        for ( symop, delta ) in self.fractional_differences_between( other = other ):
            yield ( symop, self.crystal.cell.orthogonalize( delta ) )


    def operations_between(self, other):

        for ( symop, delta ) in self.orthogonal_differences_between( other = other ):
            if scitbx.matrix.col( delta ).length_sq() <= self.tolerance_sq:
                yield symop


    def transformed(self, op, detail = simple_position_differences):

        return self.__class__(
            centre = scitbx.matrix.col( op * self.centre ),
            crystal = self.crystal,
            tolerance_sq = self.tolerance_sq,
            detail = detail,
            )


    def shifted(self, vector, detail = simple_position_differences):

        return self.__class__(
            centre = self.centre + vector,
            crystal = self.crystal,
            tolerance_sq = self.tolerance_sq,
            detail = detail,
            )


    def reference_setting(self, detail = simple_position_differences):

        if self.crystal.is_reference_setting():
            return self

        else:
            return self.with_changed_basis(
                cbop = self.crystal.change_of_basis_to_reference(),
                detail = detail, 
                )
    
    
    def with_changed_basis(self, cbop, detail = simple_position_differences):

        return self.__class__(
            centre = scitbx.matrix.col(
                cbop.transform_position( vector = self.centre )
                ),
            crystal = cbop.transformed,
            tolerance_sq = self.tolerance_sq,
            detail = detail,
            )


    def non_spatial_equivalence(self, other):

        return True


    def non_positional_equivalence(self, other, symop, euclop, cbop):

        return True


    def positional(self):

        return self


    @classmethod
    def new(cls, centre, crystal, tolerance, detail = simple_position_differences):

        return cls(
            centre = centre,
            crystal = crystal,
            tolerance_sq = tolerance ** 2,
            detail = detail,
            )


def mod_short(number):

    frac = number % 1.0 
    return frac - 1.0 if 0.5 < frac else frac
    
    
def mod_short_coords(coords):

    return [ mod_short( number = r ) for r in coords ]


class basic_rotation_differences(object):
    """
    Basic way of calculating rotation differences
    """

    def __init__(self, rotation, pgops):

        self.rotation = rotation
        self.pgops = pgops


    def __call__(self, other):

        transposed = self.rotation.transpose()
        
        for pgop in self.pgops:
            rotmat = other.rotation * pgop.r * transposed
            yield ( pgop, rotmat )
            

class simple_rotation_differences(object):
    """
    Simple (i.e. not cached) way of calculating rotation differences
    """

    def __init__(self, rotation, pgops):

        self.transposed = rotation.transpose()
        self.pgops = pgops


    def __call__(self, other):

        for pgop in self.pgops:
            rotmat = other.rotation * pgop.r * self.transposed
            yield ( pgop, rotmat )


class cached_rotation_differences(object):
    """
    Precalculate matrices
    """

    def __init__(self, rotation, pgops):

        transposed = rotation.transpose()
        self.data = [ ( pgop, pgop.r * transposed ) for pgop in pgops ]


    def __call__(self, other):

        for ( pgop, rotmat ) in self.data:
            yield ( pgop, other.rotation * rotmat )


class Orientation(object):
    """
    Information about the orientation of an object
    'rotation' and 'pgops' is in orthogonal coordinates
    'tolerance' is in radians
    """

    def __init__(self, rotation, pgops, cos_tolerance, detail = simple_rotation_differences):

        self.rotation = rotation
        self.pgops = pgops
        self.cos_tolerance = cos_tolerance
        self.rotation_differences_between = detail( rotation = self.rotation, pgops = self.pgops )


    def operations_between(self, other):

        for ( pgop, rotmat ) in self.rotation_differences_between( other = other ):
            if self.cos_tolerance <= r3_rotation_cos_rotation_angle_from_matrix( rotmat ):
                yield pgop


    def rotated(self, op, detail = basic_rotation_differences):

        return self.__class__(
            rotation = op * self.rotation,
            pgops = self.pgops,
            cos_tolerance = self.cos_tolerance,
            detail = detail,
            )


    def non_spatial_equivalence(self, other):

        return True


    def non_positional_equivalence(self, other, symop, euclop, cbop):

        joint = cbop.inverse().transform_r_orth(
            matrix = symop.orth_rotation * euclop.orth_rotation,
            )
        return any(
            self.operations_between( other = other.rotated( op = joint ) )
            )


    def with_changed_basis(self, cbop, detail = basic_rotation_differences):

        return self.__class__(
            rotation = cbop.transform_orientation( matrix = self.rotation ),
            pgops = self.pgops,
            cos_tolerance = self.cos_tolerance,
            detail = detail,
            )


    @classmethod
    def new(cls, rotation, pgops, tolerance, detail = simple_rotation_differences):

        return cls(
            rotation = rotation,
            pgops = pgops,
            cos_tolerance = math.cos( tolerance ),
            detail = detail,
            )


# Match quantities - composite types
def phaser_style_orientation_tolerance(extent, dmin):

    return 2.0 * math.atan( dmin / ( 4.0 * extent ) )


def phaser_style_position_tolerance(dmin):

    return dmin / 5.0


class MRPeakPO(object):
    """
    An object that has a combination of orientation and position
    Matching starts with position, and then with orientation
    """

    def __init__(self, orientation, position):

        self.orientation = orientation
        self.position = position


    def operations_between(self, other):

        for symop in self.position.operations_between( other = other.position ):
            rotated = other.orientation.rotated( op = symop.orth_rotation )

            for pgop in self.orientation.operations_between( other = rotated ):
                yield ( symop, pgop )


    def non_spatial_equivalence(self, other):

        return True


    def non_positional_equivalence(self, other, symop, euclop, cbop):

        return self.orientation.non_positional_equivalence(
            other = other.orientation,
            symop = symop,
            euclop = euclop,
            cbop = cbop,
            )


    def shifted(self, vector):

        return self.__class__(
            orientation = self.orientation,
            position = self.position.shifted( vector = vector ),
            )


    def reference_setting(self):

        if self.position.crystal.is_reference_setting():
            return self

        else:
            return self.__class__(
                orientation = self.orientation.with_changed_basis(
                    cbop = self.crystal.change_of_basis_to_reference()
                    ),
                position = self.position.reference_setting(),
                )


    def positional(self):

        return self.position


    @property
    def crystal(self):

        return self.position.crystal


    @classmethod
    def phaser_style(
        cls,
        rotation,
        translation,
        crystal,
        ensemble,
        dmin,
        multiplier = 1.0,
        detail_position = simple_position_differences,
        detail_orientation = simple_rotation_differences,
        ):

        orth_rotmat = scitbx.matrix.sqr( rotation )

        return cls(
            orientation = Orientation.new(
                rotation = orth_rotmat,
                pgops = ensemble.pgops,
                tolerance = multiplier * phaser_style_orientation_tolerance(
                    extent = ensemble.extent,
                    dmin = dmin
                    ),
                detail = detail_orientation
                ),
            position = Position.new(
                centre = (
                    scitbx.matrix.col(
                        crystal.cell.fractionalize( orth_rotmat * ensemble.centre )
                        )
                    + scitbx.matrix.col( translation )
                    ),
                crystal = crystal,
                tolerance = multiplier * phaser_style_position_tolerance( dmin = dmin ),
                detail = detail_position,
                )
            )


class Label(object):
    """
    Information about some property of an object
    """

    def __init__(self, value, equality):

        self.value = value
        self.equality = equality


    def equals(self, other):

        return self.equality( self.value, other.value )


class Labelled(object):
    """
    An object that has a structural and a label component
    Matching starts with label comparison
    """

    def __init__(self, label, spatial):

        self.label = label
        self.spatial = spatial


    def operations_between(self, other):

        if self.non_spatial_equivalence( other = other ):
            for op in self.spatial.operations_between( other = other.spatial ):
                yield op


    def non_spatial_equivalence(self, other):

        return self.label.equals( other.label )


    def non_positional_equivalence(self, other, symop, euclop, cbop):

        return self.spatial.non_positional_equivalence(
            other = other.spatial,
            symop = symop,
            euclop = euclop,
            cbop = cbop,
            )


    def shifted(self, vector):

        return self.__class__(
            label = self.label,
            spatial = self.spatial.shifted( vector = vector ),
            )


    def reference_setting(self):

        return self.__class__(
            label = self.label,
            spatial = self.spatial.reference_setting(),
            )


    def positional(self):

        return self.spatial.positional()


    @property
    def crystal(self):

        return self.spatial.crystal


    @classmethod
    def phaser_mr_solution(
        cls,
        rotation,
        translation,
        crystal,
        ensemble,
        dmin,
        multiplier = 1.0,
        detail_position = simple_position_differences,
        detail_orientation = simple_rotation_differences,
        ):

        return cls(
            label = Label( value = ensemble, equality = operator.eq ),
            spatial = MRPeakPO.phaser_style(
                rotation = rotation,
                translation = translation,
                crystal = crystal,
                ensemble = ensemble,
                dmin = dmin,
                multiplier = multiplier,
                detail_position = detail_position,
                detail_orientation = detail_orientation,
                )
            )


    @classmethod
    def simple_heavy_atom(cls, element, position, crystal, tolerance):

        return cls(
            label = Label( value = element, equality = operator.eq ),
            spatial = Position.new(
                centre = scitbx.matrix.col( position ),
                crystal = crystal,
                tolerance = tolerance
                )
            )


class simple_origin_shifts(object):
    """
    Simple (i.e. non-cached) origin shift calculation
    """

    def __init__(self, position, search_info):

        self.position = position
        self.search_info = search_info


    def __call__(self, other):

        for euclop in self.search_info.ops:
            transformed = other.position.transformed( op = euclop.sgop )

            for ( symop, delta ) in self.position.fractional_differences_between( other = transformed ):
                projected = self.search_info.project( delta = delta )
                yield ( euclop, symop, projected, delta )


class cached_origin_shifts(object):
    """
    Possible origin shifts are precalculated
    """

    def __init__(self, position, search_info):

        self.euclids = [
            position.transformed(
                op = euclop.sgop.inverse(),
                detail = position.fractional_differences_between.__class__
                )
            for euclop in search_info.ops
            ]
        self.search_info = search_info


    def __call__(self, other):

        for ( euclop, posi ) in zip( self.search_info.ops, self.euclids ):
            for ( symop, delta ) in posi.fractional_differences_between( other = other.position ):
                projected = self.search_info.project( delta = delta )
                yield ( euclop, symop, projected, delta )


class OriginSearchEntity(object):
    """
    Information about an object that allows to do origin search
    """

    def __init__(
        self,
        entity,
        multiplier = 1.0,
        detail_origin = simple_origin_shifts,
        detail_position = simple_position_differences,
        ):

        self.entity = entity
        self.cbop = self.entity.crystal.change_of_basis_to_reference()
        self.position = self.entity.positional().with_changed_basis(
            cbop = self.cbop,
            detail = detail_position,
            )
        self.multiplier = multiplier
        self.tolerance_sq = multiplier * self.position.tolerance_sq
        self.projected_differences_between = detail_origin(
            position = self.position,
            search_info = self.cbop.transformed.origin_search_info()
            )


    def potential_origins_between(self, other):

        assert self.cbop == other.cbop
        
        if self.entity.non_spatial_equivalence( other = other.entity ):
            for ( euclop, symop, delta, shift ) in self.projected_differences_between( other = other ):
                diff = sum(
                  [ c ** 2 for c in self.entity.crystal.cell.orthogonalize( delta ) ]
                  )
                
                if diff <= self.tolerance_sq:
                    if self.entity.non_positional_equivalence(
                        other = other.entity,
                        symop = symop,
                        euclop = euclop,
                        cbop = self.cbop,
                        ):
                        transf = self.cbop.inverse().transform_position( vector = shift )
                        yield (
                            euclop,
                            symop,
                            scitbx.matrix.col( transf ),
                            )


# Origin determination and refinement
def no_origin_refinement(matched, euclop, symop, shift, left, right):

    return ( matched, euclop, symop, shift, left, right )


def origin_shifts_between(left, right, refinement = no_origin_refinement):

    for l in left:
        for r in right:
            for ( euclop, symop, shift ) in l.potential_origins_between( other = r ):
                yield refinement(
                    matched = ( l, r ),
                    euclop = euclop,
                    symop = symop,
                    shift = shift,
                    left = [ le for le in left if le is not l ],
                    right = [ ri for ri in right if ri is not r ],
                    )


def origin_shifted_overlap(origin):

    ( matched, euclop, symop, shift, left, right ) = origin
    yield matched

    left_entities = [ l.entity for l in left ]
    shifted = [ r.entity.shifted( vector = shift ) for r in right ]

    for ( l, r ) in overlap( left = left_entities, right = shifted ):
        if l is not None:
            assert l in left_entities
            l = left[ left_entities.index( l ) ]

        if r is not None:
            assert r in shifted
            r = right[ shifted.index( r ) ]

        yield ( l, r )


# Matching methods
def overlap(left, right):

    copy = set( right )

    for reference in left:
        for r in copy:
            ops = reference.operations_between( other = r )

            try:
                ops.next()

            except StopIteration:
                continue

            else:
                copy.remove( r )
                yield ( reference, r )
                break

        else:
            yield ( reference, None )

    for r in copy:
        yield ( None, r )


def pairs(overlaps):

    return (
        ( l, r ) for ( l, r ) in overlaps if l is not None and r is not None
        )


def is_equivalent(left, right):

    if len( left ) != len( right ):
        return False

    if any( l is None or r is None for ( l, r ) in overlap( left, right) ):
        return False

    return True


def phaser_style_equivalence_match(left, right):

    if len( left ) != len( right ):
        return False

    if not left:
        return True

    for other in right:
        for ( euclop, symop, shift ) in left[0].potential_origins_between( other = other ):
            full_shift = scitbx.matrix.col( euclop.sgop * shift )
            shifted = [ r.entity.shifted( vector = full_shift ) for r in right if r != other ]

            if is_equivalent( left = [ l.entity for l in left[1:] ], right = shifted ):
                return True

    return False