#!/usr/bin/env python
# -*- coding: utf-8 -*-

#WRITTEN BY DENE BOWDALO

###------------------------------------------------------------------------------------###
###IMPORT MODULES
###------------------------------------------------------------------------------------###

import __builtin__
import crossnetwork_processing_modules
import glob
from mpi4py import MPI
import numpy as np
import os
import Queue

###------------------------------------------------------------------------------------###
###------------------------------------------------------------------------------------###

WORKTAG = 1
DIETAG = 0

#LOAD GHOST CONFIGURATION VARIABLES
config_frame = np.load('/esarchive/obs/ghost/process/run/GHOST_configuration.npz')
run_details = config_frame['run_details'].item()
parameter_details = config_frame['parameter_details'].item()
network_details = config_frame['network_details'].item()

#MAKE RUN/PARAMETER/NETWORK DATA GLOBAL ACROSS ALL SUBMODULES
__builtin__.run_details = run_details
__builtin__.parameter_details = parameter_details
__builtin__.network_details = network_details

class Work(object):
    def __init__(self):
    
        #GET LIST OF FINAL PROCESSED REFERENCES FROM WRITTEN FINAL STATION FILES DIRECTORY 
        #THESE AGAIN DIFFER FROM REFERENCES USED BEFORE, AS SOME OF THE STATION REFERENCES MAY NOW BE SPLIT OUT (DUE TO CHANGING KEY METADATA), OR SOME MAY NOT EXIST ANY MORE (DUE TO MISSING KEY METADATA)
        final_file_process_dir = '/esarchive/obs/ghost/%s/native/%s/final_process_valid/'%(run_details['network'],parameter_details['bsc_parameter_name'])
        final_processed_references = []
        for file in glob.glob('%s*'%(final_file_process_dir)):
            final_processed_references.append(file.split('/')[-1])

        #CREATE QUEUE OF DATA TO SEND TO SLAVES
        self.work = Queue.Queue()
        for ref in final_processed_references:
            self.work.put(ref)

    def get_next(self):
        if self.work.empty():
            return None
        return self.work.get()
        
    def test_if_start_empty(self):
        if self.work.empty():
            return True
        else:
            return False
 

def do_work(work):
    modal_metadata = crossnetwork_processing_modules.process_modal_metadata(work)
    return
    
def master(comm):
    num_procs = comm.Get_size()
    status = MPI.Status()
    
    # generate work queue
    wq = Work()
    
    #test if queue is empty to start with
    if wq.test_if_start_empty():
        # if empty - Tell all the slaves to exit by sending an empty message with DIETAG
        for rank in xrange(1, num_procs):
            comm.send(0, dest=rank, tag=DIETAG)
        return

    # Seed the slaves, send one unit of work to each slave (rank)
    for rank in xrange(1, num_procs):
        work = wq.get_next()
        if not work: break
        comm.send(work, dest=rank, tag=WORKTAG)
    
    # Loop over getting new work requests until there is no more work to be done
    while True:
        work = wq.get_next()
        if not work: break
    
        # Receive results from a slave
        result = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)

        # Send the slave a new work unit
        comm.send(work, dest=status.Get_source(), tag=WORKTAG)
    
    # No more work to be done, receive all outstanding results from slaves
    for rank in xrange(1, num_procs): 
        result = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)

    # Tell all the slaves to exit by sending an empty message with DIETAG
    for rank in xrange(1, num_procs):
        comm.send(0, dest=rank, tag=DIETAG)
        
    return result


def slave(comm):
    my_rank = comm.Get_rank()
    status = MPI.Status()

    while True:
        # Receive a message from the master (the work to do)
        work = comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
        # Check the tag of the received message
        if status.Get_tag() == DIETAG: break 

        # Do the work and send it back
        result = do_work(work)
        comm.send(result, dest=0)
        
def main():
    comm = MPI.COMM_WORLD
    my_rank = comm.Get_rank()
    my_name = MPI.Get_processor_name()
    
    if my_rank == 0:
        result = master(comm)
        print 'M1 DONE'
    else:
        slave(comm)


if __name__ == '__main__':
    main()