#!/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:
#       Jesus Pena Izquierdo (jesus.pena@bsc.es)
#       Joan Sala Calero (joan.sala@bsc.es)

from pathlib import Path
import numpy as np
import pandas as pd
import logging
from dqc_wrapper.utils.utils_files import format_float

"""
Get test message/result depending on results and tolerance/control limits
"""
def get_plausible_ranges(df_stats, vmin_nmax, vmax_nmax, valid_min, valid_max, units, logger):

    # Drop N/A
    df_stats = df_stats.dropna()

    # Unknown tolerance limits for the variable
    if valid_min is None and valid_max is None:
        if vmin_nmax is None and vmax_nmax is None:
            res = 'err'
            msg_limits = '[variable ranges = ({}, {}) {}]'.format(
                format_float(min(df_stats['Min']), 1), format_float(max(df_stats['Max']), 1), units
            )
            msg = 'Tolerance limits unknown. Unable to estimate the control limits.' + msg_limits
        else:
            res = check_vmin_vmax(df_stats, vmin_nmax, vmax_nmax, logger, res_val='warn')
            msg_limits = '[control limits = ({}, {}) {}; variable ranges = ({}, {}) {}]'.format(
                format_float(vmin_nmax, 1), format_float(vmax_nmax, 1), units,
                format_float(min(df_stats['Min']), 1), format_float(max(df_stats['Max']), 1), units
            )

            # Control limits
            if res == 'ok':
                msg = 'The variable is within the expected ranges. ' + msg_limits
            else:
                msg = 'Potential outliers identified. ' + msg_limits

    # Valid range estimation (only if defined, well-known variable limits)
    else:
        res = check_vmin_vmax(df_stats, valid_min, valid_max, logger, res_val='err')
        msg_limits = '[tolerance limits = ({}, {}) {}; variable ranges = ({}, {}) {}]'.format(
            format_float(valid_min, 1), format_float(valid_max, 1), units,
            format_float(min(df_stats['Min']), 1), format_float(max(df_stats['Max']), 1), units
        )

        # Tolerance limits
        if res == 'ok':
            msg = 'The variable is within the expected ranges. ' + msg_limits
        else:
            msg = 'The variable presents data values outside the specifications defined by the tolerance limits. ' + msg_limits

    return res, msg

"""
Check for min/max per file (values are given)
"""
def check_vmin_vmax(df_stats, valid_min, valid_max, logger, res_val='err', bins_res=2000.0):
    res = 'ok'

    # Min max can also be from different estimates
    if type(valid_min) == list or type(valid_min) == np.ndarray:
        valid_min = min(valid_min)
    if type(valid_max) == list or type(valid_max) == np.ndarray:
        valid_max = max(valid_max)

    # We add the histogram err component in case (bins resolution = 2000.0)
    if res_val == 'err' and not(valid_min is None) and not(valid_max is None):
        err_tol = abs(valid_max-valid_min)/bins_res
        valid_min -= err_tol
        valid_max += err_tol

    for fname, row in df_stats.iterrows():
        # Check ranges (1 or 2)
        outlier = False
        if not(valid_min is None) and float(row['Min']) < valid_min:
            res = res_val
            outlier = True
        if not(valid_max is None) and float(row['Max']) > valid_max:
            res = res_val
            outlier = True

        if outlier and logger:
            logger.error('{} contains outliers on dataType={}'.format(fname, row['dataType']))
            logger.error('observed [{} to {}] vs valid [{} to {}]'.format(format_float(row['Min']), format_float(row['Max']),
                                                                        format_float(valid_min), format_float(valid_max)))

    return res

"""
Update df_stats object row
File_dt,Min,Max,Med,q1,q3,IQR,dataType
"""
def update_stats_row(idfile, df_stats, dt, minv, maxv, xsum, hist_bins):

    # Collect basic stats for file 'i'
    Med, q1, q3, IQR = get_basic_stats(xsum, hist_bins)

    # Update content or insert
    if idfile in df_stats.index:
        df_stats.loc[idfile]['Min'] = minv
        df_stats.loc[idfile]['Max'] = maxv
        df_stats.loc[idfile]['Med'] = Med
        df_stats.loc[idfile]['q1'] = q1
        df_stats.loc[idfile]['q3'] = q3
        df_stats.loc[idfile]['IQR'] = IQR
        df_stats.loc[idfile]['dataType'] = dt
    else:
        series = pd.Series({'Min': minv, 'Max': maxv, 'Med': Med, 'q1': q1, 'q3': q3, 'IQR': IQR, 'dataType': dt})
        series.name = idfile
        df_stats = df_stats.append(series)

    return df_stats

"""
Check all datatypes dictionary and get absolute minimum and maximum values
"""
def get_min_max(minmax):
    minv = np.nan
    maxv = np.nan
    for dt, vals in minmax.items():
        # Not all datatypes are present in all grib files
        if not(np.isnan(vals[0])) and not(np.isnan(vals[1])):
            if np.isnan(minv):
                minv = min(vals)
            else:
                minv = min(vals + [minv])
            if np.isnan(maxv):
                maxv = max(vals)
            else:
                maxv = max(vals + [maxv])

    return minv, maxv

"""
This function extends the initial histogram bin edges array (x) depending on observed minimum or maximum
in include in minmax. If they are already contained within x, no extension occurred. If they are outside x, a new
extended bin edges array is created to contain both min and max values mantaining the initial x values within:

* x = edges of the histogram bins. len(x) = Nbins+1
* minmax = (observed min, observed max). Observed min and max of the new file
* k_ext = coefficient of limit extension to extend a bit beyond new observed min or max
"""
def extend_xhist(x, minmax, k_ext = 1.1):

    x_min, x_max = x[0], x[-1]
    dx = np.diff(x)[0]
    if minmax[0] < x_min:
        # Extension of the left side required
        dx_min = x_min - minmax[0]
        x_newmin = minmax[0] - dx_min*k_ext
        x_ext_min = np.arange(x_min - dx, x_newmin - dx, -dx)
        x_ext_min = np.flip(x_ext_min)
    else:
        # No extension of the left side required
        x_ext_min = []
    if minmax[1] > x_max:
        # Extension of the righ side required
        dx_max = minmax[1] - x_max
        x_newmax = minmax[1] + dx_max*k_ext
        x_ext_max = np.arange(x_max + dx, x_newmax + dx, dx)
    else:
        # No extension of the right side required
        x_ext_max = []
    x_ext = np.concatenate((x_ext_min, x, x_ext_max), axis = 0)
    return x_ext

"""
This function extend the histogram array of counts (n) to the extended histogram bin edges array (x_ext).
Array x_ext should be generated with function extend_xhist so array x is contained within x_ext
If x_ext equals x, n_ext = n. If x_ext is shorter than x, an error prompts:
    
* x = edges of the histogram bins. len(x) = Nbins+1
* n = counts for each histogram bins. len(n) = Nbins
* x_ext = array of the extended histogram bins edges
"""
def extend_nhist(x, n, x_ext):

    i = int(np.argwhere(x_ext == x[0]))
    n_ext = np.zeros(len(x_ext)-1)
    n_ext[i : i + len(n)] = n
    return n_ext

"""
Producing a list of statistical estimators (on the data)
Stats needed for the boxplot
"""
def get_basic_stats(x, b):

    # Stats from histogram (more efficient, almost same result)
    q1, Med, q3 = stats_from_hist(b, x)
    IQR = q3 - q1

    return [Med, q1, q3, IQR]

"""
Get approximate statistical indicators from histogram
- This function computes the position of the quantiles given in the list quantiles.
- It linearly interpolates from the histogram to the corresponding values (i.e. no discretization issue)  
"""
def stats_from_hist(xm, nx, quantiles=[0.25, 0.5, 0.75]):

    dx = np.diff(xm)[0]  # xbin length
    N = np.sum(nx)  # total number of elements
    rx_cumsum = np.cumsum(nx) / N  # density distribution

    x_q = np.array([])
    for qi in quantiles:
        i_qi1 = np.argwhere(rx_cumsum > qi)[0].item()
        i_qi0 = i_qi1 - 1

        rx_qi0 = rx_cumsum[i_qi0]
        rx_qi1 = rx_cumsum[i_qi1]

        drx_qi01 = rx_qi1 - rx_qi0
        drx_qi = qi - rx_qi0
        k_drx_qi = drx_qi / drx_qi01

        x_qi = xm[i_qi0] + dx * (k_drx_qi + 1 / 2)
        x_q = np.append(x_q, x_qi)

    return x_q


"""
Class used for estimating the valid ranges
"""
class valid_range:

    def __init__(self,
                 data=None,
                 file_name=None,
                 file_path=None,
                 logger=None,
                 domain='total',
                 dr_gap=0.05,
                 r_zero=0.01,
                 k_gap=1,
                 logarithmic_hist=False,
                 verbose=True
                 ):

        self.dr_gap = dr_gap  # Percentile distance as definition of outlier gap
        self.r_zero = r_zero  # Percentile definition of start/end of the data distribution
        # This depends on the number of elements, not use r_zero < 0.0001
        self.k_gap = k_gap  # Coefficient to multiply the estimated outlier distance
        self.file_name = file_name  # file name of csv file to process
        self.file_path = file_path  # path of the csv file to process
        self.domain = domain  # domain of the csv file to be processed ['total', 'S_polar','S_midlat',...]
        self.verbose = verbose
        self.log_hist = logarithmic_hist
        self.data = data

        # Read csv file
        if data is None:
            file_path = Path(self.file_path).joinpath(self.file_name)
            data0 = pd.read_csv(file_path)
            data0.set_index('File', inplace=True)
            self.data = data0.loc[self.domain]

        # Logger
        if logger is None:
            self.logger = logging.getLogger('valid_range')

    def set_parameters(self):

        xedge = self.data.index.values.astype('float64')  # bin edges
        self.xm = xedge[:-1] + np.diff(xedge) / 2  # bin centre
        self.nx = self.data.values[:-1]  # bin counts

        if self.log_hist:
            self.nx = self.data.values[:-1] + 1
            # self.nx = np.log(self.nx)
            self.nx = np.log10(self.nx)

        self.N = np.sum(self.nx)  # total number of elements
        self.quantiles = np.cumsum(self.nx) / self.N  # density distribution. percentiles

        # Percentiles selected for the gap definition
        self.q_left = [self.r_zero, self.r_zero + self.dr_gap]
        self.q_right = [1 - self.dr_gap - self.r_zero, 1 - self.r_zero]

    '''
    This function computes the position of the quantiles given in the list quantiles.
     It linearly interpolates from the histogram to the corresponding value (i.e. no discretization issue)
    '''
    def quantiles_from_hist(self, quantiles):

        dx = np.diff(self.xm)[0]  # xbin length
        N = np.sum(self.nx)  # total number of elements
        rx_cumsum = np.cumsum(self.nx) / N  # density distribution

        x_q = np.array([])
        for qi in quantiles:
            i_qi1 = np.argwhere(rx_cumsum > qi)[0].item()
            i_qi0 = i_qi1 - 1

            rx_qi0 = rx_cumsum[i_qi0]
            rx_qi1 = rx_cumsum[i_qi1]

            drx_qi01 = rx_qi1 - rx_qi0
            drx_qi = qi - rx_qi0
            k_drx_qi = drx_qi / drx_qi01

            x_qi = self.xm[i_qi0] + dx * (k_drx_qi + 1 / 2)
            x_q = np.append(x_q, x_qi)

        return x_q

    '''
    This function computes the outlier gaps for the left and right tails of the distribution.
        (r_gap => dx_left, dx_right)
        This is the minimum size of data distribution gaps in each tail to consider data clusters separated
        from the main distribution as real outliers.
        This outlier_gap is defined as the distance between especified quantiles:
        For the left outlier_gap: abs( x_q[r_zero] - x_q[r_zero+dr_gap] )
        For the right outlier_gap: abs( x_q[1-dr_gap-r_zero] - [1-r_zero] )
    '''
    def outlier_gaps(self):

        self.set_parameters()

        # exact quantile positions for left and right bands
        xq_left = self.quantiles_from_hist(quantiles=self.q_left)
        xq_right = self.quantiles_from_hist(quantiles=self.q_right)

        # outlier gap for left band
        dx_left = self.k_gap * np.diff(xq_left).item()

        # outlier gap for right band
        dx_right = self.k_gap * np.diff(xq_right).item()

        return dx_left, dx_right, xq_left, xq_right

    '''
    This function search for gaps greater than dx_left and dx_right at each side of the distribution
    '''
    def find_gaps(self):

        dx_left, dx_right, _, _ = self.outlier_gaps()

        # For left side:
        # --------------
        # Bins to the left of r_zero
        i_left = np.argwhere(self.quantiles < self.r_zero)[:, 0]
        nx_left = self.nx[i_left]
        if nx_left[0] == 0:
            # Index of last bins to the left of r_zero with nx = 0
            i_zero = np.argwhere(nx_left == 0)[-1].item()
            # Bins to the left of i_zero
            i_left0 = i_left[:i_zero]
            nx_left0 = self.nx[i_left0]
            xm_left0 = self.xm[i_left0]

            # Searching for a group of bins with nx=0 (i.e. n_gap=0) and extension < dx_left
            warn_left = True
            n_gap = 1
            j = -1  # counting down from i_zero to the first elements. From right to left
            while n_gap:
                i0_gap_left = j  # -1
                xm0_gap_left = xm_left0[i0_gap_left]  # right edge of the gap
                xm1_gap_left = xm0_gap_left - dx_left  # left edge of the gap

                if xm1_gap_left >= xm_left0[0]:
                    i1_gap_left = np.argwhere(xm_left0 - xm1_gap_left < 0)[-1].item()  # index of the left edge element
                    n_gap = np.sum(nx_left0[i1_gap_left: i0_gap_left])  # total number of counts within the gap
                    j -= 1

                else:
                    i1_gap_left = 0  # index of the left edge element
                    n_gap = np.sum(nx_left0[i1_gap_left: i0_gap_left])  # total number of counts within the gap
                    j -= 1

                    if self.verbose and warn_left:
                        xm0_gap_left, xm1_gap_left = np.nan, np.nan
                        self.logger.warning('WARNING-left: Estimated valid left range out of histogram limit. Extension of histogram left limit required')
                        warn_left = False

        else:
            xm0_gap_left, xm1_gap_left = np.nan, np.nan
            if self.verbose:
                self.logger.error('ERROR-left: Observed range goes beyond histogram left limit. Extension of histogram left limit required')

        # For right side:
        # ----------------
        # Bins to the right of 1-r_zero
        i_right = np.argwhere(self.quantiles > 1 - self.r_zero)[:, 0]
        nx_right = self.nx[i_right]

        if nx_right[-1] == 0:
            # Index of first bins to the right of 1-r_zero with nx = 0
            i_zero = np.argwhere(nx_right == 0)[0].item()
            # Bins to the right of 1-i_zero
            i_right0 = i_right[i_zero:]
            nx_right0 = self.nx[i_right0]
            xm_right0 = self.xm[i_right0]

            # Searching for a group of bins with nx=0 (i.e. n_gap=0) and extension < dx_right
            warn_right = True
            n_gap = 1
            j = 0
            while n_gap:
                i0_gap_right = j  # 0
                xm0_gap_right = xm_right0[i0_gap_right]  # left edge of the gap
                xm1_gap_right = xm0_gap_right + dx_right  # right edge of the gap

                if xm1_gap_right <= xm_right0[-1]:
                    i1_gap_right = np.argwhere(xm_right0 - xm1_gap_right > 0)[
                        0].item()  # index of the right edge element
                    n_gap = np.sum(nx_right0[i0_gap_right: i1_gap_right])  # total number of counts within the gap
                    j += 1

                else:
                    i1_gap_right = len(nx_right0)  # index of the right edge element
                    n_gap = np.sum(nx_right0[i0_gap_right: i1_gap_right])  # total number of counts within the gap
                    j += 1

                    if self.verbose and warn_right:
                        self.logger.warning('WARNING-right: Estimated valid right range out of histogram limit')
                        warn_right = False

        else:
            xm0_gap_right, xm1_gap_right = np.nan, np.nan
            if self.verbose:
                self.logger.error('ERROR-right: Observed range goes beyond histogram right limit')

        return xm1_gap_left, xm0_gap_left, xm0_gap_right, xm1_gap_right

    def get(self):

        # Get outer boundaries of the empty gaps.
        xm1_gap_left, _, _, xm1_gap_right = self.find_gaps()

        return xm1_gap_left, xm1_gap_right


"""
Class used for estimating the valid ranges 
(improved version with multiparametric approach and plausabilities)

    This method computes for each csv file a "plausibility" distribution for the valid range of each tail.
    This plausibility is the proportion of tests that consider a specific value realistic or not
    based on the hipothesys of outlier-gaps.
    Different valid ranges are estimated for different plausibility percentiles specified in parameter
    percentile for each tail (p_vmin, p_vmax).
    
"""
class multiparametric_valid_range:

    def __init__(self,
                 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),
                 logarithmic_hist=[True],  # [False, True],
                 csv_path='',
                 domain='total',
                 nbins=100,
                 extra=False,
                 percentile=np.arange(0, 105, 5)
                 ):

        self.percentile = percentile
        self.csv_path = csv_path
        self.dr_gap = dr_gap
        self.r_zero = r_zero
        self.k_gap = k_gap
        self.domain = domain
        self.log_hist = logarithmic_hist

        self.nbins = nbins  # Number of bins to compute the probability distribution in self.get()
        self.extra = extra  # Enable or disable the computation of extra parameters for the valid range estimation

    def set_vr(self, data=None, file=None, logger=None, log_scale=False):

        # Define the histogram in logarithmic scale for the outlier computation
        vr = valid_range(data=data,
                         file_name=file,
                         file_path=self.csv_path,
                         logger=logger,
                         domain=self.domain,
                         verbose=False,
                         logarithmic_hist=log_scale
                         )

        vr.set_parameters()
        self.xm = vr.xm
        self.nx = vr.nx

        return vr

    '''
    This function computes the valid ranges for the file "file".
    
    vmin_nmax, vmax_nmax are the estimated valid range with the highest repetition frequency within the
    all the possible combination of parameters.
    
    err indicates of the domain is shorter than required by the method
    '''
    def get(self, data=None, file=None, logger=None):

        err = False
        Lp = len(self.percentile)
        L_dr_gap = len(self.dr_gap)
        L_r_zero = len(self.r_zero)
        L_k_gap = len(self.k_gap)
        L_log_hist = len(self.log_hist)
        xmin = np.empty(shape=[L_dr_gap, L_r_zero, L_k_gap, L_log_hist])
        xmax = np.empty(shape=[L_dr_gap, L_r_zero, L_k_gap, L_log_hist])
        vr = self.set_vr(data=data,
                         file=file,
                         logger=logger
                         )

        for l, log_l in enumerate(self.log_hist):
            vr.log_hist = log_l
            vr.set_parameters()

            for i, dr_i in enumerate(self.dr_gap):
                for j, r_j in enumerate(self.r_zero):
                    for k, k_h in enumerate(self.k_gap):
                        vr.dr_gap = dr_i
                        vr.r_zero = r_j
                        vr.k_gap = k_h

                        # Estimating the valid ranges
                        xmin[i, j, k, l], xmax[i, j, k, l] = vr.get()

        vmin = xmin.flatten()
        vmax = xmax.flatten()
        others = [(vmin, vmax)]

        if (np.isnan(vmin).sum() > 0) | (np.isnan(vmax).sum() > 0):
            err = True
            vmin_nmax, vmax_nmax, o_vmin, o_vmax = 4 * [np.nan]
            raise Exception('Variable range out of domain')

        else:  # Computation of valid range based on maximum frequency/repetition

            n_vmin, xedge_vmin = np.histogram(vmin, bins=self.nbins)
            x_vmin = xedge_vmin[:-1] + np.diff(xedge_vmin)  # x_vmin is bin centred
            vmin_nmax = x_vmin[np.argmax(n_vmin)]

            n_vmax, xedge_vmax = np.histogram(vmax, bins=self.nbins)
            x_vmax = xedge_vmax[:-1] + np.diff(xedge_vmax)  # Cx_vmax is bin centred
            vmax_nmax = x_vmax[np.argmax(n_vmax)]

            # Count number of outliers out of estimated ranges
            o_vmin, o_vmax = self.outlier_detection(limits=[vmin_nmax, vmax_nmax],
                                                    data=data,
                                                    file=file
                                                    )
            if self.extra:
                others.append(self.get_extra([vmin, vmax]))

        return (vmin_nmax, vmax_nmax), (o_vmin, o_vmax), err, others

    '''
    If parameter extra = True, the following estimations are also computed
        p_vmin[i] is the estimated valid range for the left tail (minimum plausible value for the variable)
        such that a proportion of percentile[i] of the tests estimates plausible values above this threshold.
        p_vmax[i] is the estimated valid range for the right tail (maximum plausible value for the variable)
        such that a proportion of percentile[i] of the tests estimates plausible values below this threshold.
        o_vmin[i] is the number of outliers found considering percentile[i]
        o_vmax[i] is the number of outlier found considering percentile[i]    
    '''
    def get_extra(self, vmin_vmax):

        Lp = len(self.percentile)
        vmin, vmax = vmin_vmax

        # Preallocating percentiled valid range
        p_vmin = np.full((1, Lp), np.nan)
        p_vmax = np.full((1, Lp), np.nan)
        op_vmin = np.full((1, Lp), np.nan)
        op_vmax = np.full((1, Lp), np.nan)

        # Valid range for the corresponding percentile threshold
        p_vmin = np.percentile(vmin, q=self.percentile)
        p_vmax = np.percentile(vmax, q=self.percentile)

        # Number of outlier for the corresponding percentile threshold
        for i_p, p_i in enumerate(self.percentile):
            op_vmin[0, i_p], op_vmax[0, i_p] = self.outlier_detection([p_vmin[i_p], p_vmax[i_p]])

        # Percentiles should be reversed for the right tail to be consistent
        p_vmax = np.flip(p_vmax)
        op_vmax = np.flip(op_vmax)

        return (p_vmin, p_vmax), (op_vmin, op_vmax)

    '''
    This function return the number of outliers given the limits defined in "limits"
    '''
    def outlier_detection(self, limits, data=None, file=None):

        lim_min, lim_max = limits

        # if self.log_hist:
        #  o_min = (np.power(10, self.nx[self.xm < lim_min])-1).sum()
        #  o_max = (np.power(10, self.nx[self.xm > lim_max])-1).sum()
        # else:
        o_min = self.nx[self.xm < lim_min].sum()
        o_max = self.nx[self.xm > lim_max].sum()

        return o_min, o_max

    '''
    his function computes and store the output of the multiparametric valid range.
    Dataframes can be obtained when parameter "dataframes" is set to True. numpy array otherwise
    '''
    def get_all_files(self, files, dataframes=True):

        err_list = []
        i_err = 0
        Lp = len(self.percentile)
        Lv = len(files)

        # Preallocating percentiled valid range
        # vmin_nmax = np.full(Lv, np.nan)
        # vmax_nmax = np.full(Lv, np.nan)

        out = pd.DataFrame(columns=['vr_min', 'vr_max', 'outliers_min', 'outliers_max', 'error'])
        for i_v, file_i in enumerate(files):

            print(i_v, file_i)
            (vmin_nmax_i, vmax_nmax_i), (o_vmin, o_vmax), err, others = self.get(file=file_i)
            out.loc[file_i] = [vmin_nmax_i, vmax_nmax_i, o_vmin, o_vmax, err]
            extra = None

            if self.extra:
                cols = ['p' + str(i) for i in self.percentile]
                extra_pmin = pd.DataFrame(columns=cols)
                extra_pmax = pd.DataFrame(columns=cols)
                extra_omin = pd.DataFrame(columns=cols)
                extra_omax = pd.DataFrame(columns=cols)
                extra_pmin.loc[file_i] = others[1][0][0]
                extra_pmax.loc[file_i] = others[1][0][1]
                extra_omin.loc[file_i] = others[1][1][0]
                extra_omax.loc[file_i] = others[1][1][1]
                extra = (extra_pmin, extra_pmax, extra_omin, extra_omax)

        return out, extra

"""
Min/Max functions with None initialization
"""
def min_None(minv, minc):
    if minv is None:
        return minc
    else:
        return min(minv, minc)

def max_None(maxv, maxc):
    if maxv is None:
        return maxc
    else:
        return max(maxv, maxc)