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

# Local import
from dqc_wrapper.utils.utils_files import *
from dqc_wrapper.utils.utils_xarray 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, mask_file, mask_var):
        self.testid = 'test02-spatial_completeness'
        self.logger = logging.getLogger(self.testid)                  
        self.varname = varname
        self.datatype = dt
        self.files = files
       
        # Mask (only for ocean and/or land products)
        self.mask_file = mask_file
        self.mask_var = mask_var
        self.mask = (mask_file.strip() != "")

        # Get file characteristics (resolution and step) - if GRIB <-> filtering needed (we get the highest resolution)
        self.nbp = []
        self.stp = []
        if files[0].endswith('.grib') or files[0].endswith('.grb'):
            self.nbp = grib_get_max_resolution(files[0])
            self.stp = grib_get_stepType(files[0])

    """
    Loop over datatypes, return missing_values/unnecessary_values + error message if any
    """
    def loop_over_datasets(self, ds, mask_arr):
        mv = 0
        uv = 0
        for d in ds:
            try:
                mvt, uvt = find_missing_unnecessary_values(d, mask_arr, self.varname)
                mv += mvt
                uv += uvt
            except Exception as e:
                return -1, -1, str(e)

        return mv, uv, ''

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

        # Read data into Xarray
        msg = ''
        try:            
            ds = open_dataset(filepath, datatypes=self.datatype, nbpoints=self.nbp, steptypes=self.stp)
        except Exception as e:
            return -1, -1, str(e)

        # Check if there is an ocean mask
        if self.mask:

            # Read mask
            try:
                mask = open_dataset(self.mask_file)[0] # Only one datatype + resolution
                mask_arr = get_mask(mask, self.mask_var)
            except Exception as e:
                return -1, -1, str(e)
   
            try:
                # Given two arrays (data and mask) return missing/unnecessary values (for every cds product type)
                mv, uv, msg = self.loop_over_datasets(ds, mask_arr)

                # Try to avoid memory issues (force delete)
                del mask_arr
                del mask
                del ds
            except ValueError:
                raise Exception('The mask provided does not fit the data')

        # No mask, all values should be filled
        else:
            mv, uv, msg = self.loop_over_datasets(ds, None)

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

    """
    General loop per file
    """
    def check_spatial_completeness(self):
       
        # Progress init
        i = 0
        n = len(self.files) # read/calc_stats

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

        for f in self.files:
            # Progress update               
            i += 1
            print_progress(i, n)                        

            # Get number of missing values
            mvd, uvd, msg = self.missing_values(f)

            # 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_mv != '':
            self.logger.error('Files with missing values (NaN where val should be) - [mv]:\n')
            self.logger.error(' ' + errors_mv)
        if errors_uv != '':
            self.logger.error('Files with unnecessary values (val where NaN should be) - [uv]:\n')
            self.logger.error(' ' + errors_uv)

        # Final message
        if mv == 0 and uv == 0:            
            msg = 'Files do not contain missing values'
            if self.mask:    
                msg += ' and match the land/sea mask'      
            return 'ok', msg
        else:            
            msg = 'Files contain missing values. Check log messages in the following pages'
            return 'err', 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__':
    #files = dir_or_file('/data/dqc_test_data/reanalysis/reanalysis-era5-single-levels/sea_surface_temperature', '*.grib', '', '', '')
    #sp = SpatialCompletenessTest(files['original'],'sst', ['em', 'an'], '', '')
    #print(sp.check_spatial_completeness())
    sys.exit()

