import copy
import datetime
import os
import sys
import time
import random

import numpy as np
import pandas as pd

from crossnetwork_processing_functions import check_GHOST_output_directories, check_GHOST_stage_temporary_directories, clean_native_directory, clean_output_resolution_directories, get_data_files, get_YYYY_from_filenames, get_YYYYMM_from_filenames, get_station_references_from_filenames, get_unique_processed_references, get_relevant_yearmonths, drop_stations, get_data_from_reference
from geographic_functions import get_stations_in_region
from submit_greasy import SubmitGreasy


class ProcessGHOSTStage:

    def __init__(self, run_details, parameter_details, network_details, pipeline_logfile):
        
        #add passed variables to self
        self.run_details = run_details
        self.parameter_details = parameter_details
        self.network_details = network_details
        self.pipeline_logfile = pipeline_logfile

    def S1(self):

        '''GHOST PIPELINE STAGE: S1

           In this stage:
             -- All raw network files for a species are parsed to station specific standardised numpy files
        '''

        #add GHOST stage to run_details
        self.run_details['GHOST_stage'] = 'S1'

        #get all relevant data files for specific parameter, per network
        data_files = get_data_files(self.run_details, self.parameter_details)

        #check you have some data to process
        if len(data_files) > 0:

            #Check/create GHOST output native directories
            check_GHOST_output_directories(self.run_details, self.parameter_details, stage_type='S')

            #Check/create/clean temporary data directories of stage
            check_GHOST_stage_temporary_directories(self.run_details, self.parameter_details, stage='S1')

            #set up parallelisation to process data files, dependant on how observations are reported natively by the data providers
            #1. If file storage type is all stations per year, parallelise by year
            #2. If file storage type is all stations per day, parallelise by month (if parallelised by day the number of files would be too extreme)
            #3. If file storage is station per year/multiyear, parallelise by station

            #if file storage type is "all_per_year", get all different data years to parallelise by 
            if self.network_details['file_storage_type'] == 'all_per_year':
                param_to_parallelise_by = get_YYYY_from_filenames(data_files, self.run_details)

            #elif file storage type is "all_per_day", get all different data months to parallelise by (if parallelised by day the number of files would be too extreme)
            elif self.network_details['file_storage_type'] == 'all_per_day':             
                param_to_parallelise_by = get_YYYYMM_from_filenames(data_files, self.run_details)

            #otherwise, get all station references to parallelise by
            else:
                #get all unique station references from data filenames
                param_to_parallelise_by = get_station_references_from_filenames(data_files, self.run_details) 

            #randomly shuffle param_to_parallelise_by
            random.shuffle(param_to_parallelise_by)

            #param_to_parallelise_by = param_to_parallelise_by[:100]
            #param_to_parallelise_by = [x for x in param_to_parallelise_by if x in ['gsn']]               

            #submit greasy job/s to process all years, or station data in parallel
            job_name = 'S1_{}_{}'.format(self.run_details['network'],self.run_details['parameter'])
            greasy_inst = SubmitGreasy(job_name,'stage',param_to_parallelise_by,self.run_details,pipeline_logfile=self.pipeline_logfile)

        #if have zero data files then stage outcome == process code: 2 (i.e. no data)
        else:
            self.pipeline_logfile.write('S1 NOT PROCESSED AS HAVE NO DATA FOR PARAMETER.\nTERMINATING PIPELINE HERE')
            self.pipeline_logfile.close()
            sys.exit()


    def S2(self):

        '''GHOST PIPELINE STAGE: S2

           In this stage:
            -- Concatenate station files, handling duplicate/overlapping files  
            -- Split station data by significantly changing locational information (creating new stations where necessary)
        '''

        #add GHOST stage to run_details
        self.run_details['GHOST_stage'] = 'S2'

        #get all unique station references from output numpy files in previous stage temporary directory 
        #the stage will be parallelised by station reference
        processed_temporary_references = get_unique_processed_references(self.run_details, self.parameter_details)

        #check you have some data to process
        if len(processed_temporary_references) > 0:

            #Check/create/clean temporary data directories of stage, and lock files of previous stage
            check_GHOST_stage_temporary_directories(self.run_details, self.parameter_details, stage='S1', remove_lock=True)
            check_GHOST_stage_temporary_directories(self.run_details, self.parameter_details, stage='S2')

            #randomly shuffle processed temporary references
            random.shuffle(processed_temporary_references)

            #submit greasy job/s to process station data in parallel
            job_name = 'S2_{}_{}'.format(self.run_details['network'],self.run_details['parameter'])
            greasy_inst = SubmitGreasy(job_name,'stage',processed_temporary_references,self.run_details,pipeline_logfile=self.pipeline_logfile)

        #if have zero station references then stage outcome == process code: 1 (i.e. no data)
        else:
            self.pipeline_logfile.write('S2 NOT PROCESSED AS HAVE NO DATA FOR PARAMETER.\nTERMINATING PIPELINE HERE')
            self.pipeline_logfile.close()
            sys.exit()


    def S3(self):

        '''GHOST PIPELINE STAGE: S3

           In this stage:
            -- Add globally gridded metadata per station
            -- Do Quality Control checks per station
            -- Add temporal classification codes per station
        '''

        #add GHOST stage to run_details
        self.run_details['GHOST_stage'] = 'S3'

        #clean native directory and lock files
        check_GHOST_stage_temporary_directories(self.run_details, self.parameter_details, stage='S2', remove_lock=True)
        clean_native_directory(self.run_details, self.parameter_details)

        #get all unique station references from output numpy files in previous stage temporary directory 
        #the stage will be parallelised by station reference
        processed_temporary_references = get_unique_processed_references(self.run_details, self.parameter_details)     

        #check you have some data to process
        if len(processed_temporary_references) > 0:

            #randomly shuffle processed temporary references
            random.shuffle(processed_temporary_references)

            #submit greasy job/s to process station data in parallel
            job_name = 'S3_{}_{}'.format(self.run_details['network'],self.run_details['parameter'])
            greasy_inst = SubmitGreasy(job_name,'stage',processed_temporary_references,self.run_details,pipeline_logfile=self.pipeline_logfile)

        #if have zero station references then stage outcome == process code: 1 (i.e. no data)
        else:
            self.pipeline_logfile.write('S3 NOT PROCESSED AS HAVE NO DATA FOR PARAMETER.\nTERMINATING PIPELINE HERE')
            self.pipeline_logfile.close()
            sys.exit()


    def T1(self):

        '''GHOST PIPELINE STAGE: T1
      
           In this stage:
            -- Iterate through defined regions and isolate stations to solely those in defined region
            -- If N unique stations per region > max_N_stations, then keep max_N_stations (using total range and N points to select which) and drop the rest
            -- Temporally average station specific region data to desired temporal resolution by year/month
        '''

        #add GHOST stage to run_details
        self.run_details['GHOST_stage'] = 'T1'

        #get all unique station references from output numpy files in previous stage temporary directory 
        #the stage will be parallelised by station reference/temporal resolution
        processed_final_references = get_unique_processed_references(self.run_details, self.parameter_details)

        #initialise arguments list as empty
        arguments = []

        #initialise lists to save station longitudes and latitudes read if not GLOBAL domain
        station_longitudes = []
        station_latitudes = []

        #iterate through defined regions
        for region in self.run_details['regions']:
    
            #isolate stations to just those in region
            if region != 'GLOBAL':
                if len(station_longitudes) == 0:
                    for station_reference in processed_final_references:
                        #read native metadata for each station reference
                        metadata, station_numpy_filename = get_data_from_reference(station_reference,self.run_details,self.parameter_details,read='metadata')
                        station_longitudes.append(metadata['longitude'][0])
                        station_latitudes.append(metadata['latitude'][0])
                relevant_references = get_stations_in_region(processed_final_references, region, station_longitudes, station_latitudes) 
            else:
                relevant_references = copy.deepcopy(processed_final_references)

            #If N unique stations > max_N_stations, then selectively keep max_N_stations stations (using temporal range and N measurements to select which), and drop the rest
            relevant_references = drop_stations(relevant_references, self.run_details, self.parameter_details)   

            #create arguments list of each unique station reference, each unique temporal resolution code and region
            for ref in relevant_references:
                for temporal_resolution_code in self.run_details['temporal_resolutions']:
                    arguments.append([ref, temporal_resolution_code, region])

        #check you have some data to process
        if len(arguments) > 0:

            #Check/create GHOST output directories
            check_GHOST_output_directories(self.run_details, self.parameter_details, stage_type='T')

            #Check/create/clean temporary data directories of stage, and lock files of previous stage
            clean_native_directory(self.run_details, self.parameter_details, remove_lock=True)
            check_GHOST_stage_temporary_directories(self.run_details, self.parameter_details, stage='T1')

            #randomly shuffle arguments list
            random.shuffle(arguments)

            #submit greasy job/s to process station data in parallel
            job_name = 'T1_{}_{}'.format(self.run_details['network'],self.run_details['parameter'])
            greasy_inst = SubmitGreasy(job_name,'stage',arguments,self.run_details,pipeline_logfile=self.pipeline_logfile)

            #have valid_temporal_resolution_codes attribute in greasy_inst object? if so, add to self
            if hasattr(greasy_inst, 'valid_temporal_resolution_codes'):
                self.valid_temporal_resolution_codes = copy.deepcopy(greasy_inst.valid_temporal_resolution_codes)
            #have valid_regions attribute in greasy_inst object? if so, add to self
            if hasattr(greasy_inst, 'valid_regions'):
                self.valid_regions = copy.deepcopy(greasy_inst.valid_regions)   

        #if have zero station references then stage outcome == process code: 1 (i.e. no data)
        else:
            self.pipeline_logfile.write('T1 NOT PROCESSED AS HAVE NO DATA FOR PARAMETER.\nTERMINATING PIPELINE HERE')
            self.pipeline_logfile.close()
            sys.exit()


    def T2(self):

        '''GHOST PIPELINE STAGE: T2
      
           In this stage:
            -- Merge temporary station specific files to 1 netCDF by year/month 
        '''

        #add GHOST stage to run_details
        self.run_details['GHOST_stage'] = 'T2'  

        #clean output resolution directories
        clean_output_resolution_directories(self.run_details, self.parameter_details)

        #create list to store all relevant yearmonths + valid temporal resolution codes
        arguments = []

        #iterate through valid temporal resolutions
        for temporal_resolution_code in self.valid_temporal_resolution_codes:

            #iterate through valid regions
            for region in self.valid_regions:

                #get list of yearmonth strings for which in those time periods exist valid processed station data, of the specific temporal resolution, for the defined region
                yearmonth_strings = get_relevant_yearmonths(self.run_details, self.parameter_details, temporal_resolution_code, region, stage_type='T')      

                #iterate through yearmonths
                for yearmonth_ii, yearmonth_str in enumerate(yearmonth_strings):
                    #append yearmonth_str, temporal resolution code and region  
                    arguments.append([yearmonth_str, temporal_resolution_code, region])

        #check you have some data to process
        if len(arguments) > 0:

            #randomly shuffle arguments list
            random.shuffle(arguments)

            #submit greasy job/s to process yearmonth data in parallel
            job_name = 'T2_{}_{}'.format(self.run_details['network'],self.run_details['parameter'])
            greasy_inst = SubmitGreasy(job_name,'stage',arguments,self.run_details,pipeline_logfile=self.pipeline_logfile)

        #if have zero yearmonths with data then cannot process stage
        else:
            self.pipeline_logfile.write('T2 NOT PROCESSED AS HAVE NO VALID STATION DATA TO WRITE.\nTERMINATING PIPELINE HERE')
            self.pipeline_logfile.close()
            sys.exit()

    def P1(self):

        '''GHOST PIPELINE STAGE: P1
      
           In this stage:
            -- Calculation of data product from other processed components
        '''

        #add GHOST stage to run_details
        self.run_details['GHOST_stage'] = 'P1'

        #clean output resolution directories
        clean_output_resolution_directories(self.run_details, self.parameter_details)

        #create list to store all relevant yearmonths, valid temporal resolution codes, and regions
        arguments = []

        #iterate through valid temporal resolutions
        for temporal_resolution_code in self.run_details['temporal_resolutions']:

            #iterate through defined regions
            for region in self.run_details['regions']:

                #get list of yearmonth strings for which in those time periods exist valid processed station data, of the specific temporal resolution, for the defined region
                yearmonth_strings = get_relevant_yearmonths(self.run_details, self.parameter_details, temporal_resolution_code, region, stage_type='P')      

                #iterate through yearmonths
                for yearmonth_ii, yearmonth_str in enumerate(yearmonth_strings):
                    #append yearmonth_str, temporal resolution code and region  
                    arguments.append([yearmonth_str, temporal_resolution_code, region])

        #check you have some data to process
        if len(arguments) > 0:

            #randomly shuffle arguments list
            random.shuffle(arguments)

            #submit greasy job/s to process yearmonth data in parallel
            job_name = 'P1_{}_{}'.format(self.run_details['network'],self.run_details['parameter'])
            greasy_inst = SubmitGreasy(job_name,'stage',arguments,self.run_details,pipeline_logfile=self.pipeline_logfile)

        #if have zero arguments then cannot process stage
        else:
            self.pipeline_logfile.write('P1 NOT PROCESSED AS HAVE NO VALID DATA.\nTERMINATING PIPELINE HERE')
            self.pipeline_logfile.close()
            sys.exit()

    def E1(self):

        '''GHOST PIPELINE STAGE: E1
      
           In this stage:
            -- Split out EBAS data into individual native networks
        '''

        #add GHOST stage to run_details
        self.run_details['GHOST_stage'] = 'E1'

        #clean output resolution directories
        clean_output_resolution_directories(self.run_details, self.parameter_details)

        #create list to store all relevant yearmonths, valid temporal resolution codes, and regions
        arguments = []

        #iterate through valid temporal resolutions
        for temporal_resolution_code in self.run_details['temporal_resolutions']:

            #iterate through defined regions
            for region in self.run_details['regions']:

                #append temporal resolution code and region  
                arguments.append([temporal_resolution_code, region])

        #check you have some data to process
        if len(arguments) > 0:

            #randomly shuffle arguments list
            random.shuffle(arguments)

            #submit greasy job/s to process yearmonth data in parallel
            job_name = 'E1_{}_{}'.format(self.run_details['network'],self.run_details['parameter'])
            greasy_inst = SubmitGreasy(job_name,'stage',arguments,self.run_details,pipeline_logfile=self.pipeline_logfile)

        #if have zero arguments then cannot process stage
        else:
            self.pipeline_logfile.write('E1 NOT PROCESSED AS HAVE NO VALID DATA.\nTERMINATING PIPELINE HERE')
            self.pipeline_logfile.close()
            sys.exit()

    def M1(self):

        '''GHOST PIPELINE STAGE: M1
      
           In this stage:
            -- Merge data across networks
        '''

        #add GHOST stage to run_details
        self.run_details['GHOST_stage'] = 'M1'

        #clean output resolution directories
        clean_output_resolution_directories(self.run_details, self.parameter_details)

        #create list to store all relevant yearmonths, and valid temporal resolution codes
        arguments = []

        #iterate through valid temporal resolutions
        for temporal_resolution_code in self.run_details['temporal_resolutions']:

            #get list of yearmonth strings for which in those time periods exist valid processed station data, of the specific temporal resolution
            yearmonth_strings = get_relevant_yearmonths(self.run_details, self.parameter_details, temporal_resolution_code, 'GLOBAL', stage_type='M')      

            #iterate through yearmonths
            for yearmonth_ii, yearmonth_str in enumerate(yearmonth_strings):
                #append yearmonth_str, temporal resolution code  
                arguments.append([yearmonth_str, temporal_resolution_code])

        #check you have some data to process
        if len(arguments) > 0:

            #randomly shuffle arguments list
            random.shuffle(arguments)

            #submit greasy job/s to process yearmonth data in parallel
            job_name = 'M1_{}'.format(self.run_details['parameter'])
            greasy_inst = SubmitGreasy(job_name,'stage',arguments,self.run_details,pipeline_logfile=self.pipeline_logfile)

        #if have zero arguments then cannot process stage
        else:
            self.pipeline_logfile.write('M1 NOT PROCESSED AS HAVE NO VALID DATA.\nTERMINATING PIPELINE HERE')
            self.pipeline_logfile.close()
            sys.exit()