#!/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 xarray
import numpy as np

"""
Open a gridded dataset
"""
def open_dataset(filepath, datatype, c_chunks=-1):    
    ds = None

    # NetCDF: Read data into Xarray
    if filepath.endswith('.nc'):
        if c_chunks > 0:    
            ds = xarray.open_dataset(filepath, chunks=c_chunks)
        else:
            ds = xarray.open_dataset(filepath)

    # GRIB: Use cfgrib package in combination with Xarray (optional: filter by datatype)
    if filepath.endswith('.grib'):            
        if datatype == '':
            if c_chunks > 0:
                ds = xarray.open_dataset(filepath, engine='cfgrib', chunks=c_chunks)
            else:
                ds = xarray.open_dataset(filepath, engine='cfgrib')
        else:
            if c_chunks > 0:
                ds = xarray.open_dataset(filepath, engine='cfgrib', chunks=c_chunks, 
                    backend_kwargs={'filter_by_keys': {'dataType': datatype}})
            else:
                ds = xarray.open_dataset(filepath, engine='cfgrib', 
                    backend_kwargs={'filter_by_keys': {'dataType': datatype}})                
    return ds

"""
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


def compare_matrix(mask, data):

    # True where v is NaN
    nan_values = np.isnan(data)                
    # Case 1 - Missing value (data should be there) - isNan=True / isMask=True
    mv = np.count_nonzero(mask * nan_values)
    # Case 2 - Unnecessary value (data should not be there) - isNan=False / isMask=False
    uv = np.count_nonzero(~mask * ~nan_values)

    return mv, uv

"""
Mask a given matrix by array taking into account different dimensions (4D or 3D)
"""
def find_missing_unnecessary_values(ds, arr, mask2d, dimensions, varname):
    
    # Binary 0|1 mask 
    mask = mask2d.compute().astype(bool)

    # Counters (missing values and unnecessary values)
    mv = 0
    uv = 0

    for _, dd1 in ds[varname].groupby(dimensions[0]):
        
        if len(dimensions) == 2:
            for _, dd2 in dd1.groupby(dimensions[1]):
                mvd, uvd = compare_matrix(mask, dd2.values)
                # Update counters
                mv+=mvd
                uv+=uvd                
        else:
            mvd, uvd = compare_matrix(mask, dd1.values)
            # Update counters
            mv+=mvd
            uv+=uvd

    return mv, uv, mask.shape[0]*mask.shape[1]
   