#!/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, gc
import warnings
warnings.filterwarnings('ignore')

# Local imports
from dqc_wrapper.utils.utils_forms 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
---
- Generate histograms of data values per month
- If valid_min/valid_max are provided, raise error if outlier found
"""
class ValidRangesTest:

    def __init__(self, resdir, testname, files, varname, vmin, vmax, dt, forms_dir, cds_dataset, cds_variable, cds_units):
        self.resdir = resdir
        self.testid = 'test06-valid_ranges'
        self.logger = logging.getLogger(self.testid)
        self.title = testname.replace('.ini', '')
        self.output_plot_html = os.path.join(resdir, testname.replace('.ini', '-valid_ranges.html'))
        self.csv_hist = os.path.join(resdir, testname.replace('.ini', '-hist.csv'))
        self.csv_stats = os.path.join(resdir, testname.replace('.ini', '-stats.csv'))
        self.files = files
        self.varname = varname     
        self.datatype = dt
        self.stat_columns = [ 'dataType', 'Min', 'Max', 'Med', 'q1', 'q3', 'IQR' ]
        self.cds_variable = cds_variable
        self.cds_dataset = cds_dataset
        self.metadata = get_metadata(forms_dir, cds_dataset, cds_variable)
        self.units = cds_units

        # We keep statistics per hadley-cell, filename and total [order is relevant]
        self.regions = [ 'total', 'N_polar', 'N_midlat', 'N_hadley', 'S_hadley', 'S_midlat', 'S_polar' ]
        self.hadley_cells = {
            'N_polar': [60, 90],
            'N_midlat': [30, 60],
            'N_hadley': [0, 30],
            'S_hadley': [-30, 0],
            'S_midlat': [-60, -30],
            'S_polar': [-90, -60]
        }

        # User-defined valid_ranges (i.e - WMO official ranges)
        try:
            self.valid_min = float(vmin)
        except:
            self.valid_min = None
        try:
            self.valid_max = float(vmax)
        except:
            self.valid_max = None

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

        # Initialization vs sum
        if init==0:
            df_hist.loc[row] = x
        else:
            df_hist.loc[row] += x

        return x, [ arr.min(), arr.max() ]

    """
    Get Histogram bins initalization. 
    Number of bins + K = number of times (amplification of the histogram bins) 
    Initialization observed minmax
    """
    def get_bins_init(self, files, nbins=2000, Kmax=5, max_samples=24):

        # XARRAY (NetCDF)
        if files[0].endswith('.nc'):
            # Try to load data from the first file
            ds, errmsg = load_datafile(files[0])
            if errmsg != '':
                return None, None, None, errmsg

            # Load all data from all product types
            arr = get_variable(ds, self.varname)

        # ECCODES (GRIB)
        else:
            # Get a sample from every first message of every file [better, faster, stronger]
            arr = []

            # Avoid useless memory consumption (2 years is good enough - monthly, 12 years daily)
            if len(files) > max_samples:
                files = files[0:max_samples]

            for f in files:
                gf = GribFileCheck(f, dt=self.datatype)
                arr += gf.get_grib_sample() # Get extreme values for file i

        # Return histogram initialization
        minv = np.min(arr)
        maxv = np.max(arr)
        dist = abs(maxv-minv)
        hist = np.histogram_bin_edges(arr, bins=nbins, range=(minv - Kmax * dist, maxv + Kmax * dist))

        return hist, minv, maxv, ''

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

        # Store stats information per file (basename is the unique key and contains date)
        bn = [ os.path.basename(f) for f in self.files ]
        sn = []
        for b in bn:
            for d in self.datatype:
                sn.append('{}_{}'.format(b,d))

        # Load previous stats or create from scratch
        if os.path.exists(self.csv_stats):
            df_stats = pd.read_csv(self.csv_stats)
            df_stats.set_index('File_dt', inplace=True)
            df_stats = df_stats.drop(['Ind'], axis=1, errors='ignore') # Old datasets might contain this unnecessary column
        else:
            df_stats = pd.DataFrame(index=sn, columns=self.stat_columns)

        # Load previous histograms or create from scratch -> Rows = files + hadley_cells
        if os.path.exists(self.csv_hist):
            df_hist = pd.read_csv(self.csv_hist)
            df_hist.set_index('File', inplace=True)
            hist_bins = [float(i) for i in df_hist.columns.tolist()]
            hist_bins.append(hist_bins[-1] + (hist_bins[1] - hist_bins[0])) # Number of bins has to have +1
            errmsg = ''
        else:
            # Init histogram bins with a sample
            hist_bins, minv, maxv, errmsg = self.get_bins_init(self.files)
            if errmsg != '':
                return 'err', errmsg

            df_hist = pd.DataFrame(
                index=bn + list(self.hadley_cells.keys()),
                columns=hist_bins[:-1]
            )
       
        # Generate histograms, store them in a dataframe
        ind = 0
        s = 0

        for f in self.files:

            # Progress update
            i+=1
            print_progress(i, n)
            minmax = None

            # NetCDF - read all datatypes individually
            if f.endswith('.nc'):

                # Load xarray dataset
                ds, errmsg = load_datafile(f)
                if errmsg != '':
                    return 'err', errmsg

                # Latitude histograms - subset and cumulative sum for 6-hadley cells
                latid = get_lat_name(ds)
                first = True

                for hc, lats in self.hadley_cells.items():
                    hc_hist, mnmx = self.get_hist_hadley_xarray(
                        df_hist, hist_bins, ds, ind,  lats[0],  lats[1], hc, latid
                    )

                    # Compute histogram (global/month) as sum of all latitude ranges (init vs accum)
                    if first:
                        xsum = hc_hist
                        minv = min(mnmx)
                        maxv = max(mnmx)
                    else:
                        xsum = sum([hc_hist, xsum])
                        minv = min(minv, mnmx)
                        maxv = max(maxv, mnmx)
            else:
                # All regions at once (more efficient) - message x message
                ds = GribFileCheck(f, dt=self.datatype)
                xsum, minmax = ds.get_grib_hadley_regions(ind, hist_bins, df_hist)
                minv, maxv = get_min_max(minmax)

            """ - Extension histogram ---> to be tested
            # See if we need to extend the histogram (some variables fall out)
            x_ext = extend_xhist(xsum, [minv, maxv])

            # Store general/full histogram - see if extension of Dataframe and update of bins are needed
            if len(x_ext) == len(xsum):
                df_hist.loc[os.path.basename(f)] = xsum
            else:
                hist_bins = extend_nhist(xsum, n, x_ext)
                df_hist.loc[os.path.basename(f)] = x_ext
            """

            # Save histogram
            df_hist.loc[os.path.basename(f)] = xsum

            # Compute stats for every datatype (of the histogram - faster)
            try:
                if minmax:
                    # For every available data type
                    for dt in self.datatype:
                        mindt , maxdt = minmax[dt][0], minmax[dt][1]
                        idfile = os.path.basename(f)+'_'+dt
                        if not(mindt) is None and not(maxdt) is None:
                            df_stats = update_stats_row(idfile, df_stats, dt, minv, maxv, xsum, hist_bins)
                            s += 1
                else:
                    # No datatype distinction
                    idfile = os.path.basename(f) + '_all'
                    df_stats = update_stats_row(idfile, df_stats, 'all', minv, maxv, xsum, hist_bins)
                    s+=1
            except Exception as e:
                print(e)
                raise Exception("Unable to compute statistics for file={}. Please check GRIB content carefully.".format(os.path.basename(f)))

            # Next iteration
            del xsum
            del ds
            ind+=1

        # Compute total histogram        
        df_hist.loc['total'] = df_hist.sum().values
        
        # Save hist for future use
        df_hist.index.name = 'File'
        df_hist.sort_index(inplace=True)
        df_hist.to_csv(self.csv_hist)
        print('INFO: Saving histogram to {}'.format(self.csv_hist))

        # Save stats for future use
        df_stats.index.name = 'File_dt'
        df_stats.sort_index(inplace=True)
        df_stats.to_csv(self.csv_stats)
        print('INFO: Saving stats to {}'.format(self.csv_stats))

        # Estimated valid_min/valid_max (conservative approach) - Catch exception
        control_limits = {}
        for reg in self.regions:
            vmin_nmax, vmax_nmax, o_vmin, o_vmax, others = None, None, None, None, None
            try:
                mvr = multiparametric_valid_range(csv_path=self.resdir, logarithmic_hist=[True],
                                                  dr_gap=np.arange(0.01, 0.3, 0.01), r_zero=np.arange(0.01, 0.1, 0.01),
                                                  k_gap=np.arange(1, 6, 2), extra=False, domain=reg)
                (vmin_nmax, vmax_nmax), (o_vmin, o_vmax), err, others = mvr.get(file=self.csv_hist, logger=self.logger)
                control_limits[reg] = {
                    'vmin_nmax': vmin_nmax,
                    'vmax_nmax': vmax_nmax,
                    'o_vmin': o_vmin,
                    'o_vmax': o_vmax,
                }
            except Exception as e:
                print('ERROR: estimating ranges with multiparametic method, {}'.format(e))

        # Evaluate and return message
        try:
            res, msg = get_plausible_ranges(df_stats, control_limits['total']['vmin_nmax'], control_limits['total']['vmax_nmax'],
                                            self.valid_min, self.valid_max, self.units, self.logger)
        except Exception as e:
            print('ERROR: getting plausible ranges limits, {}'.format(e))

        # Final plots
        bp = BokehPlot(self.title, self.output_plot_html, self.csv_stats, self.csv_hist, self.cds_variable, self.cds_dataset,
                       self.valid_min, self.valid_max, control_limits, self.metadata, self.units, msg, res)
        bp.plot_histograms()
                
        # Final message
        if res == 'ok':
            self.logger.info(msg)
        elif res == 'warn':
            self.logger.warning(msg)
        else:
            self.logger.error(msg)

        return res, msg

    def main(self):
       
        # Do the actual work
        try:
            res, msg = self.check_valid_ranges()
        except Exception as e:
            res = 'err'
            msg = str(e)

        # Prepare output
        return {
            'res': res, 
            'msg': msg,
            'plot': self.output_plot_html
        }
