#!/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

from dqc_wrapper.utils.utils_grib import *

"""
Open a gridded dataset
We avoid chunks on GRIB files (not efficient)
"""
def open_dataset(filepath, datatypes=[], c_chunks=-1):
    
    # Init
    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 and resolution nbpoints)
    # - Use of chunking in xarray is not encouraged. Weird behaviour
    if filepath.endswith('.grib') or filepath.endswith('.grb'):

        # Get maximum resolution within files (to avoid errors opening with cfgrib)
        nbpoints = grib_get_resolutions(filepath)

        # Get forecast stepType if available
        steptypes = grib_get_stepType(filepath)
        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):
                filt = { 'filter_by_keys': {} }
                filt['filter_by_keys']['numberOfPoints'] = nbp

                if dt != '':
                    filt['filter_by_keys']['dataType'] = dt
                if st != '':
                    filt['filter_by_keys']['stepType'] = st

                try:         
                    ds = xarray.open_dataset(filepath, engine='cfgrib', backend_kwargs=filt) 
                    if valid_dataset(ds):
                        break # We managed to open the dataset               
                
                except Exception as e:
                    # ERR: Let's try another combination
                    logging.error('Filter keys did not work {}'.format(filt))
                    logging.error('{}'.format(e))                    
                    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):
    return type(ds) == xarray.core.dataset.Dataset

"""
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 or 3D)
"""
def find_missing_unnecessary_values(ds, arr, 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
    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
   