#!/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)
#       Alejandro Alfonso Hernandez  (alejandro.alfonso@bsc.es)

import sys
import xarray

# Local import
from dqc_wrapper.utils.utils_files import *
from dqc_wrapper.utils.utils_grib import *

"""
Class: Spatial completeness checks
---
- Checks for data gaps or missing values
"""
class SpatialCompletenessTest:

    def __init__(self, files, varname, dt, cds_dataset, system, origin, masks_dir, lsm, lnd):
        self.testid = 'test02-spatial_completeness'
        self.logger = logging.getLogger(self.testid)                  
        self.varname = varname
        self.datatype = dt
        self.files = files
        self.land_sea_mask = (lsm.strip() == 'yes')
        self.land_variable = (lnd.strip() == 'yes')

        # Mask pattern to find mask files in function of their resolution (depends on dataset)
        if 'seasonal' in cds_dataset:
            self.mask_pattern = os.path.join(masks_dir, '{}-{}-{}-land_sea_mask.grib'.format(cds_dataset, origin, system))
        else:
            if 'uerra' in cds_dataset:
                # For UERRA the origin is needed (single and soil levels)
                self.mask_pattern = os.path.join(masks_dir, '{}-{}-land_sea_mask.grib'.format(cds_dataset, origin))
            else:
                self.mask_pattern = os.path.join(masks_dir, '{}-*-land_sea_mask.grib'.format(cds_dataset))

        if self.land_sea_mask:
            self.logger.info('Mask pattern: {}'.format(self.mask_pattern))

        # Mask depends on the given dataset (only for ocean and/or land products)
        self.cds_dataset = cds_dataset

    """
    Loading DATA | Check for missing values
    Returns -1 if data problems, alongisde with an error message
    """
    def missing_values_xarray(self, filepath, mask_files):

        # Read data into Xarray
        msg = ''
        try:
            ds = xarray.open_dataset(filepath)
        except Exception as e:
            return -1, -1, str(e)

        # Check if there is an ocean mask
        try:
            if self.land_sea_mask:
                max_nbp = max(mask_files.keys())  # Get highest resolution mask
                mv, uv = find_missing_unnecessary_values(ds, mask_files[max_nbp], self.varname, self.land_variable)
            else:
                mv, uv = find_missing_unnecessary_values(ds, None, self.varname, self.land_variable)
        except Exception as e:
            return -1, -1, str(e)

        # Remove NaN values + flatten        
        return mv, uv, msg

    """
    Read masks (if needed) - Return an dictionary of all possible resolutions (useful for datatypes)
    """
    def read_masks(self):
        mask_files = {}

        # Only load stuff when needed
        if self.land_sea_mask:

            # Get all possible masks resolutions
            mask_paths = glob.glob(self.mask_pattern)
            try:
                for f in mask_paths:
                    gc = GribFileCheck(f)
                    mask = gc.get_grib_first_values(without_Nan=False)
                    mask_files[len(mask)] = mask
                    self.logger.info('INFO: {} loaded'.format(os.path.basename(f)))
            except Exception as e:
                return 'err', str(e)

            # Raise exception if mask was needed and could not find any
            if mask_files == {}:
                raise Exception('Could not find proper mask files ({})'.format(self.mask_pattern))

        return mask_files

    """
    General loop per file
    """
    def check_spatial_completeness(self):

        if self.varname == 'cbh':
            # Spatial completeness not applicable for variable cloud_base_height (cbh)
            res = 'ok'
            msg = 'Spatial completeness not applicable for this variable, as missing data are expected where the (total) cloud cover is zero'

        else:

            # Progress init
            i = 0
            n = len(self.files) # read/calc_stats

            # Gather statistics per file
            mv = 0
            uv = 0
            errors_mv = ''
            errors_uv = ''

            # Read mask (if any)
            mask_files = self.read_masks()

            # Loop over files (grib and netcdf follow different strategies)
            for f in self.files:
                # Progress update               
                i += 1
                print_progress(i, n)

                # Get number of missing values
                if f.endswith('.grib'):
                    # ECCODES / GRIB
                    gc = GribFileCheck(f)
                    try:
                        mvd, uvd = gc.get_grib_mv(mask_files, self.land_variable)
                    except Exception as e:
                        mvd = -1
                        msg = 'A proper mask is not available for this variable'
                else:
                    # XARRAY / NETCDF
                    mvd, uvd, msg = self.missing_values_xarray(f, mask_files)

                # Error occurred exit
                if mvd == -1:
                    self.logger.error(msg)
                    return 'err', msg

                # Collect stats/errors
                if mvd != 0:                                              
                    errors_mv = errors_mv + '{} ({} mv)'.format(os.path.basename(f), mvd) + ', '
                if uvd != 0:
                    errors_uv = errors_uv + '{} ({} uv)'.format(os.path.basename(f), uvd) + ', '
                
                # Update general counters
                mv += mvd
                uv += uvd

            # Dump errors
            if errors_uv != '':
                res = 'warn'
                msg = 'Files contain unexpected values. Check log messages in the following pages'
                self.logger.warning('Files with unnecessary values (val where NaN should be) - [uv]:')
                self.logger.warning(' ' + errors_uv)
            if errors_mv != '':
                res = 'err'
                msg = 'Files contain missing values. Check log messages in the following pages'
                self.logger.error('Files with missing values (NaN where val should be) - [mv]:')
                self.logger.error(' ' + errors_mv)

            # Final message
            if mv == 0 and uv == 0:
                res = 'ok'
                msg = 'Files do not contain missing values'
                if self.land_sea_mask:
                    msg += ' and match the land/sea mask'
                    
        return res, msg
        
    """
    Main function
    """
    def main(self):
       
        # Do the actual work
        res, msg = self.check_spatial_completeness()

        if res == 'ok':
            self.logger.info(msg)
        else:
            self.logger.error(msg)

        # Prepare output
        return { 
            'res': res, 
            'msg': msg
        }
    
if __name__ == '__main__':
    sys.exit()

