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

import sys
import logging
import os
import re
import subprocess
import numpy as np
import xarray
import cfgrib
import numpy as np
import pandas as pd
from scipy import stats

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

"""
Class: Valid ranges checks
---
- Checks data to detect outliers
"""
class valid_ranges_test:

    def __init__(self, logdir, files, varname, vmin, vmax, dt, cz, kstd=10):
        self.testid = 'test06-valid_ranges'
        self.logger = logging.getLogger(self.testid)
        self.logdir = logdir 
        self.files = files
        self.varname = varname     
        self.datatype = dt   
        self.chunksize = cz
        self.stat_columns = ['Min', 'Max', 'Mean', 'Kurt', 'Skew', 'Std']
        
        # Number of standard deviations
        self.kstd = kstd
        
        # Automatic versus computed valid ranges
        try:
            self.valid_min = float(vmin)
        except:
            self.valid_min = None
        try:
            self.valid_max = float(vmax)
        except:
            self.valid_max = None
    
    """
    Producing a list of statistical estimators
    """
    def get_stats(self, data):
        
        # Data to be a flatten numpy array
        if isinstance(data, np.ndarray):
            vals = data
        else:
            vals = data.compute()

        Min = np.min(vals)
        Max = np.max(vals)
        Mean = np.mean(vals)
        Std = np.std(vals)

        # Kurtosis and Skewness
        Kurt = stats.kurtosis(vals)
        Skew = stats.skew(vals)

        return [Min, Max, Mean, Kurt, Skew, Std]

    """
    Loading DATA depending on fileformat | All Timesteps | Filter out NaN
    """
    def load_datafile(self, filepath):

        # Read data into Xarray
        ds = open_dataset(filepath, self.datatype, c_chunks=self.chunksize)        
        arr = get_variable(ds, self.varname)
        
        # Remove NaN values + flatten
        arr = arr[~np.isnan(arr)]
        
        # Try to avoid memory leak
        del ds
        
        return arr

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

        # Store stats information per file       
        df = pd.DataFrame(index=self.files, columns=self.stat_columns)

        # In a CSV file    
        csv_fname = re.sub("([0-9]{8})", "stats.csv", os.path.splitext(self.files[0])[0]) # YYYYMMDD
        self.logger.info('Saving stats in {}'.format(csv_fname))
        print('INFO: Saving stats in {}'.format(csv_fname))
        print()
        
        # Gather statistics per file
        for f in self.files:
            # Progress update            
            i+=1
            print_progress(i, n)              
            
            # Load data from file
            data = self.load_datafile(f)

            # Gather statistics            
            df.loc[f] = self.get_stats(data)

            # Try to avoid memory leak
            del data

        # Save stats for future use
        df.index.name = 'File'
        df.to_csv(csv_fname)

        # Valid range estimation (averaged since we cannot load all data at once)               
        if self.valid_min is None:
            self.valid_min = df['Mean'].mean() - self.kstd * df['Std'].mean()
        if self.valid_max is None:
            self.valid_max = df['Mean'].mean() + self.kstd * df['Std'].mean()

        # Check if any file is outside the limits
        res = True
        for fname, row in df.iterrows():
            # Progress update            
            i+=1            
            print_progress(i, n)   

            # Check ranges          
            if row['Min'] < self.valid_min or row['Max'] > self.valid_max:
                self.logger.error('{} contains outliers'.format(fname))
                self.logger.error('[ {} > {} ] vs [ {} > {} ]'.format(self.valid_min, self.valid_max, row['Min'], row['Max']))
                res = False                

        # Final message
        if res:
            msg = 'Files are inside valid ranges [ from {} to  {} ]'.format(self.valid_min, self.valid_max)
            self.logger.info(msg)
            return 'ok', msg
        else:
            msg = 'Files contain unrealistic values outside of valid ranges [ from {} to  {} ]'.format(self.valid_min, self.valid_max)
            self.logger.error(msg)
            return 'err', msg

        return res

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

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