"""
Parallel management for SNNet.

Robert A McDougal, 2011
"""

import cPickle
import os
import sys
import platform
import time
import threading
import random
from fileManagement import delete_file, process_temp_dir

process_data = 'tmp_snnet_%d_%d_status.pkl'

class WorkerThread(threading.Thread):
    """Worker thread, runs a function on data from queue."""
    def __init__(self, command, items, exception_list, post_processing = None,
                 continue_after_failure=False):
        threading.Thread.__init__(self)
        self.command = command
        self.items = items
        self.exception_list = exception_list
        self.post_processing = post_processing
        self.continue_after_failure = continue_after_failure
    
    def run(self):
        while (not self.items.empty()) and (self.exception_list[0] is None or self.continue_after_failure):
            have_data = False
            try:
                data = self.items.get(timeout = 1)
                have_data = True
            except Queue.Empty:
                pass
            if have_data:
                try:
                    result = self.command(data)
                    if self.post_processing is not None:
                        self.post_processing(result)
                except:
                    self.exception_list[0] = sys.exc_info()
        

def get_process_data(source_id, temp_dir, return_id = False, uid = 0):
    """Returns pickled data from a remote process."""
    if isinstance(source_id, tuple):
        source_id = source_id[0]
    filename = os.path.join(process_temp_dir(temp_dir), process_data % (source_id, uid))
    file = open(filename, 'rb')
    data = cPickle.load(file)
    file.close()
    delete_file(filename)
    exception = data['exception']
    if exception is not None:
        raise exception[0], exception[1]
    if return_id:
        return data['return'], data['id']
    else:
        return data['return']


def launch_process(command, param, temp_dir = '.', id = 0, uid = 0):
    """Launch a new process, running function command on data param.
    
    Data is saved to temp_dir/process_data.
    
    Parameters:
        command -- the function to run
        param -- the data to pass to command
        temp_dir -- a string or function evaluating to a string
                    representing a temporary directory
        id -- an identifier to indicate which process terminated
        uid -- a random number (will be needed to get_process_data)

    Returns:
    
        pid of new process

    Note:
    
    This only works on operating systems supporting fork (e.g. OSX, Linux).
    """
    new_process = os.fork()
    if new_process == 0:
        # runs only child process
        exception = None
        result = None
        try:
            result = command(param)
        except:
            exception = sys.exc_info()[0 : 2]
        
        # pass the status to a file    
        file = open(os.path.join(process_temp_dir(temp_dir), process_data % (os.getpid(), uid)), 'wb')
        cPickle.dump({'exception': exception, 'return': result, 'id': id}, file, 2)
        file.close()

        # shutdown the process
        os._exit(os.EX_OK)
    return new_process


def parallel_run_over_queue(command, items, num_parallel = 1,
                            post_processing = None,
                            temp_dir = '.',
                            continue_after_failure=False):
    """Run command for each of the values in items.
    
    Parameters:
        command -- the function to run
        items -- a queue of items to run the command on
        num_parallel -- the maximum number to run simultaneously
        post_processing -- a function to run on the return values
        temp_dir -- a string or function evaluating to a string
                    representing a temporary directory
        continue_after_failure -- should we continue after an exception

    Note:
    
    Uses processes for Linux and OSX; otherwise uses threads.
    
    post_processing can be a list corresponding to different functions
    for different tasks (in the same order as items).
    """
    if (platform.system() in ['Darwin', 'Linux']) and (num_parallel > 1):
        # launch the initial processes
        active = 0
        id = 0
        uid = random.randint(0, 10000)
        while active < num_parallel and not(items.empty()):
            active += 1
            launch_process(command, items.get(), temp_dir, id, uid)
            id += 1
        
        # launch new processes as results come in
        while active > 0:
            try:
                result, id = get_process_data(os.wait(), temp_dir, True, uid)
                success = True
            except:
                success = False
                if not continue_after_failure:
                    raise
            if success:
                if post_processing is not None:
                    if isinstance(post_processing, list):
                        post = post_processing[id]
                    else:
                        post = post_processing
                    post(result)
            active -= 1
            if not items.empty():
                active += 1
                launch_process(command, items.get(), temp_dir, id, uid)
                id += 1
    else:
        exception_list = [None]
        num_parallel = min(num_parallel, items.qsize())
        threads = []
        
        # create and start worker threads
        for i in range(num_parallel):
            if isinstance(post_processing, list):
                post = post_processing[i]
            else:
                post = post_processing
            t = WorkerThread(command, items, exception_list, post)
            t.setDaemon(True)
            t.start()
            threads.append(t)
        
        # wait for threads to finish
        for t in threads:
            t.join()
        
        e = exception_list[0]
        if e is not None:
            raise e[0], e[1], e[2]
