from collections import Counter
import copy
import glob
import subprocess
import time

from netCDF4 import Dataset
import numpy as np

from crossnetwork_processing_functions import set_file_permissions_ownership, check_directory_existence
from crossnetwork_processing_functions import create_GHOST_stage_logfile
from GHOST_standards import standard_temporal_resolutions, EBAS_native_networks, EBAS_network_priorities

class SplitEBAS:

    '''Class which splits EBAS data for a component into individual native networks'''
    
    def __init__(self, temporal_resolution_code, region, run_details, parameter_details, network_details):

        #add passed variables to self
        self.temporal_resolution_code = temporal_resolution_code
        self.region = region
        self.run_details = run_details
        self.parameter_details = parameter_details
        self.network_details = network_details

        self.run_details['log_file_str'] += 'SPLITTING EBAS DATA\n'

    def do_work(self):
        '''Function which does executes the functions of the class in a pipeline'''

        #split EBAS data into individual native networks
        self.do_split()

    def do_split(self):
        '''Function which splits EBAS data into individual native networks'''

        #set network directory where to read/write data
        if self.region == 'GLOBAL':
            network_write_dir = self.run_details['network']
        else:
            network_write_dir = '{}_{}'.format(self.run_details['network'], self.region)

        #get output temporal resolution specific variables
        self.temporal_resolution_vars = standard_temporal_resolutions[self.temporal_resolution_code]
        print('0')
        #set relevant files
        self.relevant_files = glob.glob('{}{}/{}/{}/{}/{}*.nc'.format(self.run_details['data_root'],network_write_dir,self.run_details['ghost_version'],self.temporal_resolution_vars['temporal_resolution_path'],
                                                                      self.parameter_details['bsc_parameter_name'],self.parameter_details['bsc_parameter_name']))
        
        #if no relevant files then return logfile to inform there is no data to split
        if len(self.relevant_files) == 0:
            self.run_details['log_file_str'] += 'NO FILES TO SPLIT FOR COMPONENT'  
            create_GHOST_stage_logfile(1,self.run_details)

        #iterate through files and save associated networks per unique station references
        self.station_references = {}
        print('1')
        for relevant_file in self.relevant_files:
            root = Dataset(relevant_file)
            for station_idx, station_ref in enumerate(root['station_reference'][:]):
                if station_ref not in self.station_references:
                    self.station_references[station_ref] = []
                assoc_networks = [EBAS_native_networks[net] for net in root['associated_networks'][station_idx].split(':')[0].split(' ')]
                self.station_references[station_ref].extend(assoc_networks)
            root.close()

        print('2')
        #get chosen native network per station reference
        chosen_options = self.get_EBAS_native_network()                        
        print('3')
        #write native network data
        self.write_to_EBAS_native_network(chosen_options)
        print('4')
        #return valid process logfile (0)
        self.run_details['log_file_str'] += str((time.time() - self.run_details['stage_process_start_time'])/60.)
        print('5')
        create_GHOST_stage_logfile(0,self.run_details)

    def get_EBAS_native_network(self):

        """Get a chosen framework per station reference.
        This is done simply taking the most common occurring framework
        per station reference.
        If there is an equal count, this is broken using framework 
        priorities, 
        Also correct WMO framework if neccessary.
        """

        chosen_options = {}
        for station_ref in self.station_references.keys():
            counts = Counter(self.station_references[station_ref])
            options = []
            for n in range(len(counts)):
                most_common = counts.most_common()[n][0]
                most_common_count = counts.most_common()[n][1]
                if n == 0:
                    previous_common_count = copy.deepcopy(most_common_count)
                elif n >= 1:
                    if previous_common_count != most_common_count:
                        break
                options.append(most_common)

            # if have more than 1 framework remaining, keep that with highest priority
            chosen_option = options[np.argmin([EBAS_network_priorities.index(opt) for opt in options])]
            # correct WMO network if neccessary
            if (chosen_option == 'EBAS-WMO_WDCA') & (self.parameter_details['matrix'] == 'gas'):
                if self.parameter_details['bsc_parameter_name'] in ['sconcco','sconcch4']:
                    chosen_option = 'EBAS-WMO_WDCGG'
                    self.run_details['log_file_str'] += 'CHANGING {} CHOSEN OPTION TO WDCGG'.format(station_ref)
                else:
                    chosen_option = 'EBAS-WMO_WDCRG'
                    self.run_details['log_file_str'] += 'CHANGING {} CHOSEN OPTION TO WDCRG'.format(station_ref)

            elif (chosen_option in ['EBAS-WMO_WDCGG','EBAS-WMO_WDCRG']) & (self.parameter_details['matrix'] in ['pm','pm10','pm2.5','pm1']):
                chosen_option = 'EBAS-WMO_WDCA'
                self.run_details['log_file_str'] += 'CHANGING {} CHOSEN OPTION TO WDCA'.format(station_ref)

            chosen_options[station_ref] = chosen_option

        return chosen_options

    def write_to_EBAS_native_network(self,chosen_options):
        
        """Separate EBAS network into native frameworks"""

        for relevant_file in self.relevant_files:
            root = Dataset(relevant_file)
            file_station_refs = root['station_reference'][:]

            new_networks = np.unique(list(chosen_options.values()))
            for new_network in new_networks:
                new_relevant_file = relevant_file.replace("/EBAS/", "/{}/".format(new_network))

                relevant_station_refs = [stn for stn, net in chosen_options.items() if (net == new_network) & (stn in root['station_reference'][:])]
            
                relevant_station_inds = []
                for station_ref in relevant_station_refs:   
                    relevant_station_inds.append(np.where(root['station_reference'][:] == station_ref)[0][0])

                relevant_station_inds = np.sort(np.array(relevant_station_inds, dtype=np.int32))

                if len(relevant_station_inds) > 0:
                    #create file (and dirs)
                    check_directory_existence('/'.join(new_relevant_file.split('/')[:-1])+'/', directories_not_to_test='{}'.format(self.run_details['data_root']))
                    new_root = Dataset(new_relevant_file, 'w', format="NETCDF4")

                    #copy global attributes
                    for name in root.ncattrs():
                        if name == 'title':
                            new_root.title = root.title.replace('EBAS', new_network)
                        else:
                            new_root.setncattr(name, root.getncattr(name))
                    
                    #copy dimensions
                    for name, dimension in root.dimensions.items():
                        if name == 'station':
                            new_root.createDimension(name, len(relevant_station_inds))
                        else:
                            new_root.createDimension(name, len(dimension))
    
                    #copy variables
                    for name, variable in root.variables.items():

                        #get current data type
                        current_dtype = variable.dtype

                        #create variable
                        if '_FillValue' in root[name].ncattrs():
                            if not self.run_details['compression_level']:
                                var_root = new_root.createVariable(name, variable.datatype, variable.dimensions, fill_value=255, compression=self.run_details['compression_type'], shuffle=self.run_details['compression_shuffle'])
                            else:
                                var_root = new_root.createVariable(name, variable.datatype, variable.dimensions, fill_value=255, compression=self.run_details['compression_type'], complevel=self.run_details['compression_level'], shuffle=self.run_details['compression_shuffle'])
                        else:
                            if current_dtype == str:
                                var_root = new_root.createVariable(name, variable.datatype, variable.dimensions)
                            else:
                                if not self.run_details['compression_level']:
                                    var_root = new_root.createVariable(name, variable.datatype, variable.dimensions, compression=self.run_details['compression_type'], shuffle=self.run_details['compression_shuffle'])
                                else:
                                    var_root = new_root.createVariable(name, variable.datatype, variable.dimensions, compression=self.run_details['compression_type'], complevel=self.run_details['compression_level'], shuffle=self.run_details['compression_shuffle'])

                        #copy variable attributes
                        for att in root[name].ncattrs():
                            if att != '_FillValue':
                                var_root.setncattr(att, root[name].getncattr(att))

                        #get relevant data
                        if 'station' in root[name].dimensions:
                            if len(root[name].shape) == 1:
                                if current_dtype == str:
                                    data = np.array([root[name][ind] for ind in relevant_station_inds], dtype=object)
                                else:
                                    data = root[name][relevant_station_inds]
                            elif len(root[name].shape) == 2:  
                                if current_dtype == str:
                                    data = np.array([root[name][ind,:] for ind in relevant_station_inds], dtype=object)
                                else:
                                    data = root[name][relevant_station_inds,:]
                            elif len(root[name].shape) == 3:
                                if current_dtype == str:
                                    data = np.array([root[name][ind,:,:] for ind in relevant_station_inds], dtype=object)
                                else:
                                    data = root[name][relevant_station_inds,:,:]
                        else:
                            data = root[name][:]

                        # set new network field
                        if name == 'network':
                            data = np.array([new_network]*len(data), dtype=object) 

                        #set variable data
                        new_root[name][:] = data
                
                    #close new netCDF
                    new_root.close()

                    #apply compression (-L1 = deflation level 1), and do not write operation history (-h)
                    #and set to original name
                    submit_list = ['ncks', '-h', '-O', '-L', '1', new_relevant_file, new_relevant_file]
                    record_process =  subprocess.Popen(submit_list, stdout=subprocess.PIPE)
                    record_status = record_process.communicate()[0]
                    record_return_code = record_process.returncode

                    #give 770 permissions for file and make owner bsc32
                    set_file_permissions_ownership(new_relevant_file)   

            #close old netCDF
            root.close()