#!/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 os

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


network = 'EBAS'
parameter = 'CO'
start_year = '1970'
end_year = '2018'

run_details = {'network':network,'parameter':parameter,'start_year':start_year,'end_year':end_year}
__builtin__.run_details = run_details

###----------------------------------------------------------------------------###
#GET SPECIFIC PARAMETER AND NETWORK ASSOCIATED VARIABLES 
###----------------------------------------------------------------------------###

parameter_details = crossnetwork_processing_modules.get_parameter_details()[run_details['parameter']]
network_details = crossnetwork_processing_modules.get_network_details()
#MAKE CORE PARAMETER DATA GLOBAL ACROSS ALL SUBMODULES
__builtin__.parameter_details = parameter_details
__builtin__.network_details = network_details

###----------------------------------------------------------------------------###
#IMPORT SPECIFIC NETWORK PROCESSING MODULE FOR NETWORK BEING PROCESSED
###----------------------------------------------------------------------------###
network_processor = __import__(run_details['network']+'_processor')

WORKTAG = 1
DIETAG = 0

class Work(object):
    def __init__(self):
	
        #TEST FOR PRESENCE OF TEMPORARY STATION DATA FILE DIRECTORY. USED TO STORE INITIAL NUMPY .npz FILES OF RAW DATA BEFORE QA CHECKS.
        #IF DOES NOT EXIST CREATE IT
        temp_file_dir = '/esarchive/obs/uniyork/%s/native/%s/temporary/'%(run_details['network'],parameter_details['bsc_parameter_name'])
        does_temporary_exist = os.path.isdir(temp_file_dir)
        if does_temporary_exist == False:
            os.makedirs(temp_file_dir)
        #DELETE ALL CURRENT TEMPORARY STATION DATA FILES FOR PARAMETER/NETWORK (IN CASE OF FAILED PREVIOUS RUN AND HAVE NOT BEEN CLEANED)
        for zippath in glob.iglob(os.path.join(temp_file_dir, '*.npz')):
            os.remove(zippath)
            
        #TEST FOR PRESENCE OF FINAL STATION PROCESS INFORMATION DIRECTORY. USED FOR GETTING ALL VALID FULLY PROCESSED STATION REFERENCES 
        #IF DOES NOT EXIST CREATE IT
        final_file_process_dir = '/esarchive/obs/uniyork/%s/native/%s/final_process_valid/'%(run_details['network'],parameter_details['bsc_parameter_name'])
        does_temporary_exist = os.path.isdir(final_file_process_dir)
        if does_temporary_exist == False:
            os.makedirs(final_file_process_dir)
        #DELETE ALL CURRENT FINAL STATION FILES FOR PARAMETER/NETWORK (IN CASE OF FAILED PREVIOUS RUN AND HAVE NOT BEEN CLEANED)
        for current_file in glob.glob('%s*'%(final_file_process_dir)):
            os.remove(current_file)
	
        #GET ALL DATA FILES FOR SPECIFIC PARAMETER, FOR NETWORK 
        data_files = crossnetwork_processing_modules.get_data_files()
    
        #GET ALL DIFFERENT YEARS OF FILENAMES DATA IS SEPARATED BETWEEN
        if network_details['file_storage_type'] == 'all_per_year':
            data_years = range(int(run_details['start_year']),int(run_details['end_year']))
        elif network_details['file_storage_type'] == 'site_per_year':
            data_years = range(1)
        
        #BRANCH HERE, DEPENDENT ON HOW OBSERVATIONS ARE REPORTED
        #1. IF FILE STORAGE TYPE IS ALL SITES PER YEAR, PROCESS ALL STATIONS YEAR BY YEAR
        #2. IF FILE STORAGE TYPE IS SITE PER YEAR, PROCESS EACH ALL YEARS STATION BY STATION
    
        #ITERATE THROUGH ALL YEARS OF DATA FILES (JUST 1 YEAR IF HAVE SITES REPORTED PER YEAR)
        for year in data_years:
        #IF FILE STORAGE TYPE IS ALL SITES PER YEAR, LIMIT DATA FILES TO RELEVANT YEARLY FILE
            if network_details['file_storage_type'] == 'all_per_year':
                data_files_to_process = [s for s in data_files if str(year) in s]
                print 'PROCESSING YEAR: %s'%(year)
            #OTHERWISE DO NOT FILTER FILES BY YEAR
            elif network_details['file_storage_type'] == 'site_per_year':
                data_files_to_process = data_files
        
            #GET ALL STATION REFERENCES & REFERENCES TO CONNECT METADATA/DATA FROM DATA FILE/S
            station_references = crossnetwork_processing_modules.get_station_references(data_files_to_process) 
            #station_references = station_references[51:]
            #station_references = station_references[:3]
        
            #HOW MANY STATIONS TO PROCESS?
            print 'NUMBER OF STATIONS: %s'%(len(station_references))
    
            #LOAD EXTERNAL METADATA INTO MEMORY (*IF EXISTS FOR NETWORK)        
            if network_details['external_metadata_file'] != '':
                external_metadata = crossnetwork_processing_modules.read_external_metadata(station_references)
            else:
                external_metadata = [[] for _ in xrange(len(station_references))]
    
            #IF FILE STORAGE TYPE IS SITE PER YEAR, SPLIT DATA FILES INTO LIST OF LISTS CONTAINING STATION SPECIFIC DATA FILES
            if network_details['file_storage_type'] == 'site_per_year':
                data_files_to_process = crossnetwork_processing_modules.generate_data_files_LoL(station_references,data_files_to_process)

            #ZIP VARIABLES THAT WILL BE PROCESSED
            zipped_list = zip(range(len(station_references)),station_references,data_files_to_process,external_metadata)
    
            #CREATE QUEUE OF DATA TO SEND TO SLAVES
            #queue = []
            #q = Queue()
            #for zipped_item in zipped_list:
            #    q.put(zipped_item)
            self.work = zipped_list

    def get_next(self):
        if len(self.work) == 0:
            return None
        return self.work.pop(0)
 

def do_work(work):
    network_processor.read_data(work)
    return


def master(comm):
    num_procs = comm.Get_size()
    status = MPI.Status()
    
    # generate work queue
    wq = Work()

    # 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)
        #process_result(result)

        # 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)
        del result
        #process_result(result)

    # 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()
    #comm.Barrier()
    #start = MPI.Wtime()
    
    if my_rank == 0:
        master(comm)
    else:
        slave(comm)
    
    #comm.Barrier()
    #end = MPI.Wtime()
    #print 'time:', end - start


if __name__ == '__main__':
    main()
    
