#!/usr/bin/env python

# Copyright 2013-2019 Lukas Burget, Mireia Diez (burget@fit.vutbr.cz, mireia@fit.vutbr.cz)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#
# Revision History
#   16/07/13 01:00AM - original version
#   20/06/17 12:07AM - np.asarray replaced by .toarray()
#                    - minor bug fix in initializing q(Z)
#                    - minor bug fix in ELBO calculation
#                    - few more optimizations
#   03/10/19 02:27PM - speaker regularization coefficient Fb added
#

import numpy as np
from scipy.sparse import coo_matrix
import scipy.linalg as spl
import numexpr as ne  # the dependency on this modul can be avoided by replacing

# logsumexp_ne and exp_ne with logsumexp and np.exp

# [gamma pi Li] =
def VB_diarization(
    X,
    m,
    invSigma,
    w,
    V,
    pi=None,
    gamma=None,
    maxSpeakers=10,
    maxIters=10,
    epsilon=1e-4,
    loopProb=0.99,
    statScale=1.0,
    alphaQInit=1.0,
    downsample=None,
    VtinvSigmaV=None,
    ref=None,
    plot=False,
    sparsityThr=0.001,
    llScale=1.0,
    minDur=1,
    Fa=1.0,
    Fb=1.0,
):

    """
    This a generalized version of speaker diarization described in:

    Diez. M., Burget. L., Landini. F., Cernocky. J.
    Analysis of Speaker Diarization based on Bayesian HMM with Eigenvoice Priors

    Variable names and equation numbers refer to those used the paper

    Inputs:
    X  - T x D array, where columns are D dimensional feature vectors for T frames
    m  - C x D array of GMM component means
    invSigma - C x D array of GMM component inverse covariance matrix diagonals
    w  - C dimensional column vector of GMM component weights
    V  - R x C x D array of eigenvoices
    maxSpeakers - maximum number of speakers expected in the utterance
    maxIters    - maximum number of algorithm iterations
    epsilon     - stop iterating, if obj. fun. improvement is less than epsilon
    loopProb    - probability of not switching speakers between frames
    statScale   - deprecated, use Fa instead
    Fa          - scale sufficient statiscits collected using UBM
    Fb          - speaker regularization coefficient Fb (controls final # of speaker)
    llScale     - scale UBM likelihood (i.e. llScale < 1.0 make atribution of
                  frames to UBM componets more uncertain)
    sparsityThr - set occupations smaller that this threshold to 0.0 (saves memory
                  as the posteriors are represented by sparse matrix)
    alphaQInit  - Dirichlet concentration parameter for initializing gamma
    downsample  - perform diarization on input downsampled by this factor
    VtinvSigmaV       - C x (R**2+R)/2 matrix normally calculated by VB_diarization when
                  VtinvSigmaV is None. However, it can be pre-calculated using function
                  precalculate_VtinvSigmaV(V) and used across calls of VB_diarization.
    minDur      - minimum number of frames between speaker turns imposed by linear
                  chains of HMM states corresponding to each speaker. All the states
                  in a chain share the same output distribution
    ref         - T dim. integer vector with reference speaker ID (0:maxSpeakers)
                  per frame
    plot        - if set to True, plot per-frame speaker posteriors.

     Outputs:
     gamma  - S x T matrix of posteriors attribution each frame to one of S possible
          speakers, where S is given by opts.maxSpeakers
     pi - S dimensional column vector of ML learned speaker priors. Ideally, these
          should allow to estimate # of speaker in the utterance as the
          probabilities of the redundant speaker should converge to zero.
     Li - values of auxiliary function (and DER and frame cross-entropy between gamma
          and reference if 'ref' is provided) over iterations.
    """

    # The references to equations corresponds to
    # Diez. M., Burget. L., Landini. F., Cernocky. J.
    # Analysis of Speaker Diarization based on Bayesian HMM with Eigenvoice Priors

    D = X.shape[1]  # feature dimensionality
    C = len(w)  # number of mixture components
    R = V.shape[0]  # subspace rank
    nframes = X.shape[0]

    if VtinvSigmaV is None:
        VtinvSigmaV = precalculate_VtinvSigmaV(V, invSigma)

    V = V.reshape(V.shape[0], -1)

    if pi is None:
        pi = np.ones(maxSpeakers) / maxSpeakers
    else:
        maxSpeakers = len(pi)

    if gamma is None:
        # initialize gamma from flat Dirichlet prior with concentration parameter alphaQInit
        gamma = np.random.gamma(alphaQInit, size=(nframes, maxSpeakers))
        gamma = gamma / gamma.sum(1, keepdims=True)

    # calculate UBM mixture frame posteriors (i.e. per-frame zero order statistics)
    ll = (
        (X ** 2).dot(-0.5 * invSigma.T)
        + X.dot(invSigma.T * m.T)
        - 0.5
        * (
            (invSigma * m ** 2 - np.log(invSigma)).sum(1)
            - 2 * np.log(w)
            + D * np.log(2 * np.pi)
        )
    )
    ll *= llScale
    G = logsumexp_ne(ll, axis=1)
    zeta = exp_ne(ll - G[:, np.newaxis])
    zeta[zeta < sparsityThr] = 0.0
    zeta = zeta * statScale
    G = G * statScale

    # Kx = np.sum(zeta * (np.log(w) - np.log(zeta)), 1)
    zeta = coo_matrix(zeta)  # represent zero-order stats using sparse matrix
    print("Sparsity: ", len(zeta.row), float(len(zeta.row)) / np.prod(zeta.shape))
    LL = np.sum(G)  # total log-likelihod as calculated using UBM

    mixture_sum = coo_matrix((np.ones(C * D), (np.repeat(range(C), D), range(C * D))))

    # G = np.sum((zeta.multiply(ll - np.log(w))).toarray(), 1) + Kx  # from eq. (30) # Aleready calculated above

    # Calculate per-frame first order statistics projected into the R-dim. subspace
    # V^T \Sigma^{-1} F_m
    F_s = coo_matrix(
        (
            ((X[zeta.row] - m[zeta.col]) * zeta.data[:, np.newaxis]).flat,
            (
                zeta.row.repeat(D),
                zeta.col.repeat(D) * D + np.tile(range(D), len(zeta.col)),
            ),
        )
    )
    rho = F_s.tocsr().dot((invSigma.flat * V).T)
    del F_s
    ## The code above is only efficient implementation of the following comented code
    # rho = 0;
    # for ii in range(C):
    #  rho = rho + V[ii*D:(ii+1)*D,:].T.dot(zeta[ii,:] * invSigma[:,[ii]] *  (X - m[:,[ii]]))

    if downsample is not None:
        # Downsample zeta, rho, G and gamma by summing the statistic over 'downsample' frames
        # This speeds-up diarization for the price of lowering its frame resolution
        # downsampler = coo_matrix((np.ones(nframes), (np.ceil(np.arange(nframes)/downsample).astype(int), np.arange(nframes))))
        downsampler = coo_matrix(
            (
                np.ones(nframes),
                (
                    np.ceil(np.arange(nframes) / downsample).astype(int),
                    np.arange(nframes),
                ),
            )
        )
        zeta = downsampler.dot(zeta)
        rho = downsampler.dot(rho)
        G = downsampler.dot(G)
        gamma = downsampler.dot(gamma) / downsample
    else:
        downsampler = np.array(1)

    Li = [[LL * Fa]]  # for the 0-th iteration,
    if ref is not None:
        Li[-1] += [
            DER(downsampler.T.dot(gamma), ref),
            DER(downsampler.T.dot(gamma), ref, xentropy=True),
        ]

    ln_p = np.zeros_like(gamma)
    tr = np.eye(minDur * maxSpeakers, k=1)
    ip = np.zeros(minDur * maxSpeakers)
    for ii in range(maxIters):
        ELBO = 0  # objective function (11) (i.e. VB lower-bound on the evidence)
        sum_gamma_zeta = zeta.T.dot(
            gamma
        ).T  # corresponds to the last sum in eq. (26) for all 's'
        invLnoI_flat = sum_gamma_zeta.astype(VtinvSigmaV.dtype).dot(
            VtinvSigmaV
        )  # eq. (26) except for 'I' and the F_A F_B factors for all 's'
        sum_gamma_rho = gamma.T.dot(rho)  # summation in eq. (17)
        for sid in range(maxSpeakers):
            invL = np.linalg.inv(
                np.eye(R) + tril_to_sym(invLnoI_flat[sid]) * Fa / Fb
            )  # eq. (18) inverse
            a = invL.dot(sum_gamma_rho[sid]) * Fa / Fb  # eq. (17)
            ln_p[:, sid] = Fa * (
                G
                + rho.dot(a)
                - 0.5
                * zeta.dot(
                    mixture_sum.dot(
                        (
                            (invL + np.outer(a, a)).astype(V.dtype).dot(V)
                            * (invSigma.flat * V)
                        ).sum(0)
                    )
                )
            )  # eq. (23)
            ELBO += Fb * 0.5 * (logdet(invL) - np.sum(np.diag(invL) + a ** 2, 0) + R)

        # Construct transition probability matrix with linear chain of 'minDur'
        # states for each of 'maxSpeaker' speaker. The last state in each chain has
        # self-loop probability 'loopProb' and the transition probabilities to the
        # initial chain states given by vector '(1-loopProb) * pi'. From all other,
        # states, one must move to the next state in the chain with probability one.
        tr[minDur - 1 :: minDur, 0::minDur] = (1 - loopProb) * pi
        tr[(np.arange(1, maxSpeakers + 1) * minDur - 1,) * 2] += loopProb
        ip[::minDur] = pi
        # per-frame HMM state posteriors. Note that we can have linear chain of minDur states
        # for each speaker.
        gamma, tll, lf, lb = forward_backward(
            ln_p.repeat(minDur, axis=1), tr, ip
        )  # , np.arange(1,maxSpeakers+1)*minDur-1)

        # Right after updating q(Z), tll is E{log p(X|,Y,Z)} - KL{q(Z)||p(Z)}.
        # ELBO now contains -KL{q(Y)||p(Y)}. Therefore, ELBO+ttl is correct value for ELBO.
        ELBO += tll
        Li.append([ELBO])

        # ML estimate of speaker prior probabilities, eq. (24)
        pi = gamma[0, ::minDur] + np.exp(
            logsumexp(lf[:-1, minDur - 1 :: minDur], axis=1)[:, np.newaxis]
            + lb[1:, ::minDur]
            + ln_p[1:]
            + np.log((1 - loopProb) * pi)
            - tll
        ).sum(0)
        pi = pi / pi.sum()

        # per-frame speaker posteriors (eq. (19)), obtained by summing
        # HMM state posteriors corresponding to each speaker
        gamma = gamma.reshape(len(gamma), maxSpeakers, minDur).sum(axis=2)

        # if reference is provided, report DER, cross-entropy and plot the figures
        if ref is not None:
            Li[-1] += [
                DER(downsampler.T.dot(gamma), ref),
                DER(downsampler.T.dot(gamma), ref, xentropy=True),
            ]

            if plot:
                import matplotlib.pyplot

                if ii == 0:
                    matplotlib.pyplot.clf()
                matplotlib.pyplot.subplot(maxIters, 1, ii + 1)
                matplotlib.pyplot.plot(downsampler.T.dot(gamma), lw=2)
                matplotlib.pyplot.imshow(
                    np.atleast_2d(ref),
                    interpolation="none",
                    aspect="auto",
                    cmap=matplotlib.pyplot.cm.Pastel1,
                    extent=(0, len(ref), -0.05, 1.05),
                )
            print(ii, Li[-2])

        if ii > 0 and ELBO - Li[-2][0] < epsilon:
            if ELBO - Li[-1][0] < 0:
                print("WARNING: Value of auxiliary function has decreased!")
            break

    if downsample is not None:
        # upsample resulting gamma to match number of frames in the input utterance
        gamma = downsampler.T.dot(gamma)

    return gamma, pi, Li


def precalculate_VtinvSigmaV(V, invSigma):
    tril_ind = np.tril_indices(V.shape[0])
    VtinvSigmaV = np.empty((V.shape[1], len(tril_ind[0])), V.dtype)
    for c in range(V.shape[1]):
        VtinvSigmaV[c, :] = np.dot(
            V[:, c, :] * invSigma[np.newaxis, c, :], V[:, c, :].T
        )[tril_ind]
    return VtinvSigmaV


# Initialize gamma (per-frame speaker posteriors) from a reference
# (vector of per-frame zero based integer speaker IDs)
def frame_labels2posterior_mx(labels):
    # initialize from reference
    pmx = np.zeros((len(labels), labels.max() + 1))
    pmx[np.arange(len(labels)), labels] = 1
    return pmx


# Calculates Diarization Error Rate (DER) or per-frame cross-entropy between
# reference (vector of per-frame zero based integer speaker IDs) and gamma
# (per-frame speaker posteriors). If expected=False, gamma is converted into
# hard labels before calculating DER. If expected=TRUE, posteriors in gamma
# are used to calculate "expected" DER.
def DER(gamma, ref, expected=True, xentropy=False):
    from itertools import permutations

    if not expected:
        # replace probabiities in gamma by zeros and ones
        hard_labels = gamma.argmax(1)
        gamma = np.zeros_like(gamma)
        gamma[range(len(gamma)), hard_labels] = 1

    err_mx = np.empty((ref.max() + 1, gamma.shape[1]))
    for s in range(err_mx.shape[0]):
        tmpgamma = gamma[ref == s, :]
        err_mx[s] = (-np.log(tmpgamma) if xentropy else tmpgamma).sum(0)

    if err_mx.shape[0] < err_mx.shape[1]:
        err_mx = err_mx.T

    # try all alignments (permutations) of reference and detected speaker
    # could be written in more efficient way using dynamic programing
    acc = [
        err_mx[perm[: err_mx.shape[1]], range(err_mx.shape[1])].sum()
        for perm in permutations(range(err_mx.shape[0]))
    ]
    if xentropy:
        return min(acc) / float(len(ref))
    else:
        return (len(ref) - max(acc)) / float(len(ref))


###############################################################################
# Module private functions
###############################################################################
def logsumexp(x, axis=0):
    xmax = x.max(axis)
    x = xmax + np.log(np.sum(np.exp(x - np.expand_dims(xmax, axis)), axis))
    infs = np.isinf(xmax)
    if np.ndim(x) > 0:
        x[infs] = xmax[infs]
    elif infs:
        x = xmax
    return x


# The following two functions are only versions optimized for speed using numexpr
# module and can be replaced by logsumexp and np.exp functions to avoid
# the dependency on the module.
def logsumexp_ne(x, axis=0):
    xmax = np.array(x).max(axis=axis)
    xmax_e = np.expand_dims(xmax, axis)
    x = ne.evaluate("sum(exp(x - xmax_e), axis=%d)" % axis)
    x = ne.evaluate("xmax + log(x)")
    infs = np.isinf(xmax)
    if np.ndim(x) > 0:
        x[infs] = xmax[infs]
    elif infs:
        x = xmax
    return x


def exp_ne(x, out=None):
    return ne.evaluate("exp(x)", out=None)


# Convert vector with lower-triangular coefficients into symetric matrix
def tril_to_sym(tril):
    R = np.sqrt(len(tril) * 2).astype(int)
    tril_ind = np.tril_indices(R)
    S = np.empty((R, R))
    S[tril_ind] = tril
    S[tril_ind[::-1]] = tril
    return S


def logdet(A):
    return 2 * np.sum(np.log(np.diag(spl.cholesky(A))))


def forward_backward(lls, tr, ip):
    """
    Inputs:
        lls - matrix of per-frame log HMM state output probabilities
        tr  - transition probability matrix
        ip  - vector of initial state probabilities (i.e. statrting in the state)
    Outputs:
        sp  - matrix of per-frame state occupation posteriors
        tll - total (forward) log-likelihood
        lfw - log forward probabilities
        lfw - log backward probabilities
    """
    ltr = np.log(tr)
    lfw = np.empty_like(lls)
    lbw = np.empty_like(lls)
    lfw[:] = -np.inf
    lbw[:] = -np.inf
    lfw[0] = lls[0] + np.log(ip)
    lbw[-1] = 0.0

    for ii in range(1, len(lls)):
        lfw[ii] = lls[ii] + logsumexp(lfw[ii - 1] + ltr.T, axis=1)

    for ii in reversed(range(len(lls) - 1)):
        lbw[ii] = logsumexp(ltr + lls[ii + 1] + lbw[ii + 1], axis=1)

    tll = logsumexp(lfw[-1])
    sp = np.exp(lfw + lbw - tll)
    return sp, tll, lfw, lbw
