"""A reference implementation of the multi-processor package demonstrating pre-sending of data.

Description
===========

This implementation of the multi-processor package is used to demonstrate how to pre-send invariant data once to all slaves to minimise the data transfer for each parallelised calculation block.


Testing
=======

To run in uni-processor mode on a dual core system, change the MULTI variable to False and type::

$ python test_implementation2.py


To run in mpi4py multi-processor mode with one master and two slave processors, change the MULTI variable to True and type::

$ mpiexec -n 3 python test_implementation2.py
"""

# Python module imports.
from numpy import array, dot, float64, zeros
from numpy.linalg import norm
try:
    import cProfile as profile
except ImportError:
    import profile
import pstats
from random import gauss
import sys

# Modify the module path.
sys.path.append('..')

# relax module imports.
from multi import fetch_data, fetch_data_store, send_data_to_slaves, Application_callback, load_multiprocessor, Memo, Processor_box, Result_command, Slave_command


# Module variables.
PROFILE = False
MULTI = True
if MULTI:
    FABRIC = 'mpi4py'
    PROCESSOR_NUM = 2
else:
    FABRIC = 'uni'
    PROCESSOR_NUM = 1


def print_stats(stats, status=0):
    """Profiling printout function, sorting first by cumulative time."""

    # Sorted printout.
    pstats.Stats(stats).sort_stats('cumulative').print_stats()


def quaternion_to_R(quat, R):
    """Convert a quaternion into rotation matrix form.

    This is from Wikipedia (http://en.wikipedia.org/wiki/Rotation_matrix#Quaternion), where::

            | 1 - 2y**2 - 2z**2      2xy - 2zw          2xz + 2yw     |
        Q = |     2xy + 2zw      1 - 2x**2 - 2z**2      2yz - 2xw     |,
            |     2xz - 2yw          2yz + 2xw      1 - 2x**2 - 2y**2 |

    and where the quaternion is defined as q = (w, x, y, z).  This has been verified using Simo
    Saerkkae's "Notes on Quaternions" at http://www.lce.hut.fi/~ssarkka/.


    @param quat:    The quaternion.
    @type quat:     numpy 4D, rank-1 array
    @param R:       A 3D matrix to convert to the rotation matrix.
    @type R:        numpy 3D, rank-2 array
    """

    # Alias.
    (w, x, y, z) = quat

    # Repetitive calculations.
    x2 = 2.0 * x**2
    y2 = 2.0 * y**2
    z2 = 2.0 * z**2
    xw = 2.0 * x*w
    xy = 2.0 * x*y
    xz = 2.0 * x*z
    yw = 2.0 * y*w
    yz = 2.0 * y*z
    zw = 2.0 * z*w

    # The diagonal.
    R[0, 0] = 1.0 - y2 - z2
    R[1, 1] = 1.0 - x2 - z2
    R[2, 2] = 1.0 - x2 - y2

    # The off-diagonal.
    R[0, 1] = xy - zw
    R[0, 2] = xz + yw
    R[1, 2] = yz - xw

    R[1, 0] = xy + zw
    R[2, 0] = xz - yw
    R[2, 1] = yz + xw


def R_random_hypersphere(R):
    """Generate a random rotation matrix using 4D hypersphere point picking.

    A quaternion is generated by creating a 4D vector with each value randomly selected from a
    Gaussian distribution, and then normalising.

    @param R:       A 3D matrix to convert to the rotation matrix.
    @type R:        numpy 3D, rank-2 array
    """

    # The quaternion.
    quat = array([gauss(0, 1), gauss(0, 1), gauss(0, 1), gauss(0, 1)], float64)
    quat = quat / norm(quat)

    # Convert the quaternion to a rotation matrix.
    quaternion_to_R(quat, R)



class Main:
    """The program."""

    def __init__(self):
        """Set up some initial variables."""

        # The total number of calculations to perform by all slave processors.
        self.N = 20000

        # Variable for counting the completed calculations (to demonstrate slave->master communication).
        self.num = 0

        # The invariant data to pass to the slaves once.
        self.real_length = 2.0
        self.vect = array([1, 2, 3], float64)
        self.vect = self.vect / norm(self.vect) * self.real_length


    def run(self):
        """This required method executes the entire program."""

        # Send the invariant data to the slaves' data stores.
        send_data_to_slaves('vect', self.vect)

        # Initialise the Processor box singleton.
        processor_box = Processor_box()

        # Loop over the slaves.
        num = processor_box.processor.processor_size()
        for i in range(num):
            # Partition out the calculations to one slave.
            slave = Test_slave_command(N=self.N/num)

            # Initialise the memo object.
            memo = Test_memo(name="Memo_"+repr(i), sum_fn=self.sum_fn)

            # Queue the slave command and its memo.
            processor_box.processor.add_to_queue(slave, memo)

        # Execute the calculations, waiting for completion.
        processor_box.processor.run_queue()

        # Calculate the average length.
        ave_len = fetch_data('total_length') / self.N

        # Final program printout.
        print("\n\nTotal number of calculations: %s" % self.num)
        print("Real length: %s" % self.real_length)
        print("Averaged vector length: %s" % ave_len)


    def sum_fn(self, num):
        """Method for slave->master communication.

        This is stored in the memo object and used by the result_command on the master (itself invoked by the slave command on the slave processors) to pass the slave data to the master.

        @param num:     The number of calculations performed by a given slave processor.
        @type num:      int
        """

        # Sum the total number of calculations performed on the slaves.
        self.num += num



class Test_memo(Memo):
    """The memo object containing data and functions for the results_command."""

    def __init__(self, name, sum_fn):
        """Store some data for the result command.

        @param name:    A name for the memo.
        @type name:     str
        @param sum_fn:  A method for summing the number calculations performed by all slaves.
        @type sum_fn:   method
        """

        # Store the arguments for later use by the result_command.
        self.name = name
        self.sum_fn = sum_fn



class Test_result_command(Result_command):
    """The result command for processing the results from the slaves on the master."""

    def __init__(self, processor, memo_id=None, num=None, length=None, completed=True):
        """Store all the slave results for processing on the master.

        @param processor:   The slave processor object.
        @type processor:    Processor instance
        @keyword memo_id:   The ID of the corresponding memo object (currently serves no purpose).
        @type memo_id:      int
        @keyword num:       The number of calculations performed by the slave.  This is an example of data transfer from the slave to master processor.
        @type num:          int
        @keyword length:    The sum of vector lengths.
        @type length:       float
        @keyword completed: A flag saying if the calculation on the slave processor completed correctly.
        @type completed:    bool
        """

        # Execute the base class __init__() method (essential).
        super(Test_result_command, self).__init__(processor=processor, completed=completed)

        # Store the arguments.
        self.memo_id = memo_id
        self.num = num
        self.length = length


    def run(self, processor, memo):
        """Essential method for doing something with the results from the slave processors.

        This code will run on the master processor.


        @param processor:   The slave processor object.
        @type processor:    Processor instance
        @param memo:        The slave's corresponding memo object.
        @type memo:         Memo instance
        """

        # Random printout.
        print("%s, %s calculations completed." % (memo.name, self.num))

        # Calling a method on the master.
        memo.sum_fn(self.num)

        # Get the master processor data store (this is running on the master!).
        data_store = fetch_data_store()

        # Initialise the total length variable if the other slaves have not created it.
        if not hasattr(data_store, 'total_length'):
            data_store.total_length = 0.0

        # Sum the lengths.
        data_store.total_length += self.length



class Test_slave_command(Slave_command):
    """The slave command for use by the slave processor."""

    def __init__(self, N=0):
        """Set up the slave command object for the slave processor.

        @keyword N:     The number of calculations for the slave to perform.
        @type N:        int
        """

        # Store the argument.
        self.N = N

        # Initialise the rotation matrix.
        self.R = zeros((3, 3), float64)

        # Initialise some variables for the calculation.
        self.length = 0.0


    def run(self, processor, completed=False):
        """Essential method for performing calculations on the slave processors.

        @param processor:   The slave processor object.
        @type processor:    Processor instance
        @keyword completed: A flag specifying if the slave calculation is completed.  This is currently meaningless, but will be passed to this run() method anyway so it needs to be present.
        @type completed:    bool
        """

        # Get the invariant data from this slave's data store.
        vect = fetch_data('vect')

        # Perform some random useless time-consuming stuff.
        num_calcs = 0
        for i in range(self.N):
            # Random rotation matrix.
            R_random_hypersphere(self.R)

            # Rotate the vector.
            new_vect = dot(self.R, vect)

            # The length sum.
            self.length += norm(new_vect)

            # Keep track of the number of calculations.
            num_calcs += 1

        # Process the results on the master.
        processor.return_object(Test_result_command(processor, memo_id=self.memo_id, num=num_calcs, length=self.length))



# Set up the processor.
processor = load_multiprocessor(FABRIC, Application_callback(master=Main()), processor_size=PROCESSOR_NUM, verbosity=1)

# Run without profiling.
if not PROFILE:
    processor.run()

# Run with profiling.
else:
    # Replace the default profiling printout function.
    profile.Profile.print_stats = print_stats

    # Execute with profiling.
    profile.runctx('processor.run()', globals(), locals())
