# -*- coding: utf-8 -*-
"""
qspec.stats
===========

Module for doing statistical analysis.
"""

import numpy as np
import scipy.stats as st
import scipy.integrate as si
from scipy.optimize import root, minimize
import matplotlib.pyplot as plt

from qspec.qtypes import *
from qspec import tools
from qspec.analyze import curve_fit

__all__ = ['Observable', 'add', 'mul', 'average', 'median', 'mode_lognormal', 'estimate_skewnorm', 'propagate',
           'propagate_fit', 'combined_pdf', 'relevant_interval', 'uniform', 'uniform_pumped', 'info']


class Observable(float):
    """
    A float object which has a 'label' and optionally both a left- and right-sided or a symmetric uncertainty.
    """
    def __new__(cls, x: Union[SupportsFloat, SupportsIndex, str, bytes, bytearray],
                std: scalar = None, std_2: scalar = None, label: str = None):
        return super(Observable, cls).__new__(cls, x)

    def __init__(self, x: Union[SupportsFloat, SupportsIndex, str, bytes, bytearray],
                 std: scalar = None, std_2: scalar = None, label: str = None):
        """
        :param x: The value of the observable.
        :param label: The label of the observable.
        :param std: The left-sided 1-sigma percentile relative to 'x' if 'std_2' is specified,
         else the standard deviation of the observable.
        :param std_2: The right-sided 1-sigma percentile relative to 'x' if 'std' is specified,
         else the standard deviation of the observable.
        """
        float.__init__(float(x))
        self.label = label
        self.std = std if std is not None else std_2
        self.std_2 = std_2 if std is not None else std
        self.popt: Union[array_iter, None] = None
        self.est: Union[array_iter, None] = None
        if std_2 is not None:
            try:
                self.est = estimate_skewnorm(float(self), self.std, self.std_2)
            except ValueError as e:
                tools.printw(repr(e))
                tools.printw('The real PDF might not be sufficiently described by a skew normal distribution.')

    def __repr__(self):
        return 'Observable({}, {}, {}, \'{}\')'.format(super().__repr__(), self.std, self.std_2, self.label)

    def __str__(self):
        if self.std is None:
            return '{}'.format(super().__repr__())
        elif self.std_2 is None:
            return '{}(+-{})'.format(super().__repr__(), self.std)
        else:
            return '{}(-{}+{})'.format(super().__repr__(), self.std, self.std_2)

    def __add__(self, other):
        return propagate(add, [self, other])

    def __radd__(self, other):
        return propagate(add, [other, self])

    def __mul__(self, other):
        return propagate(mul, [self, other])

    def __rmul__(self, other):
        return propagate(mul, [other, self])

    def __pow__(self, power, modulo=None):
        return propagate(pow, [self, power])

    def __rpow__(self, power, modulo=None):
        return propagate(pow, [power, self])

    def set_popt(self, popt: array_like = None):
        """
        :param popt: The new value of 'popt'.
        :returns: None. Sets the 'popt' attribute.
        """
        if popt is None:
            self.popt = popt
        else:
            self.popt = np.asarray(popt)

    def rvs(self, size: int = 1) -> ndarray:
        """
        :param size: The defining number of random variates (default is 1).
        :returns: Random variates of the observable of given 'size'. If asymmetric uncertainties are specified,
         a skew normal distribution is assumed. However, the ratio between the left- and right-sided uncertainty
         must not exceed 1.5.
        """
        if self.std_2 is None:
            return st.norm.rvs(loc=self, scale=0. if self.std is None else self.std, size=size)
        elif self.popt is not None:
            return st.skewnorm.rvs(*self.popt, size=size)
        elif self.est is not None:
            return st.skewnorm.rvs(*self.est, size=size)
        else:
            return np.full(size, np.nan)

    def pdf(self) -> Callable:
        """
        :returns: The probability distribution function (pdf) of the observable.
        """
        if self.std_2 is None:
            return np.zeros_like if self.std is None else st.norm(loc=self, scale=self.std).pdf
        elif self.popt is not None:
            return st.skewnorm(*self.popt).pdf
        elif self.est is not None:
            return st.skewnorm(*self.est).pdf
        else:
            return lambda x: np.full_like(x, np.nan)

    def hist(self, size: int = 1000000, n_bins: int = 200):
        """
        :param size: The defining number of random variates (default is 1,000,000).
        :param n_bins: The number of bins.
        :returns: None. Plots a histogram of the observable.
        """
        y = self.rvs(size)
        n, bins, _ = plt.hist(y, bins=n_bins, density=True, label='Sample data', facecolor='lightgray')
        bins = bins[:-1] + 0.5 * (bins[-1] - bins[0]) / n_bins
        plt.plot(bins, self.pdf()(bins), 'C0', label='Est. PDF')
        plt.legend(loc=1)
        plt.show()


def add(a, b):
    """
    :param a: The first summand.
    :param b: The second summand.
    :returns: The sum of the two summands.
    """
    return a + b


def mul(a, b):
    """
    :param a: The first factor.
    :param b: The second factor.
    :returns: The product of the two factors.
    """
    return a * b


def average(a: array_like, std: array_like = None, cov: array_iter = None, axis: int = None) -> (ndarray, ndarray):
    """
    :param a: The sample data.
    :param axis: The axis along which the average is computed.
    :param std: An array of standard deviations associated with the values in 'a'.
     If specified, the weighted average of the uncorrelated sample data and its standard error is returned.
    :param cov: The covariance matrix associated with the values in 'a'.
     If specified, it overrides 'std' and the weighted average of the correlated sample data
     and its standard error is returned. If no axis is specified, 'cov' must have shape (a.size, a.size)
     with elements associated with the values in the flattened 'a',
     otherwise cov must have either shape (a.shape[axis], a.shape[axis])
     or (..., a.shape[axis - 1], a.shape[axis], a.shape[axis], a.shape[axis + 1], ...).
    :returns: The average and its standard error for a given sample 'a' along the specified 'axis'.
    """
    a = np.asarray(a, dtype=float)
    if std is None and cov is None:  # The average and its standard error.
        av = np.average(a, axis=axis)
        av_d = np.std(a, axis=axis, ddof=1)
        if axis is None:  # The shape of 'a' is ignored ('a' is flattened).
            av_d /= np.sqrt(a.size)
        else:
            av_d /= np.sqrt(a.shape[axis])
    elif std is not None and cov is None:  # The weighted average of uncorrelated data and its standard error.
        std = np.asarray(std, dtype=float)
        av, sum_of_weights = np.average(a, axis=axis, weights=1. / (std ** 2), returned=True)
        av_d = np.sqrt(1. / sum_of_weights)
    else:  # The weighted average of correlated data and its standard error.
        cov = np.asarray(cov, dtype=float)
        if axis is None:  # The shape of 'a' is ignored ('a' is flattened).
            if cov.shape != (a.size, a.size):
                raise ValueError('Shape mismatch between \'a\' {} and \'cov\' {}.'
                                 .format(a.shape, cov.shape))
            av, sum_of_weights = np.average(a, weights=1. / np.diag(cov), returned=True)
            av_d = np.sum([1. / cov[i, i] + 2. * np.sum([cov[i, j] / (cov[i, i] * cov[j, j])
                                                         for j in range(a.size) if i < j])
                           for i in range(a.size)])
        else:
            if a.size != a.shape[0] and cov.shape == (a.size, a.size):
                raise ValueError('Shape mismatch between \'a\' {} and \'cov\' {} for specified axis {}.'
                                 .format(a.shape, cov.shape, axis))
            elif cov.shape == (a.shape[axis], a.shape[axis]):
                av, sum_of_weights = np.average(a, axis=axis, weights=1. / np.diag(cov), returned=True)
                av_d = np.sum([1. / cov[i, i] + 2. * np.sum([cov[i, j] / (cov[i, i] * cov[j, j])
                                                             for j in range(a.shape[axis]) if i < j])
                               for i in range(a.shape[axis])])
            else:
                var = np.diagonal(cov, axis1=axis, axis2=axis + 1)
                axes = [ax for ax in range(len(var.shape))]
                i = axes.pop(-1)
                axes.insert(axis, i)
                var = np.transpose(var, axes=axes)
                sum_of_weights = np.sum(1. / var, axis=axis)
                av = np.sum(a / var, axis=axis) / sum_of_weights
                av_d = np.sum(cov / (np.expand_dims(var, axis=axis + 1) * np.expand_dims(var, axis=axis)),
                              axis=(axis + 1, axis))
        av_d = np.sqrt(av_d) / sum_of_weights
    return av, av_d


def median(a: array_like, axis: int = None) -> (ndarray, ndarray, ndarray):
    """
    :param a: The sample data.
    :param axis: The axis along which the three percentiles are computed.
    :returns: The median (0.5-percentile) as well as the left- (~0.1587) and right-sided (~0.8413) 1-sigma percentile
     of a given sample 'a' along the specified 'axis'.
    """
    med = np.median(a, axis=axis)
    neg = np.percentile(a, 15.8655254, axis=axis)
    pos = np.percentile(a, 84.1344746, axis=axis)
    return med, neg, pos


def estimate_skewnorm(med: scalar, per_0: scalar, per_1: scalar):
    """
    :param med: The median (0.5-percentile) of a random variable.
    :param per_0: The left-sided 1-sigma percentile (~0.1587-percentile) relative to 'med'.
    :param per_1: The right-sided 1-sigma percentile (~0.8413-percentile) relative to 'med'.
    :returns: A size-3-array of the estimated parameters (alpha, mean, std.)
     of a skew normal distribution that matches the given percentiles.
    :raises ValueError: If the ratio between the left-(right-) and right-(left-)sided uncertainty exceeds 1.5.
    """
    if per_0 == 0 or per_1 == 0:
        return None
    if per_0 / per_1 > 1.5 or per_1 / per_0 > 1.5:
        raise ValueError('The ratio between the left-(right-) and right-(left-)sided uncertainty must not exceed 1.5.')
    per = np.array([med, med - per_0, med + per_1])
    y = np.array([0.5, 0.158655254, 0.841344746])

    def f(x):
        return st.skewnorm.cdf(per, *x) - y

    def df(x):
        xi0 = -np.exp(-0.5 * ((per - x[1]) / x[2]) ** 2 * (1. + x[0] ** 2)) / (np.pi * (1. + x[0] ** 2))
        xi1 = -st.skewnorm.pdf(per, *x)
        xi2 = -st.skewnorm.pdf(per, *x) * (per - x[1]) / x[2]
        return np.array([xi0, xi1, xi2]).T

    return root(f, x0=np.array([0., med, per_0]), jac=df).x


def lognormal(x, s=0., loc=0., scale=1., a=1.):
    return a * st.lognorm.pdf(x, s=s, loc=loc, scale=scale)


def mode_lognormal(x, bins=100):
    y, edges = np.histogram(x, bins=bins)
    delta = np.mean(edges[1:] - edges[:-1])
    x = edges[:-1] + 0.5 * delta

    loc = np.min(x) - 0.1 * (np.max(x) - np.min(x))
    scale = np.mean(x - loc)
    s = np.std(np.log(x - loc))
    a = scale * np.max(y)
    p0 = [s, loc, scale, a]

    popt, pcov = curve_fit(lognormal, x, y, p0=p0)
    mode = popt[1] + popt[2] * np.exp(-popt[0] ** 2)
    mode = p0[1] + p0[2] * np.exp(-p0[0] ** 2)

    plt.bar(x, y, width=delta)
    plt.plot(x, lognormal(x, *p0), '-C1')
    plt.plot(x, lognormal(x, *popt), '-C2')
    plt.vlines(mode, 0, np.max(y), colors='k')
    plt.show()
    return mode


def propagate(f: Callable, x: array_like, x_d: array_like = None, cov: array_iter = None,
              unc_places: int = None, sample_size: int = 1000000, rtol: float = 1e-3, atol: float = None,
              force_sym: bool = False, full_output: bool = False, show: bool = False) -> (Observable, list, ndarray):
    """
    :param f: The function to compute. 'f' needs to be vectorized.
    :param x: The input values. If 'x_d' is None, the sample data will be generated with the 'Observable.rvs' function
     which considers asymmetric uncertainties. If an element is not an 'Observable', its uncertainty is assumed to be 0.
    :param x_d: The uncertainties of the input values.
    :param cov: The covariance matrix of the x values. If not None, 'x' are assumed to be distributed according to
     a multivariate normal distribution with covariance 'cov'.
    :param unc_places: The number of significant decimal places the result will be rounded to.
     If None, the result is not rounded.
    :param sample_size: The number of random variates used for the calculation. The default is 1,000,000.
    :param rtol: The relative tolerance, with respect to the median of the resulting sample,
     with which the left- and right-sided uncertainties can deviate before asymmetric uncertainties are used.
    :param atol: The absolute tolerance with which the left- and right-sided uncertainties can deviate,
     before asymmetric uncertainties are used. Overrides 'rtol'.
    :param force_sym: Whether to force symmetric uncertainties. If so, a normal distribution is assumed.
    :param full_output: Whether to return the randomly generated data samples.
    :param show: Whether to show a histogram and estimated PDFs of the computed sample data.
    :returns: An 'Observable' whose uncertainties result from the propagation of the uncertainties
     of the input values 'x' by function 'f'.
     If the uncertainties are asymmetric, the parameters of a skew normal distribution
     are estimated using least-square fitting and are stored to the observable. The value and the two uncertainties are
     the median and the left- (~0.1587) and right-sided (~0.8413) 1-sigma percentiles
     relative to the median, respectively. If the uncertainties are symmetric the observable
     is assumed to be normally distributed. The value and the single uncertainty is then calculated using
     the mean and the standard deviation of the sampled data. If 'full_output' is True, a list of the input samples
     as well as the output sample are returned along with the observable.
    """
    label = f.__name__
    n_bins = 200
    if cov is None:
        if x_d is None:
            rand_x = [x_i.rvs(sample_size) if isinstance(x_i, Observable) else x_i for x_i in x]
        else:
            x = np.asarray(x, dtype=float)
            x_d = np.asarray(x_d, dtype=float)
            if x.size != x_d.size:
                raise ValueError('x and x_d must have the same size but have sizes {} and {}.'.format(x.size, x_d.size))
            rand_x = [x_i if x_d_i == 0. else st.norm.rvs(loc=x_i, scale=x_d_i, size=sample_size)
                      for x_i, x_d_i in zip(x, x_d)]
    else:
        cov = np.asarray(cov)
        rand_x = st.multivariate_normal.rvs(mean=[float(x_i) for x_i in x], cov=cov, size=sample_size).T
    rand_y = np.asarray(f(*rand_x))

    med, per_0, per_1 = median(rand_y)
    mean = np.average(rand_y)
    std = np.std(rand_y, ddof=1)

    plt.figure(num=1, dpi=96)
    n, bins, _ = plt.hist(rand_y, bins=n_bins, density=True, label='Sample data', facecolor='lightgray')
    bins = bins[:-1] + 0.5 * (bins[-1] - bins[0]) / n_bins

    if force_sym or (atol is None and abs(per_1 - per_0) / med <= rtol) \
            or (atol is not None and abs(per_1 - per_0) <= atol):
        if unc_places is not None:
            std, dec = tools.round_to_n(std, unc_places)
            mean = np.around(mean, decimals=dec)
        ob = Observable(mean, label=label, std=float(std))
    else:
        if unc_places is not None:
            per_0, dec_0 = tools.round_to_n(per_0, unc_places)
            per_1, dec_1 = tools.round_to_n(per_1, unc_places)
            dec = np.max([dec_0, dec_1])
            med = np.around(med, decimals=dec)
        ob = Observable(med, label=label, std=float(med - per_0), std_2=float(per_1 - med))
        try:
            popt, _ = curve_fit(st.skewnorm.pdf, bins, n, p0=[0., mean, std])
            ob.set_popt(popt)
        except RuntimeError:
            print('Optimal parameters for a skew normal distribution could not be found through fitting.'
                  ' Using an estimation if possible.')
    
    if show:
        print('Median(sample data): {} (-{} +{})'.format(med, med - per_0, per_1 - med))
        plt.plot(bins, st.norm.pdf(bins, loc=mean, scale=std), 'C0', label='Est. normal')
        if ob.popt is not None:
            med, per_0, per_1 = median(st.skewnorm.rvs(*ob.popt, size=sample_size))
            print('Median(Fit. skew normal): {} (-{} +{})'.format(med, med - per_0, per_1 - med))
            plt.plot(bins, st.skewnorm.pdf(bins, *ob.popt), 'C1', label='Fit. skew normal')
        if ob.est is not None:
            med, per_0, per_1 = median(st.skewnorm.rvs(*ob.est, size=sample_size))
            print('Median(Est. skew normal): {} (-{} +{})'.format(med, med - per_0, per_1 - med))
            plt.plot(bins, st.skewnorm.pdf(bins, *ob.est), 'C2', label='Est. skew normal')
        plt.legend(loc=1)
        plt.show()
    plt.close()

    if full_output:
        return ob, rand_x, rand_y
    else:
        return ob


def propagate_fit(f: Callable, x: array_like, popt: array_like, pcov: array_like, sample_size: int = 1000000):
    """
    :param f: The function to compute. $f$ needs to be vectorized.
    :param x: The input values.
    :param popt: An array of the fitted parameters (compare curve_fit).
    :param pcov: A 2d-array of the estimated covariances (compare curve_fit).
    :param sample_size: The number of random variates used for the calculation. The default is 1,000,000.
    :returns: The median and the $1\\sigma$ percentiles of the sampled function $f$.
    """
    x = np.asarray(x)
    _x = np.expand_dims(x, axis=-1)
    _popt = st.multivariate_normal.rvs(mean=popt, cov=pcov, size=sample_size).T
    _popt = [np.expand_dims(p, axis=tuple(range(len(x.shape)))) for p in _popt]
    med, per_0, per_1 = median(f(_x, *_popt), axis=-1)
    return med, per_0, per_1


def combined_pdf(z: array_like, pdf_1: Callable = st.norm.pdf, pdf_2: Callable = st.norm.pdf,
                 loc_1: float = 0., scale_1: float = 1., loc_2: float = 0., scale_2: float = 1.,
                 operator: str = '+', n: int = 10) -> ndarray:
    """
    :param z: The quantiles of the combined probability density function (pdf).
    :param pdf_1: The pdf of the first random variate.
    :param pdf_2: The pdf of the second random variate.
    :param loc_1: The 'loc' parameter of the first pdf.
    :param scale_1: The 'scale' parameter of the first pdf.
    :param loc_2: The 'loc' parameter of the second pdf.
    :param scale_2: The 'scale' parameter of the second pdf.
    :param operator: The operator that defines the new random variate given by Z = X <operator> Y.
     Currently supported operators are {'+', '*'}.
    :param n: The precision of the numerical integration.
     The integration uses 2 ** n intervals to evaluate the integral.
    :returns: The value of the pdf at the given 'z' quantiles.
    """
    arg = np.asarray(z)
    if arg.shape == ():
        arg = np.expand_dims(arg, axis=0)
    arg = np.expand_dims(arg, axis=-1)

    if operator == '*':
        left = loc_1 * loc_2 - 5. * (loc_1 * loc_2 - (loc_1 - scale_1) * (loc_2 - scale_2))
        right = loc_1 * loc_2 + 5. * (loc_1 * loc_2 - (loc_1 - scale_1) * (loc_2 - scale_2))

        def kernel(x):
            return pdf_1(x, loc=loc_1, scale=scale_1) * pdf_2(arg / x, loc=loc_2, scale=scale_2) / abs(x)

    elif operator == '+':
        left = loc_1 + loc_2 - 10. * np.sqrt(scale_1 ** 2 + scale_2 ** 2)
        right = loc_1 + loc_2 + 10. * np.sqrt(scale_1 ** 2 + scale_2 ** 2)

        def kernel(x):
            return pdf_1(x, loc=loc_1, scale=scale_1) * pdf_2(arg - x, loc=loc_2, scale=scale_2)

    else:
        raise ValueError('Operator \'{}\' not supported.'.format(operator))

    x_range = np.linspace(left, right, 2 ** n + 1)
    dx = x_range[1] - x_range[0]
    x_range = np.expand_dims(x_range, axis=0)
    y = kernel(x_range)
    result = si.romb(y, dx, axis=1)
    return result


def relevant_interval(dist: Callable, *args, show: bool = False, **kwargs):
    """
    :param dist: The probability distribution function (pdf).
    :param args: Additional arguments for the pdf.
    :param show: Whether to show the pdf in the determined interval.
    :param kwargs: Additional keyword arguments for the pdf.
    :returns: An estimation of the interval where most of the probability lies in.
    """
    k = 10
    scale = 1.

    def cost(_x):
        if _x[0] > _x[1]:
            _y = dist(np.linspace(_x[1], _x[0], 2 ** k + 1), *args, **kwargs)
        else:
            _y = dist(np.linspace(_x[0], _x[1], 2 ** k + 1), *args, **kwargs)
        c = (1. - si.romb(_y, abs(_x[1] - _x[0]) / (2 ** k))) * scale
        if show:
            print(_x, c / scale)
        return c

    def d_cost(_x):
        _y = np.array([dist(_x[0], *args, **kwargs), -dist(_x[1], *args, **kwargs)])
        if _x[0] > _x[1]:
            _y = -_y
        return _y * scale

    x0 = np.array([-1., 1.])
    while cost(x0) < 0.1:
        if show:
            print('x0: ', x0)
        x0 *= 0.5
    while cost(x0) > 0.9:
        if show:
            print('x0: ', x0)
        x0 *= 2.
    scale = x0[1]
    ret = minimize(cost, x0, jac=d_cost).x
    if show:
        x = np.linspace(*ret, 2 ** k + 1)
        y = dist(x, *args, **kwargs)
        plt.plot(x, y, label=dist.__name__)
        plt.legend()
        plt.show()
    return ret


""" Probability density functions """


def uniform(x: array_like, width: array_like) -> ndarray:
    """
    :param x: The x quantiles.
    :param width: The width of the uniform distribution.
    :returns: The probability density at 'x'.
    """
    return st.uniform.pdf(x, loc=-0.5 * width, scale=width)


def _uniform_pumped(x, width, gamma_u, depth) -> ndarray:
    """
    :param x: The x quantiles.
    :param width: The width of the uniform distribution.
    :param gamma_u: The fwhm of the pump transition.
    :param depth: The depth of the pump dip in parts of the height of the underlying uniform distribution.
    :returns: A not normalized uniform distribution with a Lorentz shaped dip.
    """
    x = np.asarray(x)
    scalar_true = tools.check_shape((), x, return_mode=True)
    if scalar_true:
        x = np.array([x])
    uni = uniform(x, width)
    nonzero = np.nonzero(uni)
    ret = np.zeros_like(x)
    ret[nonzero] = uni[nonzero] - depth / width * (np.pi * gamma_u / 2.) * st.cauchy.pdf(x[nonzero], scale=gamma_u / 2.)
    if scalar_true:
        return ret[0]
    return ret


def uniform_pumped(x, width, gamma_u, depth) -> ndarray:
    """
    :param x: The x quantiles.
    :param width: The width of the uniform distribution.
    :param gamma_u: The fwhm of the pump transition.
    :param depth: The depth of the pump dip in parts of the height of the underlying uniform distribution.
    :returns: A uniform distribution with a Lorentz-shaped dip.
    """
    _x = np.linspace(-0.5 * width, 0.5 * width, 2 ** 10 + 1)
    _y = _uniform_pumped(_x, width, gamma_u, depth)
    norm = si.romb(_y, width / (2 ** 10))
    return _uniform_pumped(x, width, gamma_u, depth) / norm


def info():
    """
    Plots the 1-sigma percentiles (and their ratio) of a skew normal distribution relative to its median
    for different values of the parameter alpha.
    The plot shows that the ratio between the two 1-sigma percentiles cannot exceed 1.5.

    :returns: None.
    """
    neg = []
    pos = []
    a = np.linspace(-10, 10, 1001)
    for val in a:
        dist = st.skewnorm.rvs(val, loc=0, scale=1, size=100000)
        med = np.median(dist)
        neg.append(med - np.percentile(dist, 15.8655254))
        pos.append(np.percentile(dist, 84.1344746) - med)
    neg, pos = np.array(neg), np.array(pos)
    plt.title('Properties of the skew normal distribution')
    plt.xlabel(r'Parameter $\alpha$')
    plt.ylabel(r'Percentiles ($\sigma$) / Ratio of Percentiles')
    plt.ylim(0., 1.75)
    plt.plot(a, neg, label='left 1-sigma')
    plt.plot(a, pos, label='right 1-sigma')
    plt.plot(a, pos / neg, label='Ratio right/left')
    plt.plot(a, neg / pos, label='Ratio left/right')
    plt.legend(loc='lower center')
    plt.show()
