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

import sys

import warnings
warnings.filterwarnings('ignore')

# Local imports
from dqc_wrapper.utils.utils_forms import  *
from dqc_wrapper.utils.utils_xarray import *
from dqc_wrapper.utils.utils_bokeh import *
from dqc_wrapper.utils.utils_files import *
from dqc_wrapper.utils.utils_grib import *
from dqc_wrapper.utils.utils_stats import *

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

    def __init__(self, resdir, testname, files, varname, vmin, vmax, dt, cz, forms_dir, cds_dataset, cds_variable, kstd=10):
        self.testid = 'test06-valid_ranges'
        self.logger = logging.getLogger(self.testid)
        self.output_plot_html = os.path.join(resdir, testname.replace('.ini', '-valid_ranges.html'))
        self.files = files
        self.varname = varname     
        self.datatype = dt
        self.chunksize = cz
        self.stat_columns = [ 'Min', 'Max', 'Mean', 'Std', 'Med', 'q1', 'q3', 'IQR', 'skew' ]
        self.cds_variable = cds_variable
        self.cds_dataset = cds_dataset
        self.metadata = get_metadata(forms_dir, cds_dataset, cds_variable)

        # 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

    """
    Loading dataset | All Timesteps | Filter out NaN
    Tries to load data, return exception message otherwise
    """
    def load_datafile(self, filepath):

        # Read data into Xarray
        try:
            ds = open_dataset(filepath, datatypes=self.datatype)
            arr = get_variable(ds, self.varname)
        except Exception as e:
            return None, None, str(e)
        
        # Remove NaN values
        arr = arr[~np.isnan(arr)]
        
        return ds, arr, ''

    """
    Hadley cell stats (90, 60, 30) absolute latitudes (polar, mid-latitude, hadley)    
    """
    def get_hist_hadley(self, df_hist, hbins, ds, init, lat0, lat1, row, latid):
        arr = ds[self.varname].where((ds[latid] > lat0) & (ds[latid] <= lat1)).values       
        x, b = np.histogram(arr, bins=hbins)
        
        if init==0:
            df_hist.loc[row] = x
        else:
            df_hist.loc[row] += x

    """
    Get Histogram bins initalization. Number of bins + K = number of times (amplification of the histogram bins) 
    """
    def get_stats_files(self):
        # In a CSV file    
        csv_stats = re.sub("([0-9]{8})", "stats.csv", os.path.splitext(self.files[0])[0]) # YYYYMMDD
        csv_hist = re.sub("([0-9]{8})", "hist.csv", os.path.splitext(self.files[0])[0]) # YYYYMMDD
        self.logger.info('Saving stats in {}'.format(os.path.basename(csv_stats)))
        print('INFO: Saving stats in {}'.format(csv_stats))
        self.logger.info('Saving histograms in {}'.format(os.path.basename(csv_hist)))
        print('INFO: Saving histograms in {}'.format(csv_hist))        

        return csv_stats, csv_hist

    """
    Get Histogram bins initalization. Number of bins + K = number of times (amplification of the histogram bins) 
    """
    def get_bins_init(self, f, nbins=2000, Kmax=5):
        
        # Try to load data file
        _, arr, errmsg = self.load_datafile(f)
        if errmsg != '':
            return None, errmsg

        # Return histogram initialization
        minv = np.min(arr)
        maxv = np.max(arr)
        dist = abs(maxv-minv)        

        return np.histogram_bin_edges(arr, bins=nbins, range=(minv-Kmax*dist, maxv+Kmax*dist)), ''

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

        # Store stats information per file       
        bn = [ os.path.basename(f) for f in self.files ] 
        df_stats = pd.DataFrame(index=bn, columns=self.stat_columns)

        # Init histogram bins
        hist_bins, errmsg = self.get_bins_init(self.files[0])
        if errmsg != '':
            return 'err', errmsg

        # We keep statistics per latitude    
        df_hist = pd.DataFrame(index=bn+['N_polar', 'N_midlat', 'N_hadley', 'total'], columns=hist_bins[:-1])
       
        # Generate histograms, store them in a dataframe
        ind=0        
        for f in self.files:
            # Progress update
            i+=1
            print_progress(i, n)              
            
            # Read data
            ds, arr, errmsg = self.load_datafile(f)
            if errmsg != '':
                return 'err', errmsg

            # Latitude histograms
            try:
                latid = get_lat_name(ds)
                self.get_hist_hadley(df_hist, hist_bins, ds, ind,  60,  90, 'N_polar', latid)
                self.get_hist_hadley(df_hist, hist_bins, ds, ind,  30,  60, 'N_midlat', latid)
                self.get_hist_hadley(df_hist, hist_bins, ds, ind,   0,  30, 'N_hadley', latid)
                self.get_hist_hadley(df_hist, hist_bins, ds, ind, -30, -0, 'S_hadley', latid)
                self.get_hist_hadley(df_hist, hist_bins, ds, ind, -60, -30, 'S_midlat', latid)
                self.get_hist_hadley(df_hist, hist_bins, ds, ind, -90, -60, 'S_polar', latid)
            except:
                pass # non latlon dataset

            # Compute histogram (global/month)             
            x, b = np.histogram(arr, bins=hist_bins)    
            df_hist.loc[os.path.basename(f)] = x

            # Compute stats (of the histogram and the data itself)            
            df_stats.loc[os.path.basename(f)] = get_basic_stats(arr)
            
            # Free Mem
            del arr
            ind+=1

        # Compute total histogram        
        df_hist.loc['total'] = df_hist.sum().values
        
        # Save stats for future use
        csv_stats, csv_hist = self.get_stats_files()
        df_hist.index.name = 'File'
        df_hist.to_csv(csv_hist)
        df_stats.index.name = 'File'
        df_stats.to_csv(csv_stats)

        # Defaults
        res, msg = '', ''

        # Just histogram delivered, unknown ranges for the variable
        if self.valid_min is None and self.valid_max is None:
            res = 'warn'
            msg = 'Physical plausible range unknown. Observed range: [ from {} to  {} ] {}'.format(
                format_float(min(df_stats['Min'])), format_float(max(df_stats['Max'])), self.metadata['unit'])

        # Valid range estimation (only if defined, well-known variable limits)               
        else:
            for fname, row in df_stats.iterrows():
                # 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 = 'err'               

            # Conclusion plausible range            
            if res == '':
                res = 'ok'
                msg = 'No outliers found. Plausible range: [ from {} to  {} ]'.format(self.valid_min, self.valid_max)      
            else:
                msg = 'Potential outliers found. Plausible range: [ from {} to  {} ]'.format(self.valid_min, self.valid_max)

        # Final plots
        bp = BokehPlot(self.output_plot_html, csv_stats, csv_hist, self.cds_variable, self.valid_min, self.valid_max, self.metadata)
        bp.plot_histograms()
                
        # Final message
        self.logger.info(msg)
        return res, msg

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

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