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

from pathlib import Path
import numpy as np
import pandas as pd
import logging
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


"""
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, index_col=0)
            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
