import copy
import csv
import glob
import os
import subprocess 
import sys
import time

import numpy as np

from crossnetwork_processing_functions import get_default_parameters_per_network, natural_key


class SubmitGreasy:

    def __init__(self, job_name, GHOST_level, arguments, run_details, pipeline_logfile=None):

        #add passed variables to self
        self.job_name = job_name
        self.GHOST_level = GHOST_level
        self.arguments = arguments
        self.run_details = run_details
        self.pipeline_logfile = pipeline_logfile

        #create greasy arguments file
        self.create_greasy_arguments_file()

        #submit greasy job
        self.submit_greasy_job()


    def create_greasy_arguments_file(self):
        
        '''function which creates greasy arguments text file storing all different tasks to run by greasy'''

        #define list to store chunked argument files
        argument_files = []

        #if len(arguments) >= 20,000 force 'job_array_limit' to be 200, or >=10,000 set to 150 (100 is default) 
        if len(self.arguments) >= 20000:
            self.run_details['job_array_limit'] = 200
        elif len(self.arguments) >= 10000:
            self.run_details['job_array_limit'] = 150

        #define greasy arguments filenames
        if self.GHOST_level == 'pipeline':
            argument_files.append('{}/run/pipeline_arguments/{}.txt'.format(self.run_details['process_root'],self.run_details['slurm_job_ID']))
        elif self.GHOST_level == 'stage':
            #get the CPU chunk size -- set initially as miniumum number of arguments per file 
            N_arguments_per_file_minimum = copy.deepcopy(self.run_details['stage_chunk_size'])
            N_arguments_per_file = copy.deepcopy(self.run_details['stage_chunk_size'])
            #divide the number of arguments by the CPU chunk size, to determine how many argument files will be needed to submit all jobs
            N_submit_files = int(np.ceil(len(self.arguments)/self.run_details['stage_chunk_size']))
            #set argument remainder as 0 initially
            argument_remainder = 0
            #if the number of argument files is greater than the job array limit (i.e the limit on the number of argument files that can be processed simultaneously) 
            #then add adjust minimum N arguments per file appropriately (i.e. split extra arguments across the maximum number of argument files evenly)        
            if N_submit_files > self.run_details['job_array_limit']:
                #update the minimum number of arguments per file
                N_arguments_per_file_minimum = int(np.floor(len(self.arguments)/self.run_details['job_array_limit']))
                N_arguments_per_file = copy.deepcopy(N_arguments_per_file_minimum)
                #if the number of extra arguments does not divide equally between all files get the remainder
                argument_remainder = int(len(self.arguments)%self.run_details['job_array_limit'])
                #if have argument remainder then update N_arguments_per_file variable to be 1 greater than minimum for first file written (and for all files thereafter until  remainder is accounted for)
                if argument_remainder > 0:
                    N_arguments_per_file = N_arguments_per_file_minimum + 1
                    #subtract 1 from the argument remainder
                    argument_remainder -= 1
                #set N submit files as N of job array limit
                N_submit_files = copy.deepcopy(self.run_details['job_array_limit'])

            #create file which will store a list of all stage argument filenames
            greasy_file = open('{}/run/stage_arguments/{}_{}/{}.grz'.format(self.run_details['process_root'],self.run_details['network'],self.run_details['parameter'],self.run_details['GHOST_stage']), 'w')
            #create all stage argument filenames
            for ii in range(N_submit_files):
                stage_argument_fname = '{}/run/stage_arguments/{}_{}/{}_{}.txt'.format(self.run_details['process_root'],self.run_details['network'],self.run_details['parameter'],self.run_details['GHOST_stage'],ii)
                argument_files.append(stage_argument_fname)
                greasy_file.write('{}\n'.format(stage_argument_fname))
            greasy_file.close()

        #initialise variables to count N argument files written and N lines written
        file_ii = 0
        n_lines_written = 0
        #define special characters that will need to be escaped in string written        
        special_characters = ['(',')']

        #open first arguments file for writing
        arguments_file = open(argument_files[file_ii], 'w')

        #iterate through different arguments, writing line by line to arguments file
        for arguments_ii, line in enumerate(self.arguments):
            if type(line) == list:
                str_to_write = ' '.join([str(elem) for elem in line])
            else:
                str_to_write = '{}'.format(line)

            #escape certain special characters in str_to_write
            for ch in special_characters:
                str_to_write = str_to_write.replace(ch,'\{}'.format(ch))
        
            #write arguments str to current file
            if self.GHOST_level == 'pipeline':
                arguments_file.write('python -u {}/processing_modules/execute_pipeline.py {}\n'.format(self.run_details['process_root'],str_to_write))
            elif self.GHOST_level == 'stage':
                #affix slurm_job_ID, network and parameter to str_to_write 
                str_to_write = '{} {} {} {} {}'.format(self.run_details['GHOST_stage'], self.run_details['slurm_job_ID'],self.run_details['network'],self.run_details['parameter'],str_to_write) 
                arguments_file.write('python -u {}/processing_modules/execute_stage.py {}\n'.format(self.run_details['process_root'],str_to_write))

                #iterate n lines written    
                n_lines_written += 1
        
                #if have written sufficient arguments to file (and not currently processing last argument)
                #then close current file and open another
                if (n_lines_written == N_arguments_per_file) & (arguments_ii < (len(self.arguments)-1)):
                    #reset n lines written
                    n_lines_written = 0
                    #close current arguments file
                    arguments_file.close()
                    #iterate n arguments files written
                    file_ii += 1
                    #if have argument remainder, write an extra argument to the next file
                    if argument_remainder > 0:
                        N_arguments_per_file = N_arguments_per_file_minimum + 1
                        #subtract 1 from the argument remainder
                        argument_remainder -= 1
                    #otherwise N arguments to write to next file should be minimum
                    else:
                        N_arguments_per_file = copy.deepcopy(N_arguments_per_file_minimum)
                    #open new arguments file
                    arguments_file = open(argument_files[file_ii], 'w')
                
        #close current arguments file
        arguments_file.close()


    def submit_greasy_job(self):

        '''function which submits greasy job to process GHOST pipeline/stage'''

        #define directory where jobs are run from 
        self.run_dir = self.run_details['process_root']+'/run'

        #define submit/log/argument directories
        if self.GHOST_level == 'pipeline':
            self.submit_dir = '{}/run/submit_pipeline'.format(self.run_details['process_root'])
            self.log_dir =    '{}/run/pipeline_logs'.format(self.run_details['process_root'])
            self.arguments_dir = '{}/run/pipeline_arguments'.format(self.run_details['process_root'])
        elif self.GHOST_level == 'stage':
            self.submit_dir = '{}/run/submit_stage/{}_{}'.format(self.run_details['process_root'],self.run_details['network'],self.run_details['parameter'])
            self.log_dir =    '{}/run/stage_logs/{}_{}/{}'.format(self.run_details['process_root'],self.run_details['network'],self.run_details['parameter'],self.run_details['GHOST_stage'])
            self.arguments_dir = '{}/run/stage_arguments/{}_{}'.format(self.run_details['process_root'],self.run_details['network'],self.run_details['parameter'])

        #write slurm script for submission of greasy job
        self.create_slurm_submission_script()

        #start timer on job
        start_job_time = time.time()

        #submit slurm script
        submit_complete = False
        while not submit_complete:

            submit_process =  subprocess.Popen(['sbatch', '{}'.format(self.slurm_fname)], cwd=self.submit_dir, stdout=subprocess.PIPE)
            submit_status = submit_process.communicate()[0]
            submit_return_code = submit_process.returncode

            #if sbatch fails, time out for 60 seconds and then try again
            if submit_return_code != 0:
                time.sleep(60)
                continue
            else:
                submit_complete = True

            #take a 1 second pause after submittal (to help slurm scheduler)
            time.sleep(1)

        #now job has been submitted
        #monitor job in queue (every 10 seconds) until there are 0 left in the squeue
        all_tasks_finished = False

        while not all_tasks_finished:
            cmd = ['squeue', '-h', '-n', '{}'.format(self.job_name)]
            squeue_process = subprocess.Popen(cmd, stdout=subprocess.PIPE, encoding='utf8')
            squeue_status = squeue_process.communicate()[0]
            squeue_return_code = squeue_process.returncode

            #if squeue fails, time out for 60 seconds and then try again
            if squeue_return_code != 0:
                time.sleep(60)
                continue
            else:
                n_jobs_in_queue = len(squeue_status.split('\n')[:-1])
                #if number of jobs in queue > 0, then sleep for 10 seconds and then check again how many jobs there are in queue
                if n_jobs_in_queue > 0:
                    time.sleep(10)
                    continue
                else:
                    all_tasks_finished = True

        #get end job time
        self.end_job_time = time.time()

        #get total job time (minutes)
        self.total_job_time = (self.end_job_time - start_job_time)/60.

        #if stage job, and job no longer in the squeue, then now check the outcome of the job
        if self.GHOST_level == 'stage':
            self.check_GHOST_stage_outcome()

    def create_slurm_submission_script(self):

        '''function that writes a slurm submission shell script that submits a greasy job'''

        #get current BSC machine
        machine = os.environ['BSC_MACHINE']

        #create slurm_fname (job_name + 'sh.')
        self.slurm_fname = self.job_name+'.sh'
    
        #get all relevant argument files for current GHOST level 
        if self.GHOST_level == 'pipeline':
            argument_files = ['{}/{}.txt'.format(self.arguments_dir,self.run_details['slurm_job_ID'])]
        elif self.GHOST_level == 'stage':
            argument_files = sorted(glob.glob('{}/{}_*.txt'.format(self.arguments_dir,self.run_details['GHOST_stage'])), key=natural_key)

        #read how many lines are in first arguments file
        with open(argument_files[0]) as f: 
            for ii, line in enumerate(f):
                pass
            N_arguments = ii + 1

        #set the number of simultaneous tasks to run based on the GHOST level 
        #if the GHOST level == 'pipeline', then cap the number of simultaneously running tasks to be the defined pipeline CPU chunk size
        if self.GHOST_level == 'pipeline':
            max_tasks = self.run_details['pipeline_chunk_size']
        #else, if the GHOST level == 'stage', then cap the number of simultaneously running tasks to be the defined stage CPU chunk size  
        elif self.GHOST_level == 'stage':
            max_tasks = self.run_details['stage_chunk_size']

        #if the number of arguments is > capped max tasks, then set N simultaneous tasks to be the max tasks permitted
        if N_arguments > max_tasks:
            n_simultaneous_tasks = copy.deepcopy(max_tasks) 
        #else, if the number of arguments is <= capped max tasks, then set N simultaneous tasks to be N arguments
        else:
            n_simultaneous_tasks = copy.deepcopy(N_arguments)

        #if processing get maximum walltime for specific GHOST level
        if self.GHOST_level == 'pipeline':
            #if QOS == 'debug', then it is always 2 hours
            if self.run_details['qos_pipeline'] in ['gp_debug','debug']:
                if self.run_details['pipeline_max_walltime'] == 'default': 
                    max_walltime = '02:00:00'
                else:
                    max_walltime = self.run_details['pipeline_max_walltime']
            #standard queues should be 48 hours
            elif self.run_details['qos_pipeline'] in ['gp_bsces','bsc_es']:
                if self.run_details['pipeline_max_walltime'] == 'default': 
                    max_walltime = '48:00:00'
                else:
                    max_walltime = self.run_details['pipeline_max_walltime']
            #PRACE AND RES accounts should be 72 hours
            elif self.run_details['qos_pipeline'] in ['gp_resa','prace','class_a','class_b']:
                if self.run_details['pipeline_max_walltime'] == 'default': 
                    max_walltime = '72:00:00'
                else:
                    max_walltime = self.run_details['pipeline_max_walltime']
            #RESc account should be 24 hours
            elif self.run_details['qos_pipeline'] in ['gp_resc']:
                if self.run_details['pipeline_max_walltime'] == 'default': 
                    max_walltime = '24:00:00'
                else:
                    max_walltime = self.run_details['pipeline_max_walltime']

        elif self.GHOST_level == 'stage':
            #if QOS == 'debug', then it is always 2 hours
            if self.run_details['qos_stage'] in ['gp_debug','debug']:
                if self.run_details['stage_max_walltime'] == 'default': 
                    max_walltime = '02:00:00'
                else:
                    max_walltime = self.run_details['stage_max_walltime']
            #standard queues should be 48 hours
            elif self.run_details['qos_stage'] in ['gp_bsces','bsc_es']:
                if self.run_details['stage_max_walltime'] == 'default': 
                    max_walltime = '48:00:00'
                else:
                    max_walltime = self.run_details['stage_max_walltime']
            #PRACE AND RES accounts should be 72 hours
            elif self.run_details['qos_pipeline'] in ['gp_resa','prace','class_a','class_b']:
                if self.run_details['stage_max_walltime'] == 'default': 
                    max_walltime = '72:00:00'
                else:
                    max_walltime = self.run_details['stage_max_walltime']
            #RESc account should be 24 hours
            elif self.run_details['qos_pipeline'] in ['gp_resc']:
                if self.run_details['stage_max_walltime'] == 'default': 
                    max_walltime = '24:00:00'
                else:
                    max_walltime = self.run_details['stage_max_walltime']

        #create slurm submission script for specific GHOST stage
        file = open(self.submit_dir+'/'+self.slurm_fname,"w")

        file.write("#!/bin/bash\n")
        file.write("\n")
        file.write("#SBATCH --job-name={}\n".format(self.job_name))
        file.write("#SBATCH --account={}\n".format(self.run_details['bsc_account']))
        file.write("#SBATCH --ntasks={}\n".format(n_simultaneous_tasks))
        file.write("#SBATCH --cpus-per-task=1\n")
        file.write("#SBATCH --ntasks-per-node={}\n".format(n_simultaneous_tasks))
        #fix number of nodes to be 1 (for faster execution)
        file.write("#SBATCH --nodes 1\n")
        file.write("#SBATCH --time={}\n".format(max_walltime))
        if self.GHOST_level == 'pipeline':
            file.write("#SBATCH --qos={}\n".format(self.run_details['qos_pipeline']))
        elif self.GHOST_level == 'stage':
            file.write("#SBATCH --qos={}\n".format(self.run_details['qos_stage']))
            file.write("#SBATCH --array=1-{}\n".format(len(argument_files)))
            #if machine is mn5, and stage is S1, and high memory nodes are requested, set it.
            if (machine == 'mn5') & (self.run_details['GHOST_stage'] == 'S1') & (self.run_details['high_memory_S1']):
                file.write("#SBATCH --constraint=highmem\n") 
            #if machine is mn5, and stage is T2, and high memory nodes are requested, set it.
            elif (machine == 'mn5') & (self.run_details['GHOST_stage'] == 'T2') & (self.run_details['high_memory_T2']):
                file.write("#SBATCH --constraint=highmem\n") 
        if not self.run_details['error_out_file']:
            file.write("#SBATCH --output=/dev/null\n")
            file.write("#SBATCH --error=/dev/null\n")      
        file.write("\n") 
        file.write("source {}/load_modules.sh\n".format(self.run_dir))
        file.write('eval "$(conda shell.bash hook)"\n')
        file.write("conda activate ghost-env\n")
        file.write("export SLURM_CPU_BIND=none\n")
        #file.write("export GREASY_NWORKERS=$SLURM_NPROCS\n") 
        if self.GHOST_level == 'pipeline':
            file.write("export GREASY_LOGFILE={}/{}.log\n".format(self.submit_dir,self.job_name))
            file.write("\n")
            file.write("greasy {}".format(argument_files[0]))
        elif self.GHOST_level == 'stage':
            file.write("export GREASY_LOGFILE={}/{}_$SLURM_ARRAY_TASK_ID.log\n".format(self.submit_dir,self.job_name))
            file.write("arguments_store={}/{}.grz\n".format(self.arguments_dir,self.run_details['GHOST_stage']))
            file.write("argument_file=$(cat $arguments_store | awk -v var=$SLURM_ARRAY_TASK_ID 'NR==var {print $1}')\n")
            file.write("\n")
            file.write("greasy $argument_file")
    
        #close file
        file.close()

    def check_GHOST_stage_outcome(self):

        '''function that checks the outcome of a GHOST stage.

           all tasks processed in a GHOST stage return a log file prefixed with an appropriate process code

           these process codes are:
           0: Process completed without issue
           1: Process completed, but station no longer continuing in pipeline
           2: Caught error in process
           3: Uncaught error in process

           The current pipeline is terminated if there are any 2 or 3 process codes.
           The current pipeline is terminated if all parallel tasks return a process code of 1.
           Additionally the current pipeline is terminated if a submitted task has not returned a log file  
           Otherwise, the pipeline is continued.

           Also the process time and queue time of stage are calculated.
        '''

        #define special characters that need to be escaped in strings written        
        special_characters = ['(',')']

        #create list of all output stage log root filenames should be output for stage
        #gather all of the argument files used to submit jobs and extract job references used to prefix log files for specific stage
        argument_files = sorted(glob.glob('{}/{}**.txt'.format(self.arguments_dir, self.run_details['GHOST_stage'])))     
        job_references = []
        for argument_file in argument_files:
            with open(argument_file) as f:
                csv_reader = csv.reader(f, delimiter=' ') 
                for line in csv_reader:
                    if self.run_details['GHOST_stage'] in ['T1','T2','P1']:
                        job_references.append('{}_{}_{}'.format(line[7], line[8], line[9]))
                    elif self.run_details['GHOST_stage'] in ['E1','M1']:
                        job_references.append('{}_{}'.format(line[7], line[8]))
                    else:
                        job_references.append('{}'.format(line[7]))

        #replace escaped special characters in job references
        for job_reference_ii in range(len(job_references)):
            for ch in special_characters:
                job_references[job_reference_ii] = job_references[job_reference_ii].replace('\{}'.format(ch), ch)

        #get output stage log files
        output_stage_log_files = glob.glob('{}/*.out'.format(self.log_dir))
        #get output stage log file job references
        output_stage_log_file_job_references = np.array([f.split('/')[-1][2:-4] for f in output_stage_log_files])
        #get log file process codes
        process_codes = np.array([int(f.split('/')[-1][:1]) for f in output_stage_log_files], dtype=np.int32)
    
        #check if any jobs have failed, and for what reason
        #if all the process codes returned are 1 (no stations continuing in pipeline)
        if np.all(process_codes == 1):
            self.pipeline_logfile.write('{} COMPLETED ON {:.2f} MINUTES (BUT 0 STATIONS REMAIN).\nTERMINATING PIPELINE HERE'.format(self.run_details['GHOST_stage'],(self.end_job_time-self.run_details['pipeline_start_time'])/60.))
            self.pipeline_logfile.close()
            sys.exit()
        #if a process code of 2 is returned in any of the jobs (caught error)
        elif 2 in process_codes:
            caught_error_job_references = sorted(output_stage_log_file_job_references[np.where(process_codes == 2)[0]])
            self.pipeline_logfile.write('{} COMPLETED ON {:.2f} MINUTES (WITH CAUGHT ERROR/S). TERMINATING PIPELINE HERE\nJOB REFERENCES THAT FAILED ARE: {}'.format(self.run_details['GHOST_stage'],(self.end_job_time-self.run_details['pipeline_start_time'])/60.,caught_error_job_references))
            self.pipeline_logfile.close()
            sys.exit()
        #if one of the process codes are 3 (uncaught error)
        elif 3 in process_codes:
            non_caught_error_job_references = sorted(output_stage_log_file_job_references[np.where(process_codes == 3)[0]])
            self.pipeline_logfile.write('{} DID NOT FINISH DUE TO UNCAUGHT ERROR/S. TERMINATING PIPELINE HERE\nJOB REFERENCES THAT FAILED ARE: {}'.format(self.run_details['GHOST_stage'],non_caught_error_job_references))
            self.pipeline_logfile.close()
            sys.exit()            
        #if one of the jobs did not finish
        elif (len(output_stage_log_file_job_references) < len(job_references)):
            non_finished_job_references = sorted(list(set(output_stage_log_file_job_references) ^ set(job_references)))
            self.pipeline_logfile.write('{} DID NOT FINISH DUE TO UNFINISHED JOB/S. TERMINATING PIPELINE HERE\nJOB REFERENCES THAT DID NOT FINISH ARE: {}'.format(self.run_details['GHOST_stage'],non_finished_job_references))
            self.pipeline_logfile.close()
            sys.exit()
  
        #now know all jobs have completed
        #create list to store process times per successful task (only of successful jobs)
        task_process_times = []
        zero_process_code_inds = np.where(process_codes == 0)[0]
        for zero_process_code_ind in zero_process_code_inds:
            with open(output_stage_log_files[zero_process_code_ind], "rb") as f:
                f.seek(-2, 2)
                while f.read(1) != b"\n":
                    f.seek(-2, 1)
                process_time = f.read().decode('utf8').strip()
            task_process_times.append(float(process_time))
        #get all job references which finished succesfully
        success_job_references = sorted(output_stage_log_file_job_references[zero_process_code_inds])

        #calculate the process time for the stage (taken to be the maximum process time across all successful tasks)
        stage_process_time = np.max(task_process_times)

        #calculate the queue time for the stage (total stage time - stage process time)
        stage_queue_time = self.total_job_time - stage_process_time

        #check if one of the process codes is 1 (i.e. one of stations no longer continues in the pipeline)
        if 1 in process_codes:
            non_continuing_job_references = sorted(output_stage_log_file_job_references[np.where(process_codes == 1)[0]])
            if self.run_details['GHOST_stage'] in ['T2','P1','E1','M1']:
                self.pipeline_logfile.write('{} COMPLETED ON {:.2f} MINUTES ({:.2f} MINUTES TO PROCESS STAGE: {:.2f} MINUTES PROCESSING, {:.2f} MINUTES QUEUING).\nSOME JOBS HAD NO VALID DATA FOR PROCESSING: {}\n'.format(self.run_details['GHOST_stage'],(self.end_job_time-self.run_details['pipeline_start_time'])/60.,self.total_job_time,stage_process_time,stage_queue_time,non_continuing_job_references))
            else:
                self.pipeline_logfile.write('{} COMPLETED ON {:.2f} MINUTES ({:.2f} MINUTES TO PROCESS STAGE: {:.2f} MINUTES PROCESSING, {:.2f} MINUTES QUEUING).\nSOME STATIONS ARE NOT CONTINUING IN PIPELINE: {}\n'.format(self.run_details['GHOST_stage'],(self.end_job_time-self.run_details['pipeline_start_time'])/60.,self.total_job_time,stage_process_time,stage_queue_time,non_continuing_job_references))
        #otherwise, all process codes are 0 (all processed successfully)
        else:
            self.pipeline_logfile.write('{} COMPLETED SUCCESSFULLY ON {:.2f} MINUTES ({:.2f} MINUTES TO PROCESS STAGE: {:.2f} MINUTES PROCESSING, {:.2f} MINUTES QUEUING)\n'.format(self.run_details['GHOST_stage'],(self.end_job_time-self.run_details['pipeline_start_time'])/60.,self.total_job_time,stage_process_time,stage_queue_time))

        #if stage == T1, get all the temporal output resolutions/regions that some station data has been successfully processed for 
        if self.run_details['GHOST_stage'] == 'T1':
            self.valid_temporal_resolution_codes = np.unique([success_job_reference.split('_')[-2] for success_job_reference in success_job_references])
            self.valid_regions = np.unique([success_job_reference.split('_')[-1] for success_job_reference in success_job_references])