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

import numpy as np
import pandas as pd

from dqc_wrapper.utils.utils_files import format_float

"""
Check for min/max per file (values are given)
"""
def check_vmin_vmax(df_stats, valid_min, valid_max, logger, res_val='err'):
    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)

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

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

"""
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 = quantiles_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 quantiles_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

"""
Try different combinations of (dr_gap, r_zero) to find gaps in distribution and better estimate vmin/vmax
"""
def exhaustive_vmin_vmax(csv_hist, perc=[25,75]):
    mvr = multiparametric_valid_range(percentile=perc)
    (vmin, vmax), (p_vmin, p_vmax), (o_vmin, o_vmax), err = mvr.get(csv_hist)

    return p_vmin, p_vmax


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

    def __init__(self,
                 csv_hist,
                 domain='total',
                 dr_gap=0.05,
                 r_zero=0.01,
                 k_gap=1,
                 logarithmic_hist=False,
                 verbose=True
                 ):

        # Percentile distance as definition of outlier gap
        self.dr_gap = dr_gap

        # Percentile definition of start/end of the data distribution
        self.r_zero = r_zero  # This depends on the number of elements, not use r_zero < 0.0001

        # Coefficient to multiply the estimated outlier distance
        self.k_gap = k_gap

        # Absolute path of the filename
        self.csv_hist = csv_hist

        # domain of the csv file to be processed ['total', 'S_polar','S_midlat',...]
        self.domain = domain

        # Other params
        self.verbose = verbose
        self.logarithmic_hist = logarithmic_hist

        # Read csv file
        data0 = pd.read_csv(self.csv_hist, index_col=0)
        self.data = data0.loc[self.domain]
        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.logarithmic_hist:
            self.nx = self.data.values[:-1] + 1
            self.nx = np.log10(self.nx)

    """
    Compute parameters
    """
    def set_parameters(self):

        # total number of elements
        self.N = np.sum(self.nx)

        # density distribution. percentiles
        self.quantiles = np.cumsum(self.nx) / self.N

        # 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]

    """
    Obtain quantiles from the histogram distribution
    """
    def quantiles_from_hist(self, quantiles):

        # 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)
        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):

        # Compute initial params
        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).sum() > 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
            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:
                    if self.verbose:
                        raise Exception('2-left: Estimated valid left range out of histogram limit :o. Extension of histogram left limit required')
                    break
        else:
            xm0_gap_left, xm1_gap_left = np.nan, np.nan
            if self.verbose:
                raise Exception('1-left: Observed range goes beyond histogram left limit :o. 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 == 0).sum() > 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
            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:
                    if self.verbose:
                        raise Exception('2-right: Estimated valid right range out of histogram limit :o. Extension of histogram right limit required.\n')
                    break

        else:
            xm0_gap_right, xm1_gap_right = np.nan, np.nan
            if self.verbose:
                raise Exception('1-right: Observed range goes beyond histogram right limit :o. Extension of histogram right limit required\n')

        return xm1_gap_left, xm0_gap_left, xm0_gap_right, xm1_gap_right

    """
    Get outer boundaries of the empty gaps.
    """
    def get(self):
        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,
                 percentile=np.arange(0, 105, 5),
                 dr_gap=np.arange(0.01, 0.1, 0.005),
                 r_zero=np.arange(0.01, 0.1, 0.005),
                 k_gap=[1, 5],
                 logarithmic_hist=[False, True],
                 domain='total',
                 ):
        self.percentile = percentile
        self.dr_gap = dr_gap
        self.r_zero = r_zero
        self.k_gap = k_gap
        self.domain = domain
        self.log_hist = logarithmic_hist

    '''
    This function computes the valid ranges for the file "file".

    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]    
    err indicates of the domain is shorter than required by the method    
    '''
    def get(self, csv_path):

        err = False
        Lp = len(self.percentile)

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

        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])

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

            vr = valid_range(csv_path, domain=self.domain, logarithmic_hist=log_l, verbose=False)

            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()

        if (np.isnan(vmin).sum() > 0) | (np.isnan(vmax).sum() > 0):
            err = True
            raise Exception('Variable range out of domain.')
        else:
            # 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):
                o_vmin[0, i_p] = (np.power(10, vr.nx[vr.xm < p_vmin[i_p]]) - 1).sum()
                o_vmax[0, i_p] = (np.power(10, vr.nx[vr.xm > p_vmax[i_p]]) - 1).sum()

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

        return (vmin, vmax), (p_vmin, p_vmax), (o_vmin, o_vmax), err

    '''
    This 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
        p_vmin = np.full((Lv, Lp), np.nan)
        p_vmax = np.full((Lv, Lp), np.nan)
        o_vmin = np.full((Lv, Lp), np.nan)
        o_vmax = np.full((Lv, Lp), np.nan)

        for i_v, file_i in enumerate(files):

            print(i_v, file_i)
            (vmin_i, vmax_i), (p_vmin_i, p_vmax_i), (o_vmin_i, o_vmax_i), err = self.get(file_i)

            if err:
                i_err += 1
                err_list.append('OutOfDomain')
            else:
                err_list.append('')

                # Valid range for the corresponding percentile threshold
                p_vmin[i_v, :] = p_vmin_i
                p_vmax[i_v, :] = p_vmax_i

                # Number of outlier for the corresponding percentile threshold
                o_vmin[i_v, :] = o_vmin_i
                o_vmax[i_v, :] = o_vmax_i

        if dataframes:
            p_vmin = pd.DataFrame(p_vmin, columns=['p' + str(i) for i in self.percentile], index=files)
            p_vmax = pd.DataFrame(p_vmax, columns=['p' + str(i) for i in self.percentile], index=files)
            o_vmin = pd.DataFrame(o_vmin, columns=['p' + str(i) for i in self.percentile], index=files)
            o_vmax = pd.DataFrame(o_vmax, columns=['p' + str(i) for i in self.percentile], index=files)
            err_list = pd.DataFrame(err_list, columns=['Error'], index=files)

        return (p_vmin, p_vmax), (o_vmin, o_vmax), err_list
