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

#WRITTEN BY DENE BOWDALO

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

import __builtin__
from collections import Counter
import crossnetwork_processing_modules
from mpi4py import MPI
import numpy as np
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 PROCESSED TEMPORARY REFERENCES - GET REFERENCES FROM OUTPUT NUMPY FILES IN TEMPORARY DIRECTORY
        #THESE DIFFER FROM INITIAL STATION REFERENCES USED IN N1, AS SOME OF THE STATION REFERENCES MAY NO LONGER BE VALID - AS COULD HAVE HAD CORRUPT/BADLY FORMATTED/OR NO DATA  
        processed_temporary_references = crossnetwork_processing_modules.get_processed_temporary_references()
        
        #DO CHECK TO DEAL WITH NUMPY STATION SPECIFIC FILES THAT HAVE DUPLICATES
        #(THIS HAPPENS IN INSTANCES WHERE STATION CODES ARE UNKNOWN UNTIL READING FILES)
        base_references = []
        #find all instances of duplicated temporary files
        for reference in processed_temporary_references:
            base_references.append(reference.split('_dup')[0])
        duplicated_references = sorted([item for item, count in Counter(base_references).items() if count > 1])

        #CREATE QUEUE OF DATA TO SEND TO SLAVES
        self.work = Queue.Queue()
        for ref in duplicated_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):
    crossnetwork_processing_modules.resolve_duplicate_reference_files(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


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:
        master(comm)
        print 'N2 DONE'
    else:
        slave(comm)


if __name__ == '__main__':
    main()