from collections import Counter
import copy
import itertools
import os
import subprocess
import time

from netCDF4 import Dataset
import numpy as np
import pyproj
from scipy.spatial import cKDTree

from crossnetwork_processing_functions import set_file_permissions_ownership, check_directory_existence
from crossnetwork_processing_functions import create_GHOST_stage_logfile, networks_to_merge
from GHOST_standards import standard_parameters, standard_temporal_resolutions

class Merge:

    '''Class which merges data for a component across networks'''
    
    def __init__(self, yearmonth_str, temporal_resolution_code, run_details, parameter_details, network_details):

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

        self.run_details['log_file_str'] += 'MERGING DATA ACROSS NETWORK IN {}\n'.format(self.yearmonth_str)

        #set network priorities (higher number, bigger priority)
        self.network_priorities = {'AERONET_v3_lev1.5':95,
                                   'AERONET_v3_lev2.0':96,
                                   'BJMEMC':29,
                                   'CANADA_NAPS':80,
                                   'CAPMoN':10,
                                   'CHILE_SINCA':4,
                                   'CNEMC':30,
                                   'EANET':3,
                                   'EBAS-ACTRIS':52,
                                   'EBAS-AMAP':44,
                                   'EBAS-CAMP':45,
                                   'EBAS-COLOSSAL':51,
                                   'EBAS-EMEP':55,
                                   'EBAS-EUCAARI':48,
                                   'EBAS-EUSAAR':47,
                                   'EBAS-GUAN':49,
                                   'EBAS-HELCOM':46,
                                   'EBAS-HTAP':43,
                                   'EBAS-IMPACTS':50,
                                   'EBAS-IMPROVE':53,
                                   'EBAS-Independent':40,
                                   'EBAS-NILU':54,
                                   'EBAS-NOAA_ESRL':56,
                                   'EBAS-NOAA_GGGRN':57,
                                   'EBAS-OECD':42,
                                   'EBAS-UK_DECC':41,
                                   'EBAS-WMO_WDCA':58,
                                   'EBAS-WMO_WDCRG':59,
                                   'EEA_AIRBASE':84,
                                   'EEA_AQ_eReporting':85,
                                   'JAPAN_NIES':70,
                                   'MEXICO_CDMX':35,
                                   'MITECO':20,
                                   'NOAA_ISD':65,
                                   'NOAA_ISD_EU':67,
                                   'NOAA_ISD_IP':68,
                                   'NOAA_ISD_NA':66,
                                   'SEARCH':0,
                                   'UK_AIR':5,
                                   'US_EPA_AirNow_DOS':28,
                                   'US_EPA_AQS':75,
                                   'US_EPA_CASTNET':76,
                                   'US_NADP_AMNet':15,
                                   'US_NADP_AMoN':14,
                                   'WMO_WDCGG':100
                                    }

        # set networks to merge
        self.networks = networks_to_merge(self.run_details['network'])

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

        #merge network files per resolution, per speci, per yyyymm
        self.do_merge()

    def do_merge(self):
        '''Function which merges data for a component across networks'''

        #get output temporal resolution specific variables
        self.temporal_resolution_vars = standard_temporal_resolutions[self.temporal_resolution_code]

        #gather relevant files for yyyymm per network
        self.yyyymm_files = {}
        for network in self.networks: 
            yyyymm_file = self.run_details['data_root'] + '{}/{}/{}/{}/{}_{}.nc'.format(network,self.run_details['ghost_version'],
                                                                                        self.temporal_resolution_vars['temporal_resolution_path'],self.parameter_details['bsc_parameter_name'],self.parameter_details['bsc_parameter_name'],self.yearmonth_str)
            if os.path.isfile(yyyymm_file):
                self.yyyymm_files[network] = yyyymm_file

        #do merge
        self.merge_networks()

        #create successful job log log file
        self.run_details['log_file_str'] += str((time.time() - self.run_details['stage_process_start_time'])/60.)
        create_GHOST_stage_logfile(0,self.run_details)

    def merge_networks(self):
        """Function that merges a yyyymm file across networks.
        
        Any duplicate stations are removed during this process.
        This is done by geographically matching stations within a 
        certain tolerance using longitudes, latitudes and measurement altitudes 
        and then converting to ECEF coordinates.
        If any duplicate matched stations use different methods, they are not 
        classed as duplicates.

        The tolerance is calculated by allowing for a tolerance of 11m in the 
        3 independent x,y,z dimensions, as is done in GHOST to distinguish 
        unique stations.
        Using Pythagoras in 3D √(11**2 +11**2 + 11**2) = 19.053. 
        """ 

        # set tolerance for matching longitudes / latitudes / measurement_altitudes in metres
        tolerance = 19.053

        networks = []
        priorities = []
        longitudes = []
        latitudes = []
        measurement_altitudes = []
        methods = []
        station_indices = []

        for network_ii, (network, yyyymm_file) in enumerate(self.yyyymm_files.items()):
            root = Dataset(yyyymm_file)

            if network_ii == 0:
                networks.extend(np.array([network]*root.dimensions['station'].size))
                priorities.extend(np.array([self.network_priorities[network]]*root.dimensions['station'].size))
                longitudes.extend(root['longitude'][:])
                latitudes.extend(root['latitude'][:])
                measurement_altitudes.extend(root['measurement_altitude'][:])
                methods.extend(root['measurement_methodology'][:])
                station_indices.extend(np.arange(len(longitudes)))

            # compare current network stations with current stations
            else:
                # convert longitudes / latitudes / measurement_altitudes in geographic coordinates to cartesian ECEF 
                # (Earth Centred, Earth Fixed) coordinates assuming WGS84 datum and ellipsoid, and that all heights equal zero
                # ECEF coordiantes represent positions (in metres) as X, Y, Z coordinates, approximating the earth surface as an ellipsoid of revolution
                lla = pyproj.Proj(proj='latlong', ellps='WGS84', datum='WGS84')
                ecef = pyproj.Proj(proj='geocent', ellps='WGS84', datum='WGS84')
                current_x, current_y, current_z = pyproj.transform(lla, ecef, 
                    longitudes, latitudes, measurement_altitudes,
                    radians=False)

                # merge coordinates to 3D array
                current_xyz = np.column_stack((current_x, current_y, current_z))

                # get nextnetwork station information
                nextnetwork_networks = np.array([network]*root.dimensions['station'].size)
                nextnetwork_priorities = np.array([self.network_priorities[network]]*root.dimensions['station'].size)
                nextnetwork_longitudes = root['longitude'][:]
                nextnetwork_latitudes = root['latitude'][:]
                nextnetwork_measurement_altitudes = root['measurement_altitude'][:]
                nextnetwork_methods = root['measurement_methodology'][:]
                nextnetwork_station_indices = np.arange(len(nextnetwork_longitudes))

                # convert next network longitude and latitudes in geographic coordinates to cartesian ECEF 
                nextnetwork_x, nextnetwork_y, nextnetwork_z = pyproj.transform(lla, 
                    ecef, nextnetwork_longitudes, nextnetwork_latitudes, nextnetwork_measurement_altitudes,
                    radians=False)
                
                # merge coordinates to 3D array
                nextnetwork_xyz = np.column_stack((nextnetwork_x, nextnetwork_y, nextnetwork_z))            

                # get all indices of next network xyz coords, within tolerance of each first network xyz coords
                idx = cKDTree(nextnetwork_xyz).query_ball_point(current_xyz, tolerance)

                # drop all within tolerance indices where the next network method is different to current method
                for idx_ii, idx_l in enumerate(idx):
                    if len(idx_l) > 0:
                        idx_l = np.array(idx_l)
                        current_method = methods[idx_ii]
                        next_methods = nextnetwork_methods[idx_l]
                        idx[idx_ii] = list(idx_l[np.where(next_methods == current_method)[0]])

                # get all indices where have non-duplicated and duplicated matched indices
                unique_idx, unique_idx_counts = np.unique(list(itertools.chain(*idx)), return_counts=True)
                nondup_idx = unique_idx[np.where(unique_idx_counts == 1)[0]]
                dup_idx = unique_idx[np.where(unique_idx_counts > 1)[0]]

                # gather list of indices for current network and next network of already resolved within tolerance indices
                cn_wtol_inds = []
                nn_wtol_inds = []

                # resolve all duplicated matched indices
                idx, unresolved_dup_idx, cn_wtol_inds, nn_wtol_inds = self.resolve_duplicate_spatial_colocation_matches(idx, 
                                                                               nondup_idx, dup_idx, 
                                                                               current_xyz, nextnetwork_xyz,
                                                                               cn_wtol_inds, nn_wtol_inds)

                # pass through again to resolve all unresovered duplicated indices after first pass
                if len(unresolved_dup_idx) > 0:
                    idx, _, cn_wtol_inds, nn_wtol_inds = self.resolve_duplicate_spatial_colocation_matches(idx, 
                                                                  nondup_idx, unresolved_dup_idx, 
                                                                  current_xyz, nextnetwork_xyz,
                                                                  cn_wtol_inds, nn_wtol_inds)

                # iterate though next network within tolerance indices, per each current network coord 
                for idx_ii, idx_l in enumerate(idx):
                    # if position has already been resolved (through resolving duplicates) then continue
                    if idx_ii in cn_wtol_inds:
                        continue
                    
                    # no matches, then append nothing
                    elif len(idx_l) == 0:
                        continue

                    # just 1 match, then append
                    elif len(idx_l) == 1:
                        cn_wtol_inds.append(idx_ii) 
                        nn_wtol_inds.append(idx_l[0])

                    # more than 1 match, then find which match has closest distance, then append
                    else:
                        # find the dists between all relevant current network xyz coords and next network duplicate xyz coord
                        idx_l = np.array(idx_l)
                        dists = cKDTree([current_xyz[idx_ii]]).query(nextnetwork_xyz[idx_l], k=1)[0]
                        ordered_idx_l = idx_l[np.argsort(dists)]
                        cn_wtol_inds.append(idx_ii) 
                        nn_wtol_inds.append(ordered_idx_l[0])

                # check if have zero matching stations
                if len(cn_wtol_inds) > 0: 

                    # order matched indices for first species in ascending order, and order next speci indices in smae way
                    cn_wtol_inds, nn_wtol_inds = list(zip(*sorted(zip(cn_wtol_inds, nn_wtol_inds))))

                    # for stations within tolerance,
                    # if new network has higher priority than current network, then overwrite station with new network station
                    # otherwise discard station data
                    for fs_wtol_ind, ns_wtol_ind in zip(cn_wtol_inds, nn_wtol_inds):
                        if priorities[fs_wtol_ind] < nextnetwork_priorities[ns_wtol_ind]:
                            networks[fs_wtol_ind] = nextnetwork_networks[ns_wtol_ind]
                            priorities[fs_wtol_ind] = nextnetwork_priorities[ns_wtol_ind]
                            longitudes[fs_wtol_ind] = nextnetwork_longitudes[ns_wtol_ind]
                            latitudes[fs_wtol_ind] = nextnetwork_latitudes[ns_wtol_ind]
                            measurement_altitudes[fs_wtol_ind] = nextnetwork_measurement_altitudes[ns_wtol_ind]
                            methods[fs_wtol_ind] = nextnetwork_methods[ns_wtol_ind]
                            station_indices[fs_wtol_ind] = nextnetwork_station_indices[ns_wtol_ind]

                # for stations not exceeding tolerance, simply append stations data to current ones
                nn_otol_inds = np.setdiff1d(np.arange(nextnetwork_longitudes.shape[0]), np.array(nn_wtol_inds))
                networks.extend(nextnetwork_networks[nn_otol_inds])
                priorities.extend(nextnetwork_priorities[nn_otol_inds])
                longitudes.extend(nextnetwork_longitudes[nn_otol_inds])
                latitudes.extend(nextnetwork_latitudes[nn_otol_inds])
                measurement_altitudes.extend(nextnetwork_measurement_altitudes[nn_otol_inds])
                methods.extend(nextnetwork_methods[nn_otol_inds])
                station_indices.extend(nextnetwork_station_indices[nn_otol_inds])

        # now merge stations to a netCDF for yyyymm
        if len(station_indices) > 0:
            self.write_to_netCDF(np.array(networks), np.array(station_indices))

        return

    def resolve_duplicate_spatial_colocation_matches(self, idx, nondup_idx, dup_idx, 
                                                     current_xyz, nextnetwork_xyz,
                                                     cn_wtol_inds, nn_wtol_inds):

        """ Function that resolves duplicate indices found during spatial colocation of 2 networks.

            In some cases the stations within the tolerance can match for multiple stations. 
            In order to resolve this, for each duplicated station, it is set to match with the closest station to it
            in terms of 3D distance. This is done iteratively until there are no more duplicates.

            :param idx: per current xyz coords, a list of indices of next network xyz coords within tolerance
            :type idx: array 
            :param nondup_idx: next network idx coords that are not duplicated across idx array 
            :type nondup_idx: array
            :param dup_idx: next network idx coords that are duplicated across idx array 
            :type dup_idx: array
            :param current_xyz: ECEF coordinates for current stations
            :type current_xyz: array
            :param nextnetwork_xyz: ECEF coordinates for next network stations
            :type nextnetwork_xyz: array
            :param cn_wtol_inds: current station indices within tolerance (i.e. have paired match)
            :type cn_wtol_inds: list
            :param nn_wtol_inds: next network station indices within tolerance (i.e. have paired match)
            :type nn_wtol_inds: list
            :return: idx, unresolved_dup_idx, cn_wtol_inds, nn_wtol_inds 
            :rtype: array, list, list, list
        """

        # resolve all duplicated matched indices by finding for which index the distance is closest
        unresolved_dup_idx = []
        for dup_index in dup_idx:

            # get all relevant first speci indices for which contain the duplicate next speci index
            relevant_idx = np.array([idx_ii for idx_ii, idx_l in enumerate(idx) if dup_index in idx_l])
            # find the dists between all relevant first speci xyz coords and next speci duplicate xyz coord
            dists = cKDTree([nextnetwork_xyz[dup_index]]).query(current_xyz[relevant_idx], k=1)[0]
            # order the relevant first speci indices by dists (closest first)
            ordered_relevant_idx = relevant_idx[np.argsort(dists)]
            # iterate through ordered relevant first speci indices until find an ind for which can claim duplicate ind
            # keep going in order iteratively until have exhausted all options for duplicate ind
            for ordered_relevant_index_ii, ordered_relevant_index in enumerate(ordered_relevant_idx):
            
                # if current relevant first speci index has no competing indices, then append matched index,
                # remove duplicate index from other match lists, and then break out of iteration
                if len(idx[ordered_relevant_index]) == 1:
                    cn_wtol_inds.append(ordered_relevant_index)
                    nn_wtol_inds.append(dup_index)
                    for next_ordered_relevant_index in ordered_relevant_idx[ordered_relevant_index_ii+1:]:
                        idx[next_ordered_relevant_index].remove(dup_index)
                    break

                # otherwise, 
                # find the dists between the relevant first speci xyz coord and next speci duplicate xyz coords
                else:
                    idx_l = np.array(idx[ordered_relevant_index])
                    dists = cKDTree([current_xyz[ordered_relevant_index]]).query(nextnetwork_xyz[idx_l], k=1)[0]
                    ordered_idx_l = idx_l[np.argsort(dists)]

                    # if the closest index is the dup index, or is another non-duplicate, then append it
                    # if the closest index is the dup index, then remove duplicate index from other match lists 
                    # and then break out of iteration
                    if (ordered_idx_l[0] == dup_index) or (ordered_idx_l[0] in nondup_idx):
                        cn_wtol_inds.append(ordered_relevant_index)
                        nn_wtol_inds.append(ordered_idx_l[0])
                        if ordered_idx_l[0] == dup_index:
                            for next_ordered_relevant_index in ordered_relevant_idx[ordered_relevant_index_ii+1:]:
                                idx[next_ordered_relevant_index].remove(dup_index)
                            break

                    # else, if the closest index is another duplicated index,
                    # then cannot currently resolve position, so come back to this later
                    else:
                        unresolved_dup_idx.append(dup_index)
                        break

        return idx, unresolved_dup_idx, cn_wtol_inds, nn_wtol_inds 


    def write_to_netCDF(self, networks, station_indices):
        """write a merged netCDF file for a specific resolution / speci / yyyymm,
        for all networks
        """

        # set merged filename
        merged_filename = self.run_details['data_root'] + '{}/{}/{}/{}/{}_{}.nc'.format(self.run_details['network'],self.run_details['ghost_version'],
                                                                                        self.temporal_resolution_vars['temporal_resolution_path'],self.parameter_details['bsc_parameter_name'],self.parameter_details['bsc_parameter_name'],self.yearmonth_str)

        #create file (and dirs)
        check_directory_existence('/'.join(merged_filename.split('/')[:-1])+'/', directories_not_to_test='{}{}/'.format(self.run_details['data_root'],self.run_details['network']))
        
        # start new netCDF file
        merged_root = Dataset(merged_filename, 'w', format="NETCDF4")

        # get all unique networks giving data
        unique_networks = np.unique(networks)

        # get year and month
        year = str(self.yearmonth_str[:4])
        month = str(self.yearmonth_str[4:])

        # iterate through unique networks
        station_start_n = 0 
        for unique_network_ii, unique_network in enumerate(unique_networks):

            # get relevant indices for stations in network
            relevant_station_inds = np.where(networks == unique_network)[0]

            # get indices of relevant stations inside native network netCDF
            network_station_indices = station_indices[relevant_station_inds]

            # load network netCDF
            root = Dataset(self.run_details['data_root'] + '{}/{}/{}/{}/{}_{}.nc'.format(unique_network,self.run_details['ghost_version'],
                                                                                         self.temporal_resolution_vars['temporal_resolution_path'],self.parameter_details['bsc_parameter_name'],self.parameter_details['bsc_parameter_name'],self.yearmonth_str))

            #copy global attributes (on first pass)
            if unique_network_ii == 0:
                for name in root.ncattrs():
                    if name == 'title':
                        if self.run_details['network'] == 'GHOST-PUBLIC':
                            merged_root.title = 'Merged {} data, across all public networks in the GHOST dataset, in {}-{}.'.format(standard_parameters[self.run_details['parameter']]['long_parameter_name'], year, month)
                        else:
                            merged_root.title = 'Merged {} data, across all networks in the GHOST dataset, in {}-{}.'.format(standard_parameters[self.run_details['parameter']]['long_parameter_name'], year, month)
                    else:
                        merged_root.setncattr(name, root.getncattr(name))
            
                #copy dimensions (altering station dimension) (on first pass)
                for name, dimension in root.dimensions.items():
                    if name == 'station':
                        merged_root.createDimension(name, len(networks))
                    else:
                        merged_root.createDimension(name, len(dimension))

            #copy variables
            for name, variable in root.variables.items():

                #get current var dtype
                current_dtype = variable.dtype

                #create variable (on first pass)
                if unique_network_ii == 0:
                    if '_FillValue' in root[name].ncattrs():
                        var_root = merged_root.createVariable(name, variable.datatype, variable.dimensions, fill_value=255, zlib=True, complevel=1, shuffle=True)
                    else:
                        var_root = merged_root.createVariable(name, variable.datatype, variable.dimensions, zlib=True, complevel=1, shuffle=True)

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

                #copy and set 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 network_station_indices], dtype=object)
                        else:
                            data = root[name][network_station_indices]
                        merged_root[name][station_start_n:station_start_n+len(network_station_indices)] = data
                    elif len(root[name].shape) == 2:  
                        if current_dtype == str:
                            data = np.array([root[name][ind,:] for ind in network_station_indices], dtype=object)
                        else:
                            data = root[name][network_station_indices,:]
                        merged_root[name][station_start_n:station_start_n+len(network_station_indices),:] = data
                    elif len(root[name].shape) == 3:
                        if current_dtype == str:
                            data = np.array([root[name][ind,:,:] for ind in network_station_indices], dtype=object)
                        else:
                            data = root[name][network_station_indices,:,:]
                        merged_root[name][station_start_n:station_start_n+len(network_station_indices),:,:] = data
                # only write time variable on first pass
                elif unique_network_ii == 0:
                    merged_root[name][:] = root[name][:]

            #iterate station start n
            station_start_n += len(network_station_indices)

            #close network netCDF
            root.close()
            
        #close new netCDF
        merged_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', merged_filename, merged_filename]
        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(merged_filename)   