#!/usr/bin/env python
#
# Data Quality Checker also known as DQC
#
# Read the README.md for more detailed information
#
#       Barcelona Supercomputing Center - Centro Nacional de Supercomputacion
#       Earth Sciences Department - Computational Earth Sciences
#     
#     Authors: 
#       Joan Sala Calero (joan.sala@bsc.es)

import itertools
import logging
import xarray
import numpy as np


"""
Get filter_by_keys dictionary for xarray/cfgrib
- Data type
- Step type
- Number of data points
"""
def get_filter_keys(dt, st, nbp):
    # Empty filter
    filt = {'filter_by_keys': {}}

    # Number of points (resolution)
    if nbp != '':
        filt['filter_by_keys']['numberOfPoints'] = nbp

    # Data type (cds product type)
    if dt != '':
        filt['filter_by_keys']['dataType'] = dt

    # Step type for accumulated variables
    if st != '':
        filt['filter_by_keys']['stepType'] = st

    return filt

"""
Open a gridded dataset
- Return an array of dataset objects
"""
def open_dataset(filepath, datatypes=[], steptypes=[], nbpoints=[]):
    # Init
    ds = []

    # NetCDF:
    # - Read data directly into Xarray
    if filepath.endswith('.nc'):
        ds = [ xarray.open_dataset(filepath) ]

    # GRIB:
    #  - Use cfgrib package in combination with Xarray (optional: filter by datatype and resolution nbpoints)
    if filepath.endswith('.grib') or filepath.endswith('.grb'):

        # Optional: Step type (avgas, avgad)
        if len(steptypes) > 1:
            st = steptypes[0]
        else:
            st = ''

        if datatypes == [] and nbpoints == []:
            try:
                ds = [ xarray.open_dataset(filepath, engine='cfgrib') ]
            except Exception as e:
                logging.error('{}'.format(e))
        else:
            # - Filter keys is sometimes necessary to distinguish messages (numberOfPoints + dataTypes)
            # - Higher resolution has preference
            for nbp, dt in itertools.product(sorted(nbpoints, reverse=True), datatypes):

                try:
                    filt = get_filter_keys(dt, st, nbp)
                    ds.append(xarray.open_dataset(filepath, engine='cfgrib', backend_kwargs=filt))
                    logging.info('opened dataset with dt={}, nbp={}'.format(dt, nbp))
                except Exception as e:
                    # Skip error: Let's try another combination
                    pass

    # For debugging purposes
    if not (valid_dataset(ds)):
        raise Exception('Could not open dataset {} with dataTypes={}'.format(filepath, datatypes))

    return ds

"""
Check if dataset has been opened properly
"""
def valid_dataset(ds):
    # All datasets opened
    if len(ds) > 0:
        valid = True
        for d in ds:
            valid = valid and (type(d) == xarray.core.dataset.Dataset)
    else:
        valid = False
    return valid

"""
Read a variable from a given dataset
"""
def get_variable(dataset, variable):

    arr = None
    try:        
        arr = dataset[variable].data
    except: 
        raise Exception('Variable {} not present in dataset, please check your configuration (available vars: {})'.format(
            variable, dataset.data_vars)) 

    return arr

"""
Read a mask (2d coordinates), given the other dimensions
"""
def get_mask(dataset, variable):
 
    arr = None
    try:
        if len(dataset[variable].shape) == 2:
            arr2d = dataset[variable].data
        elif len(dataset[variable].shape) == 3:
            arr2d = dataset[variable].data[0,:,:]
        else:
            raise Exception('Mask file provided is not well-formed'.format(variable))     
    except KeyError:        
        raise Exception('Variable {} not present in dataset, please check your configuration'.format(variable)) 
    
    return arr2d

"""
Get mask dimensions (except coordinates)
"""
def get_dimensions(dataset):    
    
    dims = []
    coord_names = [ 'latitude', 'lat', 'longitude', 'lon', 'x', 'y']

    for d in dataset.dims:      
        if not(d in coord_names):
            dims.append(d)
    
    return dims

"""
Get latitude name for a given dataset
"""
def get_lat_name(dataset):       
   
    coord_names = [ 'latitude', 'lat', 'y']
    for d in dataset.coords:      
        if (d in coord_names):
            return d
    
    return None

"""
Check spatial completeness with a land/sea mask (percentage of water)
"""
def compare_matrix(mask, data):

    # True where v is NaN
    nan_values = np.isnan(data)

    # Case 1 - SEA - Missing value (data should be there) - isNan=True / isMask=False (sea is 0)    
    mv = np.count_nonzero(nan_values[np.where(mask == 0.0)])

    # Case 2 - LAND - Unnecessary value (data should not be there) - isNan=False / isMask=True (land is 0% water)    
    uv = np.count_nonzero(~nan_values[np.where(mask == 1.0)])
    
    return mv, uv

"""
Count missing_values (mv) and unnecessary_values (uv)
- Compare with land/sea mask
- If no mask, just count NaN
"""
def count_mv_uv(mask, data):
    if not(mask is None):
        mv, uv = compare_matrix(mask, data) 
    else:
        uv = 0
        mv = np.count_nonzero(np.isnan(data))

    return mv, uv

"""
Mask a given matrix by array taking into account different dimensions (4D,3D,2D)
"""
def find_missing_unnecessary_values(ds, mask2d, varname):
    
    # Binary 0|1 mask - Sea is 0 (everything else >0)  
    # If no mask, no action
    if not(mask2d is None):        
        if isinstance(mask2d, np.ndarray):
            mask2d = mask2d 
        else:
            mask2d = mask2d.compute()

    # Get dataset dimensions (except XY)
    dimensions = get_dimensions(ds)

    # Counters (missing values and unnecessary values) 
    # Loop over dimensions (apologies for the code)   
    mv = 0
    uv = 0    

    # 4D - Dimensions (d1,d2,x,y)
    if len(dimensions) == 2:
        for _, dd1 in ds[varname].groupby(dimensions[0]):           
            for _, dd2 in dd1.groupby(dimensions[1]):
                mvd,uvd = count_mv_uv(mask2d, dd2.values)
                # Update counter
                mv+=mvd
                uv+=uvd

    ## 3D - Dimensions (d1,x,y)    
    if len(dimensions) == 1:
        for _, dd1 in ds[varname].groupby(dimensions[0]):
            mvd,uvd = count_mv_uv(mask2d, dd1.values)
            # Update counter
            mv+=mvd
            uv+=uvd
            
    ## 2D - Dimensions (x,y)    
    if len(dimensions) == 0:
        mvd,uvd = count_mv_uv(mask2d, ds[varname].values)
        # Update counter
        mv+=mvd
        uv+=uvd

    return mv, uv
