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

"""
Loading dataset (NetCDF)
Tries to load data, return exception message otherwise
"""
def load_datafile(filepath):
    # Read data into Xarray
    try:
        ds = xarray.open_dataset(filepath)
    except Exception as e:
        return None, str(e)

    return ds, ''

"""
Read a variable from a given dataset (NetCDF)
"""
def get_variable(dataset, variable):
    arr = None
    try:        
        arr = dataset[variable].data.flatten()
    except: 
        raise Exception('Variable {} not present in dataset, please check your configuration (available vars: {})'.format(
            variable, dataset.data_vars)) 

    return arr[~np.isnan(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, is_land):

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

    # Case 1 - Missing value (data should be there)
    # Case 2 - Unnecessary value (data should not be there)
    # Mask is %of water (0 is land for sure and 1 is water for sure)
    if is_land:
        mv = np.count_nonzero(nan_values[np.where(mask == 1.0)])
        uv = np.count_nonzero(~nan_values[np.where(mask == 0.0)])
    else:
        mv = np.count_nonzero(nan_values[np.where(mask == 0.0)])
        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, is_land):
    if not(mask is None):
        mv, uv = compare_matrix(mask, data, is_land)
    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, is_land):
    
    # 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, is_land)
                # 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, is_land)
            # Update counter
            mv+=mvd
            uv+=uvd
            
    ## 2D - Dimensions (x,y)
    if len(dimensions) == 0:
        mvd,uvd = count_mv_uv(mask2d, ds[varname].values, is_land)
        # Update counter
        mv+=mvd
        uv+=uvd

    return mv, uv