import numpy as np
import scipy.sparse as sp


class Delta:
    """Class for deltas in MICE

    Parameters
    ----------
    x : array_like
        Point where to estimate the gradient .
    c : float
        Cost of each evaluation of this Delta object.
    x_l1 : array_like, default=None
        If set, this Delta object goal is to estimate the difference of the
        gradients evaluated at 'x' and 'x_l1', otherwise, this Delta will
        estimate the gradient at 'x'.
    m_min : int, default=5
        Minimum sample size for this Delta.

    Methods
    -------

    __call__()
        Prints all the information stored in the object.

    update_delta(mice, m)
        Updates Delta object to the new value for sample size.

    restart(mice)
        Generates a new instance of the Delta class with the bottom of the
        hierarchy at 'x' and 'c'=1.
    """

    def __init__(self, x, sampler, c=2, x_l1=None, m_min=5):
        self.x_l = x
        self.x_l1 = x_l1
        self.f_delta = np.array([]).reshape(0, len(x))
        self.f_l = np.array([]).reshape(0, len(x))
        self.f_delta_av = np.array(0.)
        self.v_l = None
        self.v_batch = None
        self.m = 0
        self.c = c
        self.m_min = m_min
        self.sampler = sampler

    def update_delta(self, mice, m):
        """
        Updates Delta object to the new value for sample size.

        Parameters
        ----------
        mice : object from class MICE
            Object from class MICE to which this Delta object belongs.
        m : int
            New sample size for this Delta.
        """
        if self.m < m:
            m_to_sample = np.ceil(m - self.m).astype('int')
            samples = self.sampler(m_to_sample)
            if self.c == 1:
                self.f_delta = np.vstack(
                    [self.f_delta, mice.grad(self.x_l, samples)])
                mice.counter += self.c * m_to_sample
                self.v_batch = mice.var(self.f_delta)
            else:
                new_f_l = mice.grad(self.x_l, samples)
                new_f_delta = new_f_l - mice.grad(self.x_l1, samples)
                self.f_l = np.vstack([self.f_l, new_f_l])
                self.f_delta = np.vstack([self.f_delta, new_f_delta])
                mice.counter += self.c * m_to_sample
                self.v_batch = mice.var(self.f_l)
            self.v_l = mice.var(self.f_delta)
            self.f_delta_av = mice.aggr(self.f_delta)
            self.m = len(self.f_delta)
        return

    def restart(self, mice):
        """
        Generates a new instance of the Delta class with the bottom of the
        hierarchy at 'x' and 'c'=1. All the evaluations in the current Delta
        object are passed to the new instance.

        Parameters
        ----------
        mice : object from class MICE
            Object from class MICE to which this Delta object belongs.

        Returns
        -------
        object : instance of Delta class
            Instance of Delta class representing a restart at 'x'.
        """
        new_delta = mice.create_delta(self.x_l, c=1)
        # new_delta = deepcopy(self)
        new_delta.x_l1 = None
        new_delta.f_delta = new_delta.f_l
        new_delta.f_delta_av = mice.aggr(new_delta.f_delta)
        new_delta.v_l = self.v_batch
        new_delta.v_batch = self.v_batch
        # new_delta.v_batch = self.v_batch
        # new_delta.f_l = self.f_l
        new_delta.m = self.m
        new_delta.m_min = mice.m_restart_min
        new_delta.c = 1
        new_delta.sampler = self.sampler
        return new_delta

    def __call__(self):
        """
        Prints all the information stored in the object.
        Returns
        -------

        """
        for key in self.__dict__.keys():
            print(f'{key}: {self.__dict__[key]}')


class LightDelta:
    """Class for low memory deltas in MICE

    Parameters
    ----------
    x : array_like
        Point where to estimate the gradient .
    c : float
        Cost of each evaluation of this Delta object.
    x_l1 : array_like, default=None
        If set, this Delta object goal is to estimate the difference of the
        gradients evaluated at 'x' and 'x_l1', otherwise, this Delta will
        estimate the gradient at 'x'.
    m_min : int, default=5
        Minimum sample size for this Delta.

    Methods
    -------

    __call__()
        Prints all the information stored in the object.

    update_delta(mice, m)
        Updates Delta object to the new value for sample size.

    restart(mice)
        Generates a new instance of the Delta class with the bottom of the
        hierarchy at 'x' and 'c'=1.
    """

    def __init__(self, x, sampler, c=2, x_l1=None, m_min=5):
        self.x_l = x
        self.x_l1 = x_l1
        self.f_l = np.array(0.)
        self.f_delta_av = np.array(0.)
        self.v_l = None
        self.v_batch = None
        self.m2_del = np.array(0.)
        self.m2_l = np.array(0.)
        self.m = 0
        self.c = c
        self.m_min = m_min
        self.sampler = sampler

    def update_delta(self, mice, m):
        """Updates Delta object to the new value for sample size.

        Parameters
        ----------
        mice : object from class MICE
            Object from class MICE to which this Delta object belongs.
        m : int
            New sample size for this Delta.
        """
        if self.m < m:
            m_to_sample = np.ceil(m - self.m).astype('int')
            samples = self.sampler(m_to_sample)
            if self.c == 1:
                new_values = mice.grad(self.x_l, samples)
                self.f_delta_av, self.m2_del = _new_update(
                    self.m, self.f_delta_av, self.m2_del, new_values)
                self.m += m_to_sample
                mice.counter += self.c * m_to_sample
                self.v_batch = self.m2_del / (self.m - 1)
            else:
                new_f_ls = mice.grad(self.x_l, samples)
                new_f_l1s = mice.grad(self.x_l1, samples)
                self.f_l, self.m2_l = _new_update(
                    self.m, self.f_l, self.m2_l, new_f_ls)
                self.f_delta_av, self.m2_del = _new_update(
                    self.m, self.f_delta_av, self.m2_del, new_f_ls - new_f_l1s)
                self.m += m_to_sample
                mice.counter += self.c * m_to_sample
                self.v_batch = self.m2_l / (self.m - 1)
            self.v_l = self.m2_del / (self.m - 1)
        return

    def restart(self, mice):
        """
        Generates a new instance of the LightDelta class with the bottom of the
        hierarchy at 'x' and 'c'=1. All the evaluations in the current
        LightDelta object are passed to the new instance.

        Parameters
        ----------
        mice : object from class MICE
            Object from class MICE to which this LightDelta object belongs.

        Returns
        -------
        object : instance of LightDelta class
            Instance of LightDelta class representing a restart at 'x'.
        """
        new_delta = mice.create_delta(self.x_l, c=1)
        # new_delta = deepcopy(self)
        new_delta.x_l1 = None
        new_delta.f_l = self.f_l
        new_delta.f_delta_av = self.f_l
        new_delta.v_l = self.v_batch
        new_delta.v_batch = self.v_batch
        new_delta.m2_del = self.m2_l
        new_delta.m2_l = self.m2_l
        new_delta.m = self.m
        new_delta.m_min = mice.m_restart_min
        new_delta.c = 1
        new_delta.sampler = self.sampler
        return new_delta

    def __call__(self):
        for key in self.__dict__.keys():
            print(f'{key}: {self.__dict__[key]}')


class DeltaResampling:
    """Class for low memory deltas in MICE using the resampling technique

    Parameters
    ----------
    x : array_like
        Point where to estimate the gradient .
    c : float
        Cost of each evaluation of this Delta object.
    x_l1 : array_like, default=None
        If set, this Delta object goal is to estimate the difference of the
        gradients evaluated at 'x' and 'x_l1', otherwise, this Delta will
        estimate the gradient at 'x'.
    re_part : int, default=2
        Number of partitions used in resampling.
    m_min : int, default=5
        Minimum sample size for this Delta.

    Methods
    -------

    __call__()
        Prints all the information stored in the object.

    update_delta(mice, m)
        Updates Delta object to the new value for sample size.

    restart(mice)
        Generates a new instance of the Delta class with the bottom of the
        hierarchy at 'x' and 'c'=1.
    """

    def __init__(self, x, sampler, c=2, x_l1=None, re_part=2, m_min=5):
        self.x_l = x
        self.x_l1 = x_l1
        self.f_l = np.array(0.)
        # self.f_ls = [0.0 for i in range(re_part)]
        self.f_ls = np.zeros((re_part, len(x)))
        self.f_delta_av = np.zeros((len(x)))
        # self.f_deltas = [0.0 for i in range(re_part)]
        self.f_deltas = np.zeros((re_part, len(x)))
        self.f_ms = [0 for _ in range(re_part)]
        self.v_l = None
        self.v_batch = None
        self.m2_del = np.array(0.)
        self.m2_l = np.array(0.)
        self.m = 0
        self.m_prev = 0
        self.c = c
        self.m_min = m_min
        self.sampler = sampler
        self.re_part = re_part

    def update_delta(self, mice, m):
        """Updates Delta object to the new value for sample size.

        Parameters
        ----------
        mice : object from class MICE
            Object from class MICE to which this Delta object belongs.
        m : int
            New sample size for this Delta.
        """
        if self.m < m:
            m_to_sample = np.ceil(m - self.m).astype('int')
            samples = self.sampler(m_to_sample)
            if self.c == 1:
                new_values = mice.grad(self.x_l, samples)
                self.f_delta_av, self.m2_del = _new_update(
                    self.m, self.f_delta_av, self.m2_del, new_values)
                ms = np.arange(self.m, self.m + m_to_sample)
                idxs = ms % self.re_part
                for idx in range(self.re_part):
                    mask = ~(idxs == idx)
                    m_ = self.f_ms[idx]
                    m_new = mask.sum()
                    self.f_deltas[idx] = (
                                                 self.f_deltas[idx] * m_
                                                 + np.sum(new_values[mask],
                                                          axis=0)) / (
                                                     m_ + m_new)
                    self.f_ms[idx] += m_new
                self.m += m_to_sample
                mice.counter += self.c * m_to_sample
                self.v_batch = self.m2_del / (self.m - 1)
                self.f_l = self.f_delta_av
                self.f_ls = self.f_deltas
            else:
                new_f_ls = mice.grad(self.x_l, samples)
                new_f_l1s = mice.grad(self.x_l1, samples)
                self.f_l, self.m2_l = _new_update(
                    self.m, self.f_l, self.m2_l, new_f_ls)
                new_f_deltas = new_f_ls - new_f_l1s
                self.f_delta_av, self.m2_del = _new_update(
                    self.m, self.f_delta_av, self.m2_del, new_f_deltas)
                ms = np.arange(self.m, self.m + m_to_sample)
                idxs = ms % self.re_part
                for idx in range(self.re_part):
                    mask = ~(idxs == idx)
                    m_ = self.f_ms[idx]
                    m_new = mask.sum()
                    self.f_deltas[idx] = (
                                                 self.f_deltas[idx] * m_
                                                 + np.sum(new_f_deltas[mask],
                                                          axis=0)) / (
                                                     m_ + m_new)
                    self.f_ls[idx] = (
                                             self.f_ls[idx] * m_
                                             + np.sum(new_f_ls[mask],
                                                      axis=0)) / (m_ + m_new)
                    self.f_ms[idx] += m_new
                self.m += m_to_sample
                mice.counter += self.c * m_to_sample
                self.v_batch = self.m2_l / (self.m - 1)
            self.v_l = self.m2_del / (self.m - 1)
        return

    def restart(self, mice):
        """
        Generates a new instance of the DeltaResampling class with the bottom
        of the hierarchy at 'x' and 'c'=1. All the evaluations in the current
        DeltaResampling object are passed to the new instance.

        Parameters
        ----------
        mice : object from class MICE
            Object from class MICE to which this DeltaResampling object belongs.

        Returns
        -------
        object : instance of DeltaResampling class
            Instance of DeltaResampling class representing a restart at 'x'.
        """
        new_delta = mice.create_delta(self.x_l, c=1)
        # new_delta = deepcopy(self)
        new_delta.x_l1 = None
        new_delta.f_delta_av = self.f_l
        new_delta.f_deltas = self.f_ls
        new_delta.f_l = new_delta.f_delta_av
        new_delta.f_ls = new_delta.f_deltas
        new_delta.f_ms = self.f_ms
        new_delta.v_l = self.v_batch
        new_delta.v_batch = self.v_batch
        new_delta.m2_l = self.m2_l
        new_delta.m2_del = self.m2_l
        new_delta.m = self.m
        new_delta.m_min = mice.m_restart_min
        new_delta.c = 1
        new_delta.sampler = self.sampler
        return new_delta

    def __call__(self):
        for key in self.__dict__.keys():
            print(f'{key}: {self.__dict__[key]}')


class DeltaSparse:
    """Class for low memory deltas in MICE using the resampling technique and
    sparsity

    Parameters
    ----------
    x : array_like
        Point where to estimate the gradient .
    c : float
        Cost of each evaluation of this Delta object.
    x_l1 : array_like, default=None
        If set, this Delta object goal is to estimate the difference of the
        gradients evaluated at 'x' and 'x_l1', otherwise, this Delta will
        estimate the gradient at 'x'.
    re_part : int, default=2
        Number of partitions used in resampling.
    m_min : int, default=5
        Minimum sample size for this Delta.

    Methods
    -------

    __call__()
        Prints all the information stored in the object.

    update_delta(mice, m)
        Updates Delta object to the new value for sample size.

    restart(mice)
        Generates a new instance of the Delta class with the bottom of the
        hierarchy at 'x' and 'c'=1.
    """

    def __init__(self, x, sampler, c=2, x_l1=None, re_part=2, m_min=5):
        self.x_l = x
        self.x_l1 = x_l1
        self.f_l = np.array(0.)
        self.dim = x.shape[1]
        # self.f_ls = [0.0 for i in range(re_part)]
        self.f_ls = np.zeros((re_part, self.dim))
        self.f_delta_av = sp.csr_matrix((1, self.dim))
        # self.f_deltas = [0.0 for i in range(re_part)]
        self.f_deltas = np.zeros((re_part, self.dim))
        # self.f_deltas = sp.lil_matrix((re_part, self.dim))
        self.f_ms = [0 for _ in range(re_part)]
        self.v_l = None
        self.v_batch = None
        self.m2_del = np.array(0.)
        self.m2_l = np.array(0.)
        self.m = 0
        self.m_prev = 0
        self.c = c
        self.m_min = m_min
        self.sampler = sampler
        self.re_part = re_part

    def update_delta(self, mice, m):
        """Updates Delta object to the new value for sample size.

        Parameters
        ----------
        mice : object from class MICE
            Object from class MICE to which this Delta object belongs.
        m : int
            New sample size for this Delta.
        """
        if self.m < m:
            m_to_sample = np.ceil(m - self.m).astype('int')
            samples = self.sampler(m_to_sample)
            if self.c == 1:
                new_values = mice.grad(self.x_l, samples)
                self.f_delta_av, self.m2_del = _sparse_update(
                    self.m, self.f_delta_av, self.m2_del, new_values)
                ms = np.arange(self.m, self.m + m_to_sample)
                idxs = ms % self.re_part
                for idx in range(self.re_part):
                    mask = ~(idxs == idx)
                    m_ = self.f_ms[idx]
                    m_new = mask.sum()
                    self.f_deltas[idx] = (
                        self.f_deltas[idx] * m_
                        + new_values[mask].sum(axis=0)) / (m_ + m_new)
                    # self.f_deltas[idx] += (new_values[mask].sum(
                    #     axis=0) - m_new * self.f_deltas[idx]) / (m_ + m_new)
                    self.f_ms[idx] += m_new
                self.m += m_to_sample
                mice.counter += self.c * m_to_sample
                self.v_batch = self.m2_del / (self.m - 1)
                self.f_l = self.f_delta_av
                self.f_ls = self.f_deltas
            else:
                new_f_ls = mice.grad(self.x_l, samples)
                new_f_l1s = mice.grad(self.x_l1, samples)
                self.f_l, self.m2_l = _sparse_update(
                    self.m, self.f_l, self.m2_l, new_f_ls)
                new_f_deltas = new_f_ls - new_f_l1s
                self.f_delta_av, self.m2_del = _sparse_update(
                    self.m, self.f_delta_av, self.m2_del, new_f_deltas)
                ms = np.arange(self.m, self.m + m_to_sample)
                idxs = ms % self.re_part
                for idx in range(self.re_part):
                    mask = ~(idxs == idx)
                    m_ = self.f_ms[idx]
                    m_new = mask.sum()
                    self.f_deltas[idx] = (
                                                 self.f_deltas[idx] * m_
                                                 + new_f_deltas[mask].sum(
                                             axis=0)) / (m_ + m_new)
                    self.f_ls[idx] = (
                                             self.f_ls[idx] * m_
                                             + new_f_ls[mask].sum(axis=0)) / (
                                                 m_ + m_new)
                    self.f_ms[idx] += m_new
                self.m += m_to_sample
                mice.counter += self.c * m_to_sample
                self.v_batch = self.m2_l / (self.m - 1)
            self.v_l = self.m2_del / (self.m - 1)
        return

    def restart(self, mice):
        """
        Generates a new instance of the DeltaResampling class with the bottom
        of the hierarchy at 'x' and 'c'=1. All the evaluations in the current
        DeltaResampling object are passed to the new instance.

        Parameters
        ----------
        mice : object from class MICE
            Object from class MICE to which this DeltaResampling object belongs.

        Returns
        -------
        object : instance of DeltaResampling class
            Instance of DeltaResampling class representing a restart at 'x'.
        """
        new_delta = mice.create_delta(self.x_l, c=1)
        # new_delta = deepcopy(self)
        new_delta.x_l1 = None
        new_delta.f_delta_av = self.f_l
        new_delta.f_deltas = self.f_ls
        new_delta.f_l = new_delta.f_delta_av
        new_delta.f_ls = new_delta.f_deltas
        new_delta.f_ms = self.f_ms
        new_delta.v_l = self.v_batch
        new_delta.v_batch = self.v_batch
        new_delta.m2_l = self.m2_l
        new_delta.m2_del = self.m2_l
        new_delta.m = self.m
        new_delta.m_min = mice.m_restart_min
        new_delta.c = 1
        new_delta.sampler = self.sampler
        return new_delta

    def __call__(self):
        for key in self.__dict__.keys():
            print(f'{key}: {self.__dict__[key]}')


class DeltaBayHess:
    """Class for low memory deltas in MICE for Bayesian Hessian inverse
    using the resampling technique.

    Parameters
    ----------
    x : array_like
        Point where to estimate the gradient .
    c : float
        Cost of each evaluation of this Delta object.
    x_l1 : array_like, default=None
        If set, this Delta object goal is to estimate the difference of the
        gradients evaluated at 'x' and 'x_l1', otherwise, this Delta will
        estimate the gradient at 'x'.
    re_part : int, default=2
        Number of partitions used in resampling.
    m_min : int, default=5
        Minimum sample size for this Delta.

    Methods
    -------

    __call__()
        Prints all the information stored in the object.

    update_delta(mice, m)
        Updates Delta object to the new value for sample size.

    restart(mice)
        Generates a new instance of the Delta class with the bottom of the
        hierarchy at 'x' and 'c'=1.
    """

    def __init__(self, x, sampler, c=2, x_l1=None, re_part=2, m_min=5):
        self.x_l = x
        self.x_l1 = x_l1
        self.f_l = np.array(0.)
        # self.f_ls = [0.0 for i in range(re_part)]
        self.f_ls = np.zeros((re_part, len(x)))
        self.f_delta_av = np.zeros((len(x)))
        # self.f_deltas = [0.0 for i in range(re_part)]
        self.f_deltas = np.zeros((re_part, len(x)))
        self.f_ms = [0 for _ in range(re_part)]
        self.v_l = None
        self.v_batch = None
        self.m2_del = np.zeros(len(x))
        self.m2_l = np.zeros(len(x))
        self.m = 0
        self.c = c
        self.m_min = m_min
        self.sampler = sampler
        self.re_part = re_part

    def update_delta(self, mice, m):
        """Updates Delta object to the new value for sample size.

        Parameters
        ----------
        mice : object from class MICE
            Object from class MICE to which this Delta object belongs.
        m : int
            New sample size for this Delta.
        """
        if self.m < m:
            m_to_sample = np.ceil(m - self.m).astype('int')
            samples = self.sampler(m_to_sample)

            # # test
            #
            # av = np.zeros(2)
            # # m2 = np.zeros(2)
            # m2 = 0
            # k = 0
            # num_batch = 1
            # for k in range(num_batch):
            #     # av, m2 = _new_update(k*1000/num_batch, av, m2, samples[k::num_batch])
            #     av, m2 = _update_bay(k*1000/num_batch, av, m2, samples[k::num_batch])
            #     # av, m2 = _new_update(k, av, m2, [samples[k]])
            #     # av, m2 = update_stats(k, av, m2, samples[k])
            #     # av, m2 = _update_bay(k, av, m2, samples[k::100])
            #     k += 1
            #
            # print(m2/999)
            # print(np.cov(samples.T, ddof=1))
            # print(np.diag(np.cov(samples.T, ddof=1)).sum())

            if self.c == 1:
                new_values = mice.grad(self.x_l, samples)
                self.f_delta_av, self.m2_del = _update_bay(
                    self.m, self.f_delta_av, self.m2_del, new_values)
                ms = np.arange(self.m, self.m + m_to_sample)
                idxs = ms % self.re_part
                for idx in range(self.re_part):
                    mask = ~(idxs == idx)
                    m_ = self.f_ms[idx]
                    m_new = mask.sum()
                    self.f_deltas[idx] = \
                        (self.f_deltas[idx] * m_ + np.sum(new_values[mask],
                                                          axis=0)) / (
                                    m_ + m_new)
                    self.f_ms[idx] += m_new
                self.m += m_to_sample
                mice.counter += self.c * m_to_sample
                self.v_batch = np.sum(self.m2_del) / (self.m - 1)
                self.f_l = self.f_delta_av
                self.f_ls = self.f_deltas
            else:
                new_f_ls = mice.grad(self.x_l, samples)
                new_f_l1s = mice.grad(self.x_l1, samples)
                self.f_l, self.m2_l = _update_bay(
                    self.m, self.f_l, self.m2_l, new_f_ls)
                new_f_deltas = new_f_ls - new_f_l1s
                self.f_delta_av, self.m2_del = _update_bay(
                    self.m, self.f_delta_av, self.m2_del, new_f_deltas)
                ms = np.arange(self.m, self.m + m_to_sample)
                idxs = ms % self.re_part
                for idx in range(self.re_part):
                    mask = ~(idxs == idx)
                    m_ = self.f_ms[idx]
                    m_new = mask.sum()
                    self.f_deltas[idx] = \
                        (self.f_deltas[idx] * m_ + np.sum(new_f_deltas[mask],
                                                          axis=0)) / (
                                    m_ + m_new)
                    self.f_ls[idx] = \
                        (self.f_ls[idx] * m_ + np.sum(new_f_ls[mask],
                                                      axis=0)) / (m_ + m_new)
                    self.f_ms[idx] += m_new
                self.m += m_to_sample
                mice.counter += self.c * m_to_sample
                self.v_batch = np.sum(self.m2_l) / (self.m - 1)
            self.v_l = np.sum(self.m2_del) / (self.m - 1)
        return

    def restart(self, mice):
        """
        Generates a new instance of the DeltaBayHess class with the bottom of
        the hierarchy at 'x' and 'c'=1. All the evaluations in the current
        DeltaBayHess object are passed to the new instance.

        Parameters
        ----------
        mice : object from class MICE
            Object from class MICE to which this DeltaBayHess object belongs.

        Returns
        -------
        object : instance of DeltaBayHess class
            Instance of DeltaBayHess class representing a restart at 'x'.
        """
        new_delta = mice.create_delta(self.x_l, c=1)
        # new_delta = deepcopy(self)
        new_delta.x_l1 = None
        new_delta.f_delta_av = self.f_l
        new_delta.f_deltas = self.f_ls
        new_delta.f_l = new_delta.f_delta_av
        new_delta.f_ls = new_delta.f_deltas
        new_delta.f_ms = self.f_ms
        new_delta.v_l = self.v_batch
        new_delta.v_batch = self.v_batch
        new_delta.m2_l = self.m2_l
        new_delta.m2_del = self.m2_l
        new_delta.m = self.m
        new_delta.m_min = mice.m_restart_min
        new_delta.c = 1
        new_delta.sampler = self.sampler
        return new_delta

    def __call__(self):
        for key in self.__dict__.keys():
            print(f'{key}: {self.__dict__[key]}')


def _new_update(count, mean, m2, new_values):
    counts = np.arange(count + 1, count + len(new_values) + 1)
    size = (len(new_values[0]), 1)
    d_mean = (mean * count + np.cumsum(new_values, axis=0)) / \
             np.tile(counts, size).T
    # D_m = np.vstack(mean, d_mean)
    factors = np.divide(counts, (counts - 1), where=(counts > 1))
    m2 += ((np.linalg.norm(new_values - d_mean, axis=1)) ** 2 * factors).sum()
    return d_mean[-1], m2


def _sparse_update(count, mean, m2, new_values):
    m2_ = m2 * (count - 1)
    for new_value in new_values:
        count += 1
        diff = new_value - mean
        mean += diff / count
        diff2 = new_value - mean
        m2_ += (diff @ diff2.T).todense()[0, 0]
    return mean, m2_ / (count - 1)


def _sparse_cum_mean(mean, count, new_values):
    for datum in new_values:
        mean += new_values * (count / (count + 1))
        count += 1
    return


def _update_bay(count, mean, m2, new_values):
    counts = np.arange(count + 1, count + len(new_values) + 1)
    size = (len(new_values[0]), 1)
    d_mean = (mean * count + np.cumsum(new_values, axis=0)) \
             / np.tile(counts, size).T
    # D_m = np.vstack(mean, d_mean)
    factors = np.divide(counts, (counts - 1), where=(counts > 1))
    # m2 += ((np.linalg.norm(new_values-d_mean, axis=1))**2*factors).sum()
    m2 += np.sum((new_values - d_mean) ** 2 * np.tile(factors, size).T, axis=0)
    # set_trace()
    return d_mean[-1], m2
