# -*- coding: utf-8 -*-
#
# This file is part of s4d.
#
# s4d is a python package for speaker diarization.
# Home page: http://www-lium.univ-lemans.fr/s4d/
#
# s4d is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 3 of the License,
# or (at your option) any later version.
#
# s4d is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with s4d.  If not, see <http://www.gnu.org/licenses/>.


"""
Copyright 2014-2021 Sylvain Meignier
"""

import logging
import numpy as np


class GaussFull:
    """
    Class to model a speaker by a gaussian with full covariance
    """
    def __init__(self, name, dim):
        self.logger = logging.getLogger(__name__)
        self.name = name
        self.count = 0
        self.dim = dim
        self.stat1 = np.zeros(dim)
        self.stat2 = np.zeros((dim, dim))
        self.cov_log_det = 0;
        self.mu = None
        self.mu_dot = np.NAN
        self.cov = None
        self.partial_bic = np.NaN

    def add(self, features):
        """
        Accumulate statistics for *features*
        :param features: numpy.ndarray

        """
        self.count += features.shape[0]  # add the number of features
        self.stat1 += features.sum(axis=0)
        self.stat2 += np.dot(features.T, features)

    def _cov_log_det(self):
        """
        Compute the log det of the covariance matrix
        :return:  float
        """
        s, d = np.linalg.slogdet(self.cov)
        return d
        # cov_chol, lower = scipy.linalg.cho_factor(self.cov)
        # return 2.0 * numpy.sum(numpy.log(numpy.diagonal(cov_chol)))

    def compute(self):
        """
        Compute the mean and covariance according the statistique, the log det
         of the covariance and the partial BIC :math:`PBIC`.

        :math:`PBIC_{x}  = \\frac{n_x}{2} \\log|\\Sigma_x|`

        """
        self.mu = self.stat1 / self.count
        tmp = self.mu[:, np.newaxis]
        self.mu_dot = np.dot(tmp, tmp.T)
        self.cov = self.stat2 / self.count - self.mu_dot
        self.cov_log_det = self._cov_log_det()
        self.partial_bic = self.cov_log_det * 0.5 * self.count

    @classmethod
    def merge(cls, m1, m2):
        """
        Merge two models *m1* and *m2*. Compute the new mean (*mu*),
        covariance (*cov*) and PBIC *partial_bic*

        :param m1: a GaussFull object
        :param m2: a GaussFull object
        :return: a GaussFull object
        """
        m = GaussFull(m1.name, m1.dim)
        m.count = m1.count + m2.count
        m.stat1 = m1.stat1 + m2.stat1
        m.stat2 = m1.stat2 + m2.stat2
        m.mu = m.stat1 / m.count
        tmp = m.mu[:, np.newaxis]
        m.cov = m.stat2 / m.count - np.dot(tmp, tmp.T)
        m.cov_log_det = m._cov_log_det()
        m.partial_bic = m.cov_log_det * 0.5 * m.count
        return m

    @classmethod
    def merge_partial_bic(cls, m1, m2):
        """
        Merge statistic accumulators of two a GaussFull objects, compute
        the new log det.

        :param m1: a GaussFull object
        :param m2: a GaussFull object
        :return: the log det
        """
        try:
            count = m1.count + m2.count
            mu = ((m1.stat1 + m2.stat1) / count)[:, np.newaxis]
            cov = (m1.stat2 + m2.stat2) / count - np.dot(mu, mu.T)
            s, d = np.linalg.slogdet(cov)
            # cov_chol, lower = scipy.linalg.cho_factor(cov, overwrite_a=True, check_finite=False)
            # d = 2.0 * numpy.sum(numpy.log(numpy.diagonal(cov_chol)))
            d *= 0.5 * count
            return d
        except:
            logging.warning('Det problem set to NaN ', m1.name, m2.name)
            return np.nan

    @classmethod
    def cst_bic(cls, dim, alpha):
        """
        Compute the BIC constant:

            :math:`cst  = \\frac{1}{2} \\alpha \\left(d + \\frac{d(d+1)}{2}\\right)`

        where :math:`d`is the feature dimension (*dim*)
        and :math:`alpha` a threshold (*alpha*)

        :param dim: the feature dimension
        :param alpha: the threshold
        :return: the constant
        """
        return 0.5 * alpha * (dim + (0.5 * ((dim + 1) * dim)))




