# -------------------------------- WebLogo --------------------------------

#  Copyright (c) 2003-2004 The Regents of the University of California.
#  Copyright (c) 2005 Gavin E. Crooks
#  Copyright (c) 2006-2011, The Regents of the University of California, through
#  Lawrence Berkeley National Laboratory (subject to receipt of any required
#  approvals from the U.S. Dept. of Energy).  All rights reserved.

#  This software is distributed under the new BSD Open Source License.
#  <http://www.opensource.org/licenses/bsd-license.html>
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#  (1) Redistributions of source code must retain the above copyright notice,
#  this list of conditions and the following disclaimer.
#
#  (2) Redistributions in binary form must reproduce the above copyright
#  notice, this list of conditions and the following disclaimer in the
#  documentation and or other materials provided with the distribution.
#
#  (3) Neither the name of the University of California, Lawrence Berkeley
#  National Laboratory, U.S. Dept. of Energy nor the names of its contributors
#  may be used to endorse or promote products derived from this software
#  without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.

import random
from math import exp, log, sqrt
from typing import TYPE_CHECKING, Tuple

import numpy as np
import scipy.optimize
from numpy import asarray, float64, shape, zeros
from scipy.special import digamma, gamma, gammaincc, polygamma

if TYPE_CHECKING:
    from numpy.typing import ArrayLike


class Dirichlet(object):
    """The Dirichlet probability distribution. The Dirichlet is a continuous
    multivariate probability distribution across non-negative unit length
    vectors. In other words, the Dirichlet is a probability distribution of
    probability distributions. It is conjugate to the multinomial
    distribution and is widely used in Bayesian statistics.

    The Dirichlet probability distribution of order K-1 is

     p(theta_1,...,theta_K) d theta_1 ... d theta_K =
        (1/Z) prod_i=1,K theta_i^{alpha_i - 1} delta(1 -sum_i=1,K theta_i)

    The normalization factor Z can be expressed in terms of gamma functions:

      Z = {prod_i=1,K Gamma(alpha_i)} / {Gamma( sum_i=1,K alpha_i)}

    The K constants, alpha_1,...,alpha_K, must be positive. The K parameters,
    theta_1,...,theta_K are nonnegative and sum to 1.

    Status:
        Alpha
    """

    __slots__ = (
        "alpha",
        "_total",
        "_mean",
    )

    def __init__(self, alpha: "ArrayLike") -> None:
        """
        Args:
            - alpha  -- The parameters of the Dirichlet prior distribution.
                        A vector of non-negative real numbers.
        """
        # TODO: Check that alphas are positive
        # TODO: what if alpha's not one dimensional?
        self.alpha = asarray(alpha, float64)

        self._total = sum(alpha)
        self._mean: np.ndarray = self.alpha / self._total

    def sample(self) -> np.ndarray:
        """Return a randomly generated probability vector.

        Random samples are generated by sampling K values from gamma
        distributions with parameters a=\alpha_i, b=1, and renormalizing.

        Ref:
            A.M. Law, W.D. Kelton, Simulation Modeling and Analysis (1991).
        Authors:
            Gavin E. Crooks <gec@compbio.berkeley.edu> (2002)
        """
        alpha = self.alpha
        K = len(alpha)
        theta = zeros((K,), float64)

        for k in range(K):
            theta[k] = random.gammavariate(alpha[k], 1.0)
        theta /= sum(theta)

        return theta

    def mean(self) -> np.ndarray:
        return self._mean

    def covariance(self) -> np.ndarray:
        alpha = self.alpha
        A = sum(alpha)
        # A2 = A * A
        K = len(alpha)
        cv = zeros((K, K), float64)

        for i in range(K):
            cv[i, i] = alpha[i] * (1.0 - alpha[i] / A) / (A * (A + 1.0))

        for i in range(K):
            for j in range(i + 1, K):
                v = -alpha[i] * alpha[j] / (A * A * (A + 1.0))
                cv[i, j] = v
                cv[j, i] = v
        return cv

    def mean_x(self, x: "ArrayLike") -> float:
        x = asarray(x, float64)
        if shape(x) != shape(self.alpha):
            raise ValueError("Argument must be same dimension as Dirichlet")
        return sum(x * self.mean())

    def variance_x(self, x: "ArrayLike") -> float:
        x = asarray(x, float64)
        if shape(x) != shape(self.alpha):
            raise ValueError("Argument must be same dimension as Dirichlet")

        cv = self.covariance()
        var = np.dot(np.dot(np.transpose(x), cv), x)
        return var

    def mean_entropy(self) -> float:
        """Calculate the average entropy of probabilities sampled
        from this Dirichlet distribution.

        Returns:
            The average entropy.

        Ref:
            Wolpert & Wolf, PRE 53:6841-6854 (1996) Theorem 7
            (Warning: this paper contains typos.)
        Status:
            Alpha
        Authors:
            GEC 2005

        """
        # TODO: Optimize
        alpha = self.alpha
        A = float(sum(alpha))
        ent = 0.0
        for a in alpha:
            if a > 0:
                ent += -1.0 * a * digamma(1.0 + a)  # FIXME: Check
        ent /= A
        ent += digamma(A + 1.0)
        return ent

    def variance_entropy(self) -> float:
        """Calculate the variance of the Dirichlet entropy.

        Ref:
            Wolpert & Wolf, PRE 53:6841-6854 (1996) Theorem 8
            (Warning: this paper contains typos.)
        """
        alpha = self.alpha
        A = float(sum(alpha))
        A2 = A * (A + 1)
        L = len(alpha)

        dg1 = zeros((L), float64)
        dg2 = zeros((L), float64)
        tg2 = zeros((L), float64)

        for i in range(L):
            dg1[i] = digamma(alpha[i] + 1.0)
            dg2[i] = digamma(alpha[i] + 2.0)
            tg2[i] = polygamma(1, alpha[i] + 2.0)

        dg_Ap2 = digamma(A + 2.0)
        tg_Ap2 = polygamma(1, A + 2.0)

        mean = self.mean_entropy()
        var = 0.0

        for i in range(L):
            for j in range(L):
                if i != j:
                    var += (
                        ((dg1[i] - dg_Ap2) * (dg1[j] - dg_Ap2) - tg_Ap2)
                        * (alpha[i] * alpha[j])
                        / A2
                    )
                else:
                    var += (
                        ((dg2[i] - dg_Ap2) ** 2 + (tg2[i] - tg_Ap2))
                        * (alpha[i] * (alpha[i] + 1.0))
                        / A2
                    )

        var -= mean**2
        return var

    def mean_relative_entropy(self, pvec: "ArrayLike") -> float:
        pvec = asarray(pvec)
        ln_p = np.log(pvec)
        return -self.mean_x(ln_p) - self.mean_entropy()

    def variance_relative_entropy(self, pvec: "ArrayLike") -> float:
        pvec = asarray(pvec)
        ln_p = np.log(pvec)
        return self.variance_x(ln_p) + self.variance_entropy()

    def interval_relative_entropy(
        self, pvec: "ArrayLike", frac: float
    ) -> Tuple[float, float]:
        pvec = asarray(pvec)
        mean = self.mean_relative_entropy(pvec)
        variance = self.variance_relative_entropy(pvec)
        sd = sqrt(variance)

        # If the variance is small, use the standard 95%
        # confidence interval: mean +/- 1.96 * sd
        if mean / sd > 3.0:
            return max(0.0, mean - sd * 1.96), mean + sd * 1.96

        g = Gamma.from_mean_variance(mean, variance)
        low_limit = g.inverse_cdf((1.0 - frac) / 2.0)
        high_limit = g.inverse_cdf(1.0 - (1.0 - frac) / 2.0)

        return low_limit, high_limit


class Gamma(object):
    """The gamma probability distribution. (Not to be confused with the
    gamma function.)


    """

    __slots__ = "alpha", "beta"

    def __init__(self, alpha: float, beta: float) -> None:
        if alpha <= 0.0:
            raise ValueError("alpha must be positive")
        if beta <= 0.0:
            raise ValueError("beta must be positive")
        self.alpha = alpha
        self.beta = beta

    @classmethod
    def from_shape_scale(cls, shape: float, scale: float) -> "Gamma":
        return cls(shape, 1.0 / scale)

    @classmethod
    def from_mean_variance(cls, mean: float, variance: float) -> "Gamma":
        alpha = mean**2 / variance
        beta = alpha / mean
        return cls(alpha, beta)

    def mean(self) -> float:
        return self.alpha / self.beta

    def variance(self) -> float:
        return self.alpha / (self.beta**2)

    def sample(self) -> float:
        return random.gammavariate(self.alpha, 1.0 / self.beta)

    def pdf(self, x: float) -> float:
        if x == 0.0:
            return 0.0
        a = self.alpha
        b = self.beta
        return (x ** (a - 1.0)) * exp(-b * x) * (b**a) / gamma(a)

    def cdf(self, x: float) -> float:
        return 1.0 - gammaincc(self.alpha, self.beta * x)

    def inverse_cdf(self, p: float) -> float:
        def rootof(x: float) -> float:
            return self.cdf(exp(x)) - p

        root = scipy.optimize.newton(rootof, log(self.mean()))
        return exp(root)
