import copy
import os
import shutil
import sys
import time
import subprocess

import numpy as np
import pandas as pd

from crossnetwork_processing_functions import (get_configuration_variables, get_default_parameters_per_network, 
                                               get_default_product_parameters_per_network, 
                                               get_default_temporal_resolutions_per_matrix, 
                                               check_GHOST_pipeline_outcome)
from GHOST_standards import standard_parameters, standard_temporal_resolutions
from submit_greasy import SubmitGreasy

#start timer on GHOST processing
start_GHOST_time = time.time()

#get slurm job ID --> passed from job submittal script
slurm_job_ID = sys.argv[1]
    
#get configuration variables
run_details, _, _ = get_configuration_variables(slurm_job_ID, startup=True)

#make a copy of the configuration file indexed by the slurm job ID
shutil.copy2(run_details['process_root']+'/run/GHOST_configuration.py',run_details['process_root']+'/run/GHOST_configurations/configuration_{}.py'.format(slurm_job_ID))

#create variable to track submit iterations through 
submit_iteration = 0

#get yearmonth strs between start date and end date
yearmonth_strs = [date.strftime("%Y%m") for date in pd.date_range('{}-{}'.format(run_details['start_date'][:4], run_details['start_date'][4:6]), '{}-{}'.format(run_details['end_date'][:4], run_details['end_date'][4:6]), freq='MS', inclusive='left')]

#iterate through submitting pipelines until have nothing left to process
while True:

    #if first iteration, then gather arguments for pipelines
    if submit_iteration == 0:

        #initialise list to store arguments to submit pipelines
        arguments = []

        #create a dictionary of parameter dependencies
        #(i.e. parameters which can only be processed after others are processed first)
        dependencies = {}
        for network in run_details['networks']:
            dependencies[network] = {}
            for parameter in get_default_product_parameters_per_network(network):
                dependencies[network][parameter] = get_default_product_parameters_per_network(network)[parameter]['components']

        # create dict to store current stage of pipeline per network and parameter, and if pipeline is stil active or not 
        pipeline_status = {}
        #get which stages to first process
        #handle different types of stages in order by priority

        # 1. processing stages
        processing_stages = [stage for stage in run_details['stages_to_process'] if stage in ['S1','S2','S3','T1','T2']]
        if len(processing_stages) > 0:
            current_stages_to_process = ','.join(processing_stages)

        #2. product stage
        elif 'P1' in run_details['stages_to_process']: 
            current_stages_to_process = 'P1'

        #3. split EBAS stage
        elif 'E1' in run_details['stages_to_process']: 
            current_stages_to_process = 'E1'
        
        #4. merge stage - overwrite networks with merge networks
        elif 'M1' in run_details['stages_to_process']: 
            current_stages_to_process = 'M1'

        #iterate through networks and gather arguments to submit pipelines
        #first, handle case where current stage is not merge stage
        if current_stages_to_process != 'M1':
            #iterate through networks
            for network in run_details['networks']:

                #add dict for pipeline status per network  
                pipeline_status[network] = {}

                #get parameters to process across all stages
                #get default parameters per network if set
                parameter_defaults = get_default_parameters_per_network(network)
                parameter_defaults_product = get_default_product_parameters_per_network(network)
                #take those given otherwise
                if run_details['parameters'] != 'default':
                    parameters = copy.deepcopy(run_details['parameters'])

                #have processing stages?
                if len(processing_stages) > 0:
                    #set stage status
                    status = 'alive'

                    #set parameters
                    if run_details['parameters'] == 'default':
                        parameters = copy.deepcopy(parameter_defaults)

                    #add information about stage to pipeline_status
                    for parameter in parameters:
                        #ensure parameter can be processed in processing stages
                        if parameter in parameter_defaults:
                            #create dictionary in pipeline_status for network/parameter
                            pipeline_status[network][parameter] = {'current_stages':current_stages_to_process, 'stages_remaining':[], 'processed_stages':[], 'status':status, 'pipeline_start_time':0}

                #have product stage?
                if 'P1' in run_details['stages_to_process']:
                    #set stage status
                    if current_stages_to_process == 'P1':
                        status = 'alive'
                    else:
                        status = 'hold'

                    #set parameters
                    if run_details['parameters'] == 'default':
                        parameters = copy.deepcopy(parameter_defaults_product)

                    #add information about stage to pipeline_status
                    for parameter in parameters:

                        #ensure parameter is a product, to allow product stage to be processed
                        if parameter in parameter_defaults_product:
                            if parameter not in pipeline_status[network]:
                                #create dictionary in pipeline_status for network/parameter
                                pipeline_status[network][parameter] = {'current_stages':'P1', 'stages_remaining':[], 'processed_stages':[], 'status':status, 'pipeline_start_time':0}
                            else:
                                #append stage to stages_remaining
                                pipeline_status[network][parameter]['stages_remaining'].append('P1')

                #have split EBAS stage, and network is EBAS?
                if ('E1' in run_details['stages_to_process']) & (network == 'EBAS'): 
                    #set stage status
                    if current_stages_to_process == 'E1':
                        status = 'alive'
                    else:
                        status = 'hold'

                    #set parameters
                    if run_details['parameters'] == 'default':
                        parameters = np.unique(parameter_defaults + parameter_defaults_product)

                    #add information about stage to pipeline_status
                    for parameter in parameters:
                        if parameter not in pipeline_status[network]:
                            #create dictionary in pipeline_status for network/parameter
                            pipeline_status[network][parameter] = {'current_stages':'E1', 'stages_remaining':[], 'processed_stages':[], 'status':status, 'pipeline_start_time':0}
                        else:
                            #append stage to stages_remaining
                            pipeline_status[network][parameter]['stages_remaining'].append('E1')

        #if have merge stage to process, then set setup processing for stage
        if 'M1' in run_details['stages_to_process']: 
            #iterate through merge networks
            for network in run_details['merge_networks']:

                #add dict for pipeline status per network  
                pipeline_status[network] = {}

                #get parameters to process across all stages
                #get default parameters per network if set
                if run_details['parameters'] == 'default':
                    parameters = list(standard_parameters.keys())
                #take those given otherwise
                else:
                    parameters = copy.deepcopy(run_details['parameters'])

                #set stage status
                #if processing merge stage now, create logfile
                if current_stages_to_process == 'M1':
                    status = 'alive'
                else:
                    status = 'hold'

                #add information about stage to pipeline_status, and create pipeline_logfile
                for parameter in parameters:
                    if parameter not in pipeline_status[network]:
                        #create dictionary in pipeline_status for network/parameter
                        pipeline_status[network][parameter] = {'current_stages':'M1', 'stages_remaining':[], 'processed_stages':[], 'status':status, 'pipeline_start_time':0}
                    else:
                        #append stage to stages_remaining
                        pipeline_status[network][parameter]['stages_remaining'].append('M1')

        #iterate through pipeline status dict, and create arguments to submit pipelines which are alive
        #iterate through networks

        for network in pipeline_status:
            #iterate through parameters
            for parameter in pipeline_status[network]:

                # get parameter details
                parameter_details = standard_parameters[parameter]

                #create pipeline_logfile
                pipeline_logfile = open('{}/run/pipeline_logs/{}_{}.out'.format(run_details['process_root'],network,parameter),"w")
                pipeline_logfile.write('**************************************************************************************************************\n')        
                pipeline_logfile.write('BEGINNING PROCESSING OF STAGES FOR {}:{}\n'.format(network,parameter))
                pipeline_logfile.write('**************************************************************************************************************\n')
                pipeline_logfile.close()

                #determine if the current parameter has any dependencies to be processed first, if so put on hold
                #if no dependencies are queued then write logfile for parameter informing that it cannot be processed
                if (pipeline_status[network][parameter]['current_stages'] == 'P1') or ('P1' in pipeline_status[network][parameter]['stages_remaining']):
                    #get dependencies for parameter
                    param_dependencies = dependencies[network][parameter]
                    #iterate through dependent parameters
                    for param in param_dependencies:

                        #if dependent parameter not set to be processed, check if previously processed files exist for current date range 
                        if param not in pipeline_status[network]:

                            # get dependent parameter details
                            param_details = standard_parameters[param]

                            #check if dependent parameter has been previously processed and exist for current date range
                            exist_files = False

                            #get temporal resolutions
                            if run_details['temporal_resolutions'] == 'default':
                                temporal_resolution_codes = get_default_temporal_resolutions_per_matrix(parameter_details['matrix'])
                            else:
                                temporal_resolution_codes = copy.deepcopy(run_details['temporal_resolutions'])
                            temporal_resolutions = [standard_temporal_resolutions[temporal_resolution_code]['temporal_resolution_path'] for temporal_resolution_code in temporal_resolution_codes] 

                            #check if file exist for yearmonth and temporal resolution
                            #iterate through yearmonths
                            for yearmonth_str in yearmonth_strs:
                                if exist_files:
                                    break
                                #iterate through temporal rersolutions
                                for temporal_resolution in temporal_resolutions:

                                    #check if file exists
                                    if os.path.isfile(run_details['data_root']+'{}/{}/{}/{}/{}_{}.nc'.format(network,run_details['ghost_version'],temporal_resolution,param_details['bsc_parameter_name'],param_details['bsc_parameter_name'],yearmonth_str)):
                                        exist_files = True
                                        break

                            #if not all neccessary files exist, then do not submit pipeline for parameter, and terminate pipeline logfile
                            if not exist_files:
                                pipeline_status[network][parameter]['status'] = 'killed'
                                pipeline_logfile = open('{}/run/pipeline_logs/{}_{}.out'.format(run_details['process_root'],network,parameter),"a")
                                pipeline_logfile.write('PIPELINE NOT PROCESSED AS DEPENDENT PARAMETERS NOT PROCESSED.\nTERMINATING PIPELINE HERE')
                                pipeline_logfile.close()
                                break

                        #if dependent parameter set to be processed, and current stage is P1, then put current parameter on hold
                        elif pipeline_status[network][parameter]['current_stages'] == 'P1':
                            pipeline_status[network][parameter]['status'] = 'hold'

                #add argument for pipeline if it is alive
                if pipeline_status[network][parameter]['status'] == 'alive':
                    #add start time to pipeline_status if starting pipeline
                    if len(pipeline_status[network][parameter]['processed_stages']) == 0:
                        pipeline_status[network][parameter]['pipeline_start_time'] = '{:.0f}'.format(time.time())
                    #add argument
                    arguments.append([slurm_job_ID, network, parameter, current_stages_to_process, pipeline_status[network][parameter]['pipeline_start_time']])

                    
    #print('PIPELINE STATUS:', pipeline_status)
    #print('ARGUMENTS TO SUBMIT:', arguments)
    #print()

    #if no more pipelines to submit, break out of loop
    if len(arguments) == 0:
        break

    #submit pipeline for jobs (and monitor outcomes)
    SubmitGreasy('GPipe_{}'.format(slurm_job_ID), 'pipeline', arguments, run_details)

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

    #now determine if submitted pipelines were successful, and gather arguments to submit new pipelines

    #iterate through arguments and check if pipeline did not finish correctly
    for arg_ii, (slurm_job_ID, network, parameter, stages_to_process, pipeline_start_time) in enumerate(arguments):

        #get pipeline logfile name
        pipeline_logfile = '{}/run/pipeline_logs/{}_{}.out'.format(run_details['process_root'],network,parameter)

        #check if pipeline did not finish correctly
        lines = subprocess.check_output(['tail', '-n', '2', pipeline_logfile], encoding='utf8').strip()
        #if pipeline did not finish correctly then do not submit more jobs for that pipeline
        if 'TERMINATING PIPELINE' in lines:
            pipeline_status[network][parameter]['status'] = 'killed'
            
    #iterate through pipeline status dict, and create arguments to submit next set of pipelines
    arguments = []

    #iterate through network
    for network in pipeline_status:
        #iterate through parameter
        for parameter in pipeline_status[network]:

            #move all alive pipelines to next stages
            if pipeline_status[network][parameter]['status'] == 'alive':

                #add previous stages to processed_stages
                pipeline_status[network][parameter]['processed_stages'].append(pipeline_status[network][parameter]['current_stages'])
                #check if have stages remaining, if so, then update current_stages and stages_remaining
                if len(pipeline_status[network][parameter]['stages_remaining']) > 0:
                    pipeline_status[network][parameter]['current_stages'] = pipeline_status[network][parameter]['stages_remaining'][0] 
                    pipeline_status[network][parameter]['stages_remaining'] = pipeline_status[network][parameter]['stages_remaining'][1:]
                #else, if have no stages remaining, finalise pipeline logfile, and continue to next parameter
                else:
                    pipeline_status[network][parameter]['current_stages'] = [] 
                    pipeline_status[network][parameter]['stages_remaining'] = []
                    pipeline_status[network][parameter]['status'] = 'finished'
                    total_pipeline_time = (time.time() - float(pipeline_status[network][parameter]['pipeline_start_time']))/60.
                    pipeline_logfile = open('{}/run/pipeline_logs/{}_{}.out'.format(run_details['process_root'],network,parameter),"a")
                    pipeline_logfile.write('\n') 
                    pipeline_logfile.write('PROCESSING OF ALL DESIRED STAGES FOR {}:{} COMPLETE IN {:.2f} MINUTES'.format(network,parameter,total_pipeline_time))
                    pipeline_logfile.close()
                    continue

                #for the product stage, check on dependent parameters first
                if pipeline_status[network][parameter]['current_stages'] == 'P1':
                    #get dependencies for parameter
                    param_dependencies = dependencies[network][parameter]
                    #iterate through dependent parameters
                    for param in param_dependencies:

                        #if dependent parameter set to be processed, check if pipeline finished or was killed
                        if param in pipeline_status[network]:

                            #check if dependent pipeline was killed early, if so, terminate pipeline logfile
                            if pipeline_status[network][param]['status'] == 'killed':
                                pipeline_status[network][parameter]['status'] = 'killed'
                                pipeline_logfile = open('{}/run/pipeline_logs/{}_{}.out'.format(run_details['process_root'],network,parameter),"a")
                                pipeline_logfile.write('P1 NOT PROCESSED AS DEPENDENT PARAMETERS DID NOT FINISH PROCESSING CORRECTLY.\nTERMINATING PIPELINE HERE')
                                pipeline_logfile.close()
                                break

                            #else, ensure dependent parameter stage is past product stage, otherwise stage is put on hold
                            elif (('S1' in pipeline_status[network][param]['stages_remaining']) or ('S2' in pipeline_status[network][param]['stages_remaining']) or
                                ('S3' in pipeline_status[network][param]['stages_remaining']) or ('T1' in pipeline_status[network][param]['stages_remaining']) or
                                ('T2' in pipeline_status[network][param]['stages_remaining']) or ('P1' in pipeline_status[network][param]['stages_remaining'])):
                                pipeline_status[network][parameter]['status'] = 'hold'
                                break

            #check if on hold pipelines can be started
            elif pipeline_status[network][parameter]['status'] == 'hold':

                #set variable to determine if ready to start pipeline
                ready_to_start = True

                #check if dependent components have now been processed
                if pipeline_status[network][parameter]['current_stages'] == 'P1':
                    #get dependencies for parameter
                    param_dependencies = dependencies[network][parameter]
                    #iterate through dependent parameters
                    for param in param_dependencies:
                        
                        #if dependent parameter set to be processed, check if pipeline finished or was killed
                        if param in pipeline_status[network]:

                            #check if dependent pipeline was killed early, if so, terminate pipeline logfile
                            if pipeline_status[network][param]['status'] == 'killed':
                                pipeline_status[network][parameter]['status'] = 'killed'
                                pipeline_logfile = open('{}/run/pipeline_logs/{}_{}.out'.format(run_details['process_root'],network,parameter),"a")
                                pipeline_logfile.write('P1 NOT PROCESSED AS DEPENDENT PARAMETERS DID NOT FINISH PROCESSING CORRECTLY.\nTERMINATING PIPELINE HERE')
                                pipeline_logfile.close()
                                ready_to_start = False
                                break

                            #check if have any stages remaining to process, if so do not resume pipeline
                            elif (('S1' in pipeline_status[network][param]['stages_remaining']) or ('S2' in pipeline_status[network][param]['stages_remaining']) or
                                ('S3' in pipeline_status[network][param]['stages_remaining']) or ('T1' in pipeline_status[network][param]['stages_remaining']) or
                                ('T2' in pipeline_status[network][param]['stages_remaining']) or ('P1' in pipeline_status[network][param]['stages_remaining'])):
                                ready_to_start = False
                                break
                    
                #for merge stage, check if all pipelines for parameter are finished
                elif pipeline_status[network][parameter]['current_stages'] == 'M1':
                    #iterate through networks
                    for net in pipeline_status:
                        # ensure network is not a merge network 
                        if net not in run_details['merge_networks']:

                            #if pipeline has been killed, then cannot process merge stage. terminate pipeline logfile
                            if pipeline_status[net][parameter]['status'] == 'killed':
                                pipeline_status[network][parameter]['status'] = 'killed'
                                pipeline_logfile = open('{}/run/pipeline_logs/{}_{}.out'.format(run_details['process_root'],network,parameter),"a")
                                pipeline_logfile.write('M1 NOT PROCESSED AS AT LEAST 1 NETWORK PROCESSING PIPELINE FAILED.\nTERMINATING PIPELINE HERE')    
                                pipeline_logfile.close()
                                ready_to_start = False
                                break

                            #elif pipeline is not finished then continue to hold on processing merge stage 
                            elif pipeline_status[net][parameter]['status'] != 'finished':
                                ready_to_start = False
                                break
                
                # if pipeline can be started, update pipeline status
                if ready_to_start:
                    pipeline_status[network][parameter]['status'] = 'alive'
                    
            #add argument for pipeline if it is alive
            if pipeline_status[network][parameter]['status'] == 'alive':
                #add start time to pipeline_status if starting pipeline
                if len(pipeline_status[network][parameter]['processed_stages']) == 0:
                    pipeline_status[network][parameter]['pipeline_start_time'] = '{:.0f}'.format(time.time())
                #add argument for pipeline
                arguments.append([slurm_job_ID, network, parameter, pipeline_status[network][parameter]['current_stages'], pipeline_status[network][parameter]['pipeline_start_time']])

    #add 1 to submit_iteration
    submit_iteration += 1

#end timer on GHOST processing
run_details['end_GHOST_time'] = (time.time() - start_GHOST_time)/60.

#print('PIPELINE STATUS:', pipeline_status)

#check summary outcome of all pipelines
check_GHOST_pipeline_outcome(pipeline_status, run_details)