#!/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)
#       Ivan Jose Cernicharo Ortiz (ivan.cernicharo@bsc.es)

import sys
import logging
import os
import subprocess
import xarray
import numpy as np

# Local import
from utils.utils_files import *
from utils.utils_xarray import *

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

    def __init__(self, logdir, files, varname, dt, cz, mask_file, mask_var, mask_dim):
        self.testid = 'test02-spatial-completeness'
        self.logger = logging.getLogger(self.testid)
        self.logdir = logdir                   
        self.varname = varname
        self.datatype = dt
        self.chunksize = cz 
        self.files = files       
       
        # Mask (only for ocean products)
        self.mask_file = mask_file
        self.mask_var = mask_var
        self.mask = os.path.exists(mask_file)        
        self.mask_dim = mask_dim.split(',')

    """
    Loading DATA | Check for missing values
    """
    def missing_values(self, filepath):
        
        # Read data into Xarray
        ds = open_dataset(filepath, self.datatype, c_chunks=self.chunksize)
        arr = get_variable(ds, self.varname)
        mv = 0
        uv = 0
        tot = 0
        
        # Check if there is an ocean mask
        if self.mask:
            mask = open_dataset(self.mask_file, '', c_chunks=self.chunksize)
            mask_arr = get_mask(mask, self.mask_var)
   
            try:
                mv, uv, tot = find_missing_unnecessary_values(ds, arr, mask_arr, self.mask_dim, self.varname)

                # Try to avoid memory issues (force delete)
                del mask_arr
                del mask
                del ds
                del arr
            except ValueError:
                raise Exception('The mask provided does not fit the data')
           
        # Remove NaN values + flatten        
        return mv, uv, tot

    """
    Main function
    """
    def check_spatial_completeness(self):
       
        # Progress init
        i = 0
        n = len(self.files) # read/calc_stats

        # Gather statistics per file
        mv = 0.0
        uv = 0.0
        msg = ''

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

            # Get number of missing values
            mv, uv, tot = self.missing_values(f)            
            if mv != 0 or uv != 0:                                              
                self.logger.error('File {} has {} missing values and {} unnecessary values\n'.format(f, mv, uv))

        # Final message
        if mv == 0:            
            msg = 'Files do not contain missing values (NaN)'
            self.logger.info(msg)
            return 'ok', msg
        else:            
            msg = 'Files contain missing values (NaN). Check log for details'
            self.logger.error(msg)
            return 'err', msg

        return res

    def main(self):
       
        # Do the actual work
        res, msg = self.check_spatial_completeness()

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