"""
Agitations contains all MoveGenerator classes that agitate and shake
structures such as distances, angles, etc.

.. inheritance-diagram:: fullrmc.Generators.Agitations
    :parts: 1

+------------------------------------------------------+------------------------------------------------------+
|.. figure:: distanceAgitation.png                     |.. figure:: angleAgitation.png                        |
|   :width: 375px                                      |   :width: 375px                                      |
|   :height: 300px                                     |   :height: 300px                                     |
|   :align: left                                       |   :align: left                                       |
|                                                      |                                                      |
|   Random H-H bond length agitations generated on     |   Random H-O-H angle agitation generated on water    |
|   dihydrogen molecules. At room temperature, H2      |   molecules. At room temperature, water molecule     |
|   molecule bond length fluctuates around 0.74        |   angle formed between the two vectors formed between|
|   Angstroms. Red hydrogen atoms represent the shrank |   consecutively the Oxygen atom and the two hydrogen |
|   H-H bond length molecule while blue hydrogen atoms |   atoms is about 105 deg. Shrank H-O-H angles are    |
|   represent the expanded H-H bond length molecules.  |   represented by the red hydrogen while expanded     |
|                                                      |   angles are represented in blue.                    |
|   (:class:`DistanceAgitationGenerator`)              |   (:class:`AngleAgitationGenerator`)                 |
+------------------------------------------------------+------------------------------------------------------+

.. raw:: html

    <iframe width="560" height="315"
    src="https://www.youtube.com/embed/qTJux9kZCOo?rel=0"
    frameborder="0" allowfullscreen>
    </iframe>

"""
# standard libraries imports
from __future__ import print_function
import re

# external libraries imports
import numpy as np

# fullrmc imports
from ..Globals import INT_TYPE, FLOAT_TYPE, PI, LOGGER
from ..Globals import str, long, unicode, bytes, basestring, range, xrange, maxint
from ..Core.Collection import is_number, is_integer, get_rotation_matrix, generate_random_float, generate_random_vector
from ..Core.MoveGenerator import MoveGenerator, PathGenerator


class DistanceAgitationGenerator(MoveGenerator):
    """
    Generates random agitation moves upon a distance separating two atoms
    by translating both atoms away from each other or closer to each other
    along the direction line between them. This is mainly used to shake two
    atoms bond distance by increasing and decreasing the bond length.
    Only groups of length 2 are accepted.

    :Parameters:
        #. group (None, Group): The group instance. It must contain exactly
           two indexes.
        #. amplitude (number):  The maximum translation amplitude in Angstroms
           applied on every atom.
        #. symmetric (bool): Whether to apply the same amplitude of
           translation on both atoms or not.
        #. shrink (None, bool): Whether to always shrink the distance or
           expand it.
           If True, moves will always bring atoms closer to each other.
           If False, moves will always bring atoms away from each other.
           If None, no orientation is forced, therefore atoms can randomly get
           closer to each other or away from each other.
        #. agitate (tuple): It's a tuple of two boolean values, at least one
           of them must be True.
           Whether to agitate the first atom, the second or both. This is
           useful to set an atom fixed while only
           the other succumb the agitation to adjust the distance. For
           instance in a C-H group it can be useful and
           logical to adjust the bond length by moving only the hydrogen atom
           along the bond direction.

    .. code-block:: python

        # import fullrmc modules
        from fullrmc.Engine import Engine
        from fullrmc.Generators.Agitations import DistanceAgitationGenerator

        # create engine
        ENGINE = Engine(path='my_engine.rmc')

        # set pdb file
        ENGINE.set_pdb('system.pdb')

        # Add constraints ...
        # Re-define groups if needed ...
        # Re-define groups selector if needed ...

        # set moves generators to random agitations of distance seperating two atoms.
        # Maximum agitation amplitude is set to to 0.5A
        for g in ENGINE.groups:
            if len(g)==2:
                g.set_move_generator( DistanceAgitationGenerator(amplitude=0.5) )

    """
    def __init__(self, group=None, amplitude=0.2, symmetric=True, shrink=None, agitate=(True,True)):
        super(DistanceAgitationGenerator, self).__init__(group=group)
        # set amplitude
        self.set_amplitude(amplitude)
        # set symmetric
        self.set_symmetric(symmetric)
        # set shrink
        self.set_shrink(shrink)
        # set agitated
        self.set_agitate(agitate)

    def _codify__(self, name='generator', group=None, addDependencies=True):
        assert isinstance(name, basestring), LOGGER.error("name must be a string")
        assert re.match('[a-zA-Z_][a-zA-Z0-9_]*$', name) is not None, LOGGER.error("given name '%s' can't be used as a variable name"%name)
        dependencies = 'from fullrmc.Generators import Agitations'
        code         = []
        if addDependencies:
            code.append(dependencies)
        code.append("{name} = Agitations.DistanceAgitationGenerator\
(group={group}, amplitude={amplitude}, symmetric={symmetric}, \
shrink={shrink}, agitate={agitate})"
.format(name=name, group=group, amplitude=self.amplitude,
        symmetric=self.symmetric, shrink=self.shrink, agitate=self.agitate))
        # return
        return [dependencies], '\n'.join(code)

    @property
    def amplitude(self):
        """Maximum agitation amplitude."""
        return self.__amplitude

    @property
    def shrink(self):
        """Shrink flag value."""
        return self.__shrink

    @property
    def symmetric(self):
        """Symmetric flag value."""
        return self.__symmetric

    @property
    def agitate(self):
        """Agitate tuple flags value."""
        return self.__agitate

    def check_group(self, group):
        """
        Check the generator's group.

        :Parameters:
            #. group (Group): the Group instance.
        """
        if len(group.indexes)!=2:
            return False, "two atoms are needed in a group to perform distance agitation movements."
        else:
            return True, ""

    def set_amplitude(self, amplitude):
        """
        Sets maximum translation vector allowed amplitude.

        :Parameters:
            #. amplitude (number): the maximum allowed translation
               vector amplitude.
        """
        assert is_number(amplitude), LOGGER.error("Translation amplitude must be a number")
        amplitude = float(amplitude)
        assert amplitude>0, LOGGER.error("Translation amplitude must be bigger than 0")
        self.__amplitude = FLOAT_TYPE(amplitude)

    def set_symmetric(self, symmetric):
        """
        Set symmetric flag value.

        :Parameters:
            #. symmetric (bool): Whether to apply the same amplitude of
               translation on both atoms or not.
        """
        assert isinstance(symmetric, bool), LOGGER.error("symmetric must be boolean")
        self.__symmetric = symmetric

    def set_shrink(self, shrink):
        """
        Set shrink flag value.

        :Parameters:
            #. shrink (None, bool): Whether to always shrink the distance
               or expand it.
               If True, moves will always bring atoms closer to each other.
               If False, moves will always bring atoms away from each other.
               If None, no orientation is forced, therefore distance can
               increase or decrease randomly at every step.
        """
        assert shrink in (None, True, False), LOGGER.error("shrink can only be None, True or False")
        self.__shrink = shrink

    def set_agitate(self, agitate):
        """
        Set agitate tuple value.

        :Parameters:
            #. agitate (tuple): It's a tuple of two boolean values, at least
               one of them must be True. Whether to agitate the first atom,
               the second or both. This is useful to set an atom fixed while
               only the other succumb the agitation to adjust the distance.
               For instance in a C-H group it can be useful and logical to
               adjust the bond length by moving only the hydrogen atom along
               the bond direction.
        """
        assert isinstance(agitate, (list,tuple)), LOGGER.error("agitate must be a list or a tuple")
        assert len(agitate)==2, LOGGER.error("agitate must have 2 items")
        assert [isinstance(a,bool) for a in agitate]==[True,True], LOGGER.error("agitate items must be boolean")
        assert agitate[0] or agitate[1], LOGGER.error("agitate both items can't be False")
        self.__agitate = (agitate[0], agitate[1])

    def transform_coordinates(self, coordinates, argument=None):
        """
        Translate coordinates.

        :Parameters:
            #. coordinates (np.ndarray): The coordinates on which to apply
               the translation.

        :Returns:
            #. coordinates (np.ndarray): The new coordinates after applying
               the translation.
            #. argument (object): Any python object. Not used in this
               generator.
        """
        if coordinates.shape[0]!=2:
            # atoms where removed, fall back to random translation
            return coordinates+generate_random_vector(minAmp=self.__amplitude[0],
                                                      maxAmp=self.__amplitude[1])
        else:
            # get normalized direction vector
            vector  = FLOAT_TYPE( coordinates[0,:]-coordinates[1,:] )
            vector /= FLOAT_TYPE( np.linalg.norm(vector) )
            # create amplitudes
            if self.__symmetric:
                amp0 = amp1 =  FLOAT_TYPE(generate_random_float()*self.__amplitude)
            else:
                amp0 =  FLOAT_TYPE(generate_random_float()*self.__amplitude)
                amp1 =  FLOAT_TYPE(generate_random_float()*self.__amplitude)
            # create shrink flag
            if self.__shrink is None:
                shrink = (1-2*generate_random_float())>0
            else:
                shrink = self.__shrink
            # create directions
            if shrink:
                dir0 = FLOAT_TYPE(-1)
                dir1 = FLOAT_TYPE( 1)
            else:
                dir0 = FLOAT_TYPE( 1)
                dir1 = FLOAT_TYPE(-1)
            # create translation vectors
            translationVectors      = np.empty((2,3), dtype=FLOAT_TYPE)
            translationVectors[0,:] = self.__agitate[0]*dir0*amp0*vector
            translationVectors[1,:] = self.__agitate[1]*dir1*amp1*vector
            # translate and return
            return coordinates+translationVectors


class AngleAgitationGenerator(MoveGenerator):
    """
    Generate random agitation moves upon an angle defined between two vectors
    left-central and right-central where (central, left, right) are three
    atoms. Move will be performed on left and/or right atom while central
    atom will always remain fixed. Distances between left/right and central
    atoms will remain unchanged. This is mainly used to shake bonded atoms
    angles by increasing and decreasing the bond length.
    Only groups of length 3 are accepted.

    :Parameters:
        #. group (None, Group): The group instance. It must contain exactly
           three indexes in respective order (central, left, right)
           atoms index.
        #. amplitude (number):  The maximum agitation angle amplitude in
           degrees of left and right atoms separately.
        #. symmetric (bool): Whether to apply the same amplitude of rotation
           on both left and right atoms or not.
        #. shrink (None, bool): Whether to always shrink the angle or expand it.
           If True, moves will always reduce angle.
           If False, moves will always increase angle.
           If None, no orientation is forced, therefore angle can randomly
           get wider or tighter.
        #. agitate (tuple): It's a tuple of two boolean values for
           respectively (left, right) atoms,
           at least one of them must be True. Whether to agitate the left
           atom, the right or both. This is useful to set an atom fixed while
           only the other succumb the agitation to adjust the angle.

    .. code-block:: python

        # import fullrmc modules
        from fullrmc.Engine import Engine
        from fullrmc.Generators.Agitations import AngleAgitationGenerator

        # create engine
        ENGINE = Engine(path='my_engine.rmc')

        # set pdb file
        ENGINE.set_pdb('system.pdb')

        # Add constraints ...
        # Re-define groups if needed ...
        # Re-define groups selector if needed ...

        # set moves generators to random agitations of the angle formed between
        # one central atom and other two. Maximum agitation amplitude is set to to 10.
        for g in ENGINE.groups:
            if len(g)==3:
                g.set_move_generator( AngleAgitationGenerator(amplitude=10) )

    """
    def __init__(self, group=None, amplitude=2, symmetric=True, shrink=None, agitate=(True,True)):
        super(AngleAgitationGenerator, self).__init__(group=group)
        # set amplitude
        self.set_amplitude(amplitude)
        # set symmetric
        self.set_symmetric(symmetric)
        # set shrink
        self.set_shrink(shrink)
        # set agitated
        self.set_agitate(agitate)

    def _codify__(self, name='generator', group=None, addDependencies=True):
        assert isinstance(name, basestring), LOGGER.error("name must be a string")
        assert re.match('[a-zA-Z_][a-zA-Z0-9_]*$', name) is not None, LOGGER.error("given name '%s' can't be used as a variable name"%name)
        dependencies = 'from fullrmc.Generators import Agitations'
        code         = []
        if addDependencies:
            code.append(dependencies)
        code.append("{name} = Agitations.AngleAgitationGenerator\
(group={group}, amplitude={amplitude}, symmetric={symmetric}, \
shrink={shrink}, agitate={agitate})"
.format(name=name, group=group, amplitude=self.amplitude*FLOAT_TYPE(180.)/PI,
        symmetric=self.symmetric, shrink=self.shrink, agitate=self.agitate))
        # return
        return [dependencies], '\n'.join(code)

    @property
    def amplitude(self):
        """Maximum agitation angle amplitude in rad."""
        return self.__amplitude

    @property
    def shrink(self):
        """Shrink flag value."""
        return self.__shrink

    @property
    def symmetric(self):
        """Symmetric flag value."""
        return self.__symmetric

    @property
    def agitate(self):
        """Agitate tuple flags value."""
        return self.__agitate

    def check_group(self, group):
        """
        Check the generator's group.

        :Parameters:
            #. group (Group): the Group instance.
        """
        if len(group.indexes)!=3:
            return False, "Exactly three atoms are needed in a group to perform angle agitation movements."
        else:
            return True, ""

    def set_amplitude(self, amplitude):
        """
        Set maximum allowed agitation rotation angle amplitude in degrees of
        left and right atoms separately and transforms it to rad.

        :Parameters:
            #. amplitude (number):  The maximum agitation angle amplitude in
               degrees of left and right atoms separately.
        """
        assert is_number(amplitude), LOGGER.error("Agitation angle amplitude must be a number")
        amplitude = float(amplitude)
        assert amplitude>0, LOGGER.error("Agitation angle amplitude must be bigger than 0")
        assert amplitude<=90, LOGGER.error("Agitation angle amplitude must be smaller than 90")
        self.__amplitude = FLOAT_TYPE(amplitude*PI/180.)

    def set_symmetric(self, symmetric):
        """
        Set symmetric flag value.

        :Parameters:
            #. symmetric (bool): Whether to apply the same amplitude of
               translation on both atoms or not.
        """
        assert isinstance(symmetric, bool), LOGGER.error("symmetric must be boolean")
        self.__symmetric = symmetric

    def set_shrink(self, shrink):
        """
        Set shrink flag value.

        :Parameters:
            #. shrink (None, bool): Whether to always shrink the distance or
               expand it.
               If True, moves will always bring atoms closer to each other.
               If False, moves will always bring atoms away from each other.
               If None, no orientation is forced, therefore distance can
               increase or decrease randomly at every step.
        """
        assert shrink in (None, True, False), LOGGER.error("shrink can only be None, True or False")
        self.__shrink = shrink

    def set_agitate(self, agitate):
        """
        Set agitate tuple value.

        :Parameters:
            #. agitate (tuple): It's a tuple of two boolean values, at least
               one of them must be True. Whether to agitate the first atom,
               the second or both. This is useful to set an atom fixed while
               only the other succumb the agitation to adjust the distance.
               For instance in a C-H group it can be useful and logical to
               adjust the bond length by moving only the hydrogen atom along
               the bond direction.
        """
        assert isinstance(agitate, (list,tuple)), LOGGER.error("agitate must be a list or a tuple")
        assert len(agitate)==2, LOGGER.error("agitate must have 2 items")
        assert [isinstance(a,bool) for a in agitate]==[True,True], LOGGER.error("agitate items must be boolean")
        assert agitate[0] or agitate[1], LOGGER.error("agitate both items can't be False")
        self.__agitate = (agitate[0], agitate[1])

    def transform_coordinates(self, coordinates, argument=None):
        """
        Translate coordinates.

        :Parameters:
            #. coordinates (np.ndarray): The coordinates on which to apply
               the translation.

        :Returns:
            #. coordinates (np.ndarray): The new coordinates after applying
               the translation.
            #. argument (object): Any python object. Not used in this generator.
        """
        if coordinates.shape[0]!=3:
            # atoms where removed, fall back to random translation
            return coordinates+generate_random_vector(minAmp=self.__amplitude[0],
                                                      maxAmp=self.__amplitude[1])
        else:
            # get atoms group center
            center = np.sum(coordinates, 0)/coordinates.shape[0]
            # translate to origin
            rotatedCoordinates = coordinates-center
            # get normalized direction vectors
            leftVector   = FLOAT_TYPE( rotatedCoordinates[1,:]-rotatedCoordinates[0,:] )
            leftVector  /= FLOAT_TYPE( np.linalg.norm(leftVector) )
            rightVector  = FLOAT_TYPE( rotatedCoordinates[2,:]-rotatedCoordinates[0,:] )
            rightVector /= FLOAT_TYPE( np.linalg.norm(rightVector) )
            # get rotation axis
            rotationAxis = np.cross(leftVector, rightVector)
            if rotationAxis[0]==rotationAxis[1]==rotationAxis[2]==0.:
                rotationAxis = np.array(1-2*np.random.random(3), dtype=FLOAT_TYPE)
                rotationAxis /= FLOAT_TYPE( np.linalg.norm(rotationAxis) )
            # create shrink flag
            if self.__shrink is None:
                shrink = (1-2*generate_random_float())>0
            else:
                shrink = self.__shrink
            # get rotation angles
            if self.__symmetric:
                angleLeft  = angleRight = FLOAT_TYPE(generate_random_float()*self.__amplitude)
            else:
                angleLeft  = FLOAT_TYPE(generate_random_float()*self.__amplitude)
                angleRight = FLOAT_TYPE(generate_random_float()*self.__amplitude)
            # create directions
            if shrink:
                angleLeft  *= FLOAT_TYPE(-1)
                angleRight *= FLOAT_TYPE( 1)
            else:
                angleLeft  *= FLOAT_TYPE( 1)
                angleRight *= FLOAT_TYPE(-1)
            # rotate
            if self.__agitate[0]:
                rotationMatrix = get_rotation_matrix(rotationAxis, angleLeft)
                rotatedCoordinates[1,:] = np.dot( rotationMatrix, rotatedCoordinates[1,:])
            if self.__agitate[1]:
                rotationMatrix = get_rotation_matrix(rotationAxis, angleRight)
                rotatedCoordinates[2,:] = np.dot( rotationMatrix, rotatedCoordinates[2,:])
            # translate back from center and return
            return np.array(rotatedCoordinates+center, dtype=FLOAT_TYPE)
