import numpy as np

import cleanlog

from scipy.special import gammaln as logG
from scipy.special import digamma
from scipy.special import polygamma
from sklearn.mixture import GaussianMixture

logger = cleanlog.ColoredLogger(name='PRISM')

def trigamma(x):
    return polygamma(1, x)


class BetaBinomialMixture():
    def __init__(self, n_components=1, max_iter=10000, tol=1e-3, seed=None, verbose=False):
        self.n_components = n_components
        self.max_iter = max_iter
        self.tol = tol
        self.seed = seed
        self.verbose = verbose

    def _bebin_mle(self, ns, ks, ws, n_iter=4):
        """Given depths(ns), and counts(ks) find the most likely alpha and beta,
        which are the two parameters of beta-binomial distribution."""
        a, b = 1, 1
        N = ws.sum()
        p1_bar = np.exp(1 / N * np.sum(ws * (digamma(a + ks) - digamma(a + b + ns))))
        p2_bar = np.exp(1 / N * np.sum(ws * (digamma(b + (ns - ks)) - digamma(a + b + ns))))

        # Find good initial esimates.
        a = 1 / 2 * (1 - p2_bar) / (1 - p1_bar - p2_bar)
        b = 1 / 2 * (1 - p1_bar) / (1 - p1_bar - p2_bar)

        for _ in range(n_iter):
            p1_bar = np.exp(1 / N * np.sum(ws * (digamma(a + ks) - digamma(a + b + ns))))
            p2_bar = np.exp(1 / N * np.sum(ws * (digamma(b + (ns - ks)) - digamma(a + b + ns))))

            n_log_p1 = np.sum(ws * (digamma(a + ks) - digamma(a + b + ns)))
            n_log_p2 = np.sum(ws * (digamma(b + (ns - ks)) - digamma(a + b + ns)))

            q_1 = -N * trigamma(a)
            q_2 = -N * trigamma(b)
            g_1 = N * digamma(a + b) - N * digamma(a) + n_log_p1
            g_2 = N * digamma(a + b) - N * digamma(b) + n_log_p2
            z = N * trigamma(a + b)
            b_ = (g_1 / q_1 + g_2 / q_2) / (1 / z + 1 / q_1 + 1 / q_2)

            a = a - (g_1 - b_) / q_1
            b = b - (g_2 - b_) / q_2

        return a, b

    def _bebin_likelihood(self, n, k, a, b):
        return np.exp(self._bebin_loglikelihood(n, k, a, b))

    def _bebin_loglikelihood(self, n, k, a, b):
        tmp = logG(n + 1) + logG(k + a) + logG(n - k + b) + logG(a + b) - \
            (logG(k + 1) + logG(n - k + 1) + logG(a) + logG(b) + logG(n + a + b))
        return tmp.sum(axis=1)

    def _gmm_initialize(self, n, k):
        """Initialize alphas and betas by fitting gaussian mixture model roughly.
        Alphas and betas are correspondingly computed from means and variances of each component.
        """
        # Ratio of methylated reads.
        r = k / n
        r = r.reshape(self.n_data, self.n_dim)

        # Fit gaussian mixture model. It assumes that each dimension has its own variance.
        model = GaussianMixture(n_components=self.n_components, covariance_type='diag')
        model.fit(r)

        alphas, betas = [], []
        for i in range(self.n_components):
            # Compute alpha and beta from mean and variance of GMM fit.
            mu = model.means_[i]
            var = model.covariances_[i]

            alpha = ((1 - mu) / var - 1 / mu) * mu**2
            beta = alpha * (1 / mu - 1)

            alphas.append(alpha)
            betas.append(beta)
        return np.array(alphas), np.array(betas)

    def fit(self, n, k, headers):
        if self.seed:
            np.random.seed(self.seed)
        
        self.n_data, self.n_dim = len(n), 1 if n.ndim == 1 else len(n[0])

        # If 1d data are given, reshape them into 2d matrix.
        if self.n_dim == 1:
            n = n.reshape(self.n_data, self.n_dim)
            k = k.reshape(self.n_data, self.n_dim)

        self.depths, self.counts, self.headers = n, k, headers
        self.alphas_, self.betas_ = self._gmm_initialize(n, k)
        self.pi_ = np.ones(self.n_components) / self.n_components

        prev_weighted_loglikelihood, self.converged_ = np.inf, False

        for iteration in range(self.max_iter):
            # Compute likelihood of alpha and beta w.r.t. each data point.
            l = np.array([self._bebin_likelihood(n, k, a, b) for a, b in zip(self.alphas_, self.betas_)])

            # Compute posterior.
            weighted_likelihood = l * self.pi_.reshape([self.n_components, 1])
            curr_weighted_loglikelihood = np.log(weighted_likelihood.sum(axis=0)).sum()

            # Check for convergence.
            if np.abs(curr_weighted_loglikelihood - prev_weighted_loglikelihood) < self.tol:
                if self.verbose:
                    logger.debug(f'Met convergence criterion at iteration {iteration}. Terminating.')
                self.converged_ = True
                break
            prev_weighted_loglikelihood = curr_weighted_loglikelihood

            w = weighted_likelihood / weighted_likelihood.sum(axis=0)

            # Compute relative size of clusters.
            self.pi_ = w.sum(axis=1) / self.n_data

            # Compute maximum likelihood estimates of alpha and beta.
            alphas, betas = [], []
            for i in range(self.n_components):
                alpha_list, beta_list = [], []
                for j in range(self.n_dim):
                    tmp_alpha, tmp_beta = self._bebin_mle(n[:, j], k[:, j], w[i])
                    alpha_list.append(tmp_alpha)
                    beta_list.append(tmp_beta)

                alphas.append(alpha_list)
                betas.append(beta_list)

            self.alphas_ = np.array(alphas)
            self.betas_ = np.array(betas)
        # END for iteration

        if not self.converged_:
            if self.verbose:
                logger.warning('The EM algorithm did not converge. Increase n_iter enough to ensure convergence.')

        self.log_likelihood_ = curr_weighted_loglikelihood

    def predict_proba(self, n, k):
        """Returns the posterior probabilities of each data point for each cluster."""
        l = np.array([self._bebin_likelihood(n, k, a, b) for a, b in zip(self.alphas_, self.betas_)])
        weighted_likelihood = l * self.pi_.reshape([self.n_components, 1])
        return weighted_likelihood / weighted_likelihood.sum(axis=0)

    def _n_parameters(self):
        """Returns the number of parameters estimated while fitting the model."""
        return int(2 * self.n_dim * self.n_components + (self.n_components - 1))

    def get_weights(self):
        """Returns the list of cluster weights.
        Note that a cluster weight is computed as a sum of posterior probabilities that
        each of the data point will be assiged to that cluster.
        """
        return self.pi_

    @property
    def means_(self):
        return np.array([a / (a + b) for a, b in zip(self.alphas_, self.betas_)])

    def get_means(self):
        return self.means_
    
    @property
    def dispersions_(self):
        return np.array([1 / (a + b + 1) for a, b in zip(self.alphas_, self.betas_)])

    def get_dispersions(self):
        return self.dispersions_

    def aic(self):
        return -2 * self.log_likelihood_ + 2 * self._n_parameters()
    
    def bic(self):
        return -2 * self.log_likelihood_ + np.log(self.n_data) * self._n_parameters()

    def get_n_dimensions(self):
        return self.n_dim

    def get_n_components(self):
        return self.n_components

    def get_depths(self):
        return self.depths
    
    def get_counts(self):
        return self.counts
    
    def get_headers(self):
        return self.headers
