text: # source
  # http://nbviewer.ipython.org/github/tritemio/notebooks/blob/master/Mixture_Model_Fitting.ipynb
  
  from easydev import DevTools
  
  devtools = DevTools()
  from scipy.optimize import minimize, show_options
  import scipy.stats as ss
  import numpy as np
  import pylab
  from easydev import AttrDict
  
  from . import criteria
  
  import numpy as np
  
  half_log_two_pi = 0.5 * np.log(2 * np.pi)
  
  
  class Model(object):
      """New base model"""
  
      def __init__(self):
          pass
  
      def log_density(self, data):
          raise NotImplementedError
  
      def estimate(self, data, weights):
          raise NotImplementedError
  
      def generate(self):
          raise NotImplementedError
  
      def pdf(self):
          raise NotImplementedError
  
      def __repr__(self):
          raise NotImplementedError
  
  
  class GaussianModel(Model):
      """New gaussian model not used yet"""
  
      def __init__(self, mu=1, sigma=1):
          super(GaussianModel, self).__init__()
          self.mu = mu
          self.sigma = sigma
  
      def log_density(self, data):
          # log of gaussian distribution
          res = -((data - self.mu) ** 2) / (2 * self.sigma ** 2)
          res += -np.log(self.sigma) - half_log_two_pi
          return res
  
      def estimate(self, data, weights=None):
          assert len(data.shape) == 1, "Expect 1D data!"
          # d L/ dmu ==0 and dL/dsigma==0
          if weights:
              wsum = np.sum(weights)
              self.mu = np.sum(weights * data) / wsum
              self.sigma = np.sqrt(np.sum(weights * (data - self.mu) ** 2) / wsum)
          else:
              self.mu = np.sum(data)
              self.sigma = np.sqrt(np.sum((data - self.mu) ** 2))
  
      def generate(self, N):
          return np.array(
              [scipy.stats.norm.rvs(self.mu, self.sigma) for this in range(N)]
          )
  
      def pdf(self, data):
          pass
  
  
  class PoissonModel(Model):
      """New poisson model not used yet"""
  
      def __init__(self, lmbda=10):
          super(PoissonModel, self).__init__()
          self.lmbda = lmbda
  
      def log_density(self, data):
          # log of gaussian distribution
          return -self.lmbda + data * np.log(self.lmbda)
  
      def estimate(self, data, weights=None):
          assert len(data.shape) == 1, "Expect 1D data!"
          # d L/ dmu ==0 and dL/dsigma==0
          if weights:
              wsum = np.sum(weights)
              self.lmbda = np.sum(weights * data) / wsum
          else:
              self.lmbda = np.sum(data)
  
          # Log likelihood of poisson of N iid is L = PI_i^n \lambda^{x_i}
          # e^-\lambda / x_i!  and the loglikelihood is then l = sum_i^n
          # x_i log\lmabda - n \lambda
  
      def generate(self, N):
          return np.array([scipy.stats.poisson.rvs(self.lmbda) for this in range(N)])
  
      def __repr__(self):
          return "Poisson[lmbda={lmbda:.4g}]".format(lmbda=self.lmbda)
  
  
  class GaussianMixture(object):
      """Creates a mix of Gaussian distribution
  
      .. plot::
          :width: 80%
          :include-source:
  
          from biokit.stats.mixture import GaussianMixture
          m = GaussianMixture(mu=[-1,1], sigma=[0.5, 0.5], mixture=[.2, .8], N=1000)
          m.plot()
  
      """
  
      def __init__(self, mu=[-1, 1], sigma=[1, 1], mixture=[0.5, 0.5], N=1000):
          """.. rubric:: constructor
  
          :param list mu: list of mean for each model
          :param list sigma: list of standard deviation for each model
          :param list mixture: list of amplitude for each model
  
          """
          assert len(mu) == len(sigma)
          assert len(mu) == len(mixture)
          self.mu = mu
          self.sigma = sigma
          self.mixture = mixture
          self.data = []
          self.N = N
          self.Ns = [int(x * N) for x in mixture]
          self.k = len(self.mu)
  
          if sum(self.Ns) != N:
              print("Warning: rounding mixture ratio. total N will be %s" % sum(self.Ns))
  
          for m, s, n in zip(self.mu, self.sigma, self.Ns):
              data = pylab.normal(m, s, size=n)
              self.data.extend(data)
  
      def plot(self, bins=30, density=True):
          self.hist(bins=bins, density=density)
  
      def hist(self, bins=30, density=True):
          pylab.hist(self.data, bins=bins, density=density)
  
  
  class GaussianMixtureModel(object):
      """Gaussian Mixture Model
  
      .. plot::
  
          from biokit.stats import mixture
          m = mixture.GaussianMixtureModel()
          m.pdf(1, params=[1, 0.5, 0.2, 1, 0.5, 0.8])
  
      """
  
      def __init__(self, k=2):
          self.k = k
  
      def pdf(self, x, params, normalise=True):
          """Expected parameters are
  
  
          params is a list of gaussian distribution ordered as mu, sigma, pi,
          mu2, sigma2, pi2, ...
  
          """
          assert divmod(len(params), 3)[1] == 0
          assert len(params) >= 3 * self.k
          k = len(params) / 3
  
          self.k = k
  
          pis = np.array(params[2::3])
  
          if any(np.array(pis) < 0):
              return 0
          if normalise is True:
              pis /= pis.sum()
          # !!! sum pi must equal 1 otherwise may diverge badly
          data = 0
          for i in range(0, int(k)):
              mu, sigma, pi_ = params[i * 3 : (i + 1) * 3]
              pi_ = pis[i]
              if sigma != 0:
                  data += pi_ * ss.norm.pdf(x, mu, sigma)
          return data
  
      def log_likelihood(self, params, sample):
          res = -1 * pylab.log(self.pdf(sample, params)).sum()
          return res
  
  
  class Fitting(object):
      """Base class for :class:`EM` and :class:`GaussianMixtureFitting`"""
  
      def __init__(self, data, k=2, method="Nelder-Mead"):
          """.. rubric:: constructor
  
          :param list data:
          :param int k: number of GMM to use
          :param str method: minimization method to be used (one of scipy optimise module)
  
  
          """
          self.data = np.array(data)
          self.size = float(len(self.data))
          self._k = k
          self._model = None
          # initialise the model
          self.k = k
          self.verbose = True
  
      def _get_k(self):
          return self._k
  
      def _set_k(self, k):
          assert k > 0
          gmm = GaussianMixtureModel(k=k)
          self._k = k
          self._model = gmm
  
      k = property(_get_k, _set_k)
  
      def _get_model(self):
          return self._model
  
      model = property(_get_model)
  
      def get_guess(self):
          """Random guess to initialise optimisation"""
          params = {}
          m = self.data.min()
          M = self.data.max()
          range_ = M - m
  
          mus = [m + range_ / (self.k + 1.0) * i for i in range(1, self.k + 1)]
          params["mus"] = mus
  
          sigma = range_ / float(self.k + 1) / 2.0
          params["sigmas"] = [sigma] * self.k
  
          params["pis"] = [1.0 / self.k] * self.k
  
          params = [
              [mu, sigma, pi]
              for mu, sigma, pi in zip(params["mus"], params["sigmas"], params["pis"])
          ]
          params = list(pylab.flatten(params))
          return params
  
      def plot(
          self,
          density=True,
          N=1000,
          Xmin=None,
          Xmax=None,
          bins=50,
          color="red",
          lw=2,
          hist_kw={"color": "#5F9EA0", "edgecolor": "k"},
          ax=None,
      ):
  
          if ax:
              ax.hist(self.data, density=density, bins=bins, **hist_kw)
          else:
              pylab.hist(self.data, density=density, bins=bins, **hist_kw)
          if Xmin is None:
              Xmin = self.data.min()
          if Xmax is None:
              Xmax = self.data.max()
          X = pylab.linspace(Xmin, Xmax, N)
  
          if ax:
              ax.plot(
                  X, [self.model.pdf(x, self.results.x) for x in X], color=color, lw=lw
              )
          else:
              pylab.plot(
                  X, [self.model.pdf(x, self.results.x) for x in X], color=color, lw=lw
              )
  
          K = len(self.results.x)
          # The PIs must be normalised
          for i in range(self.k):
  
              mu, sigma, pi_ = (
                  self.results.mus[i],
                  self.results.sigmas[i],
                  self.results.pis[i],
              )
              if ax:
                  ax.plot(
                      X,
                      [pi_ * ss.norm.pdf(x, mu, sigma) for x in X],
                      "k--",
                      alpha=0.7,
                      lw=2,
                  )
              else:
                  pylab.plot(
                      X,
                      [pi_ * ss.norm.pdf(x, mu, sigma) for x in X],
                      "k--",
                      alpha=0.7,
                      lw=2,
                  )
  
  
  class GaussianMixtureFitting(Fitting):
      """GaussianMixtureFitting using scipy minization
  
      .. plot::
          :width: 80%
          :include-source:
  
          from biokit.stats.mixture import GaussianMixture, GaussianMixtureFitting
          m = GaussianMixture(mu=[-1,1], sigma=[0.5,0.5], mixture=[0.2,0.8])
          mf = GaussianMixtureFitting(m.data)
          mf.estimate(k=2)
          mf.plot()
  
  
      """
  
      def __init__(self, data, k=2, method="Nelder-Mead"):
          """
  
          Here we use the function minimize() from scipy.optimization.
          The list of (currently) available minimization methods is 'Nelder-Mead' (simplex),
          'Powell', 'CG', 'BFGS', 'Newton-CG',> 'Anneal', 'L-BFGS-B' (like BFGS but bounded),
          'TNC', 'COBYLA', 'SLSQPG'.
  
          """
          super(GaussianMixtureFitting, self).__init__(data, k=k, method=method)
          self._method = method
  
      def _get_method(self):
          return self._method
  
      def _set_method(self, method):
          devtools.check_param_in_list(
              method,
              ["Nelder-Mead", "Powell", "CG", "BFGS", "Newton-CG", "Anneal", "L-BFGS-B"],
          )
          self._method = method
  
      method = property(_get_method, _set_method)
  
      def estimate(self, guess=None, k=None, maxfev=2e4, maxiter=1e3, bounds=None):
          """guess is a list of parameters as expected by the model
  
  
          guess = {'mus':[1,2], 'sigmas': [0.5, 0.5], 'pis': [0.3, 0.7]  }
  
          """
          if k is not None:
              self.k = k
  
          if guess is None:
              # estimate the mu/sigma/pis from the data
              guess = self.get_guess()
  
          res = minimize(
              self.model.log_likelihood,
              x0=guess,
              args=(self.data,),
              method=self.method,
              options=dict(maxiter=maxiter, maxfev=maxfev),
              bounds=bounds,
          )
  
          self.results = res
          pis = np.array(self.results.x[2::3])
          self.results.pis_raw = pis.copy()
          # The ratio may be negative, in which case we need to normalise.
          # An example would be to have -0.35, -0.15, which normalise would five 0.7, 0.3 as expected.
          """if sum(pis<0) > 0:
              unstable = True
              pis /= pis.sum()
              if self.verbose:
                  print("Unstable... found negative pis (k=%s)" % self.k)
          else:
              unstable = False
              pis /= pis.sum()
          """
          unstable = False
          k = len(self.results.x) / 3
          params = []
          for i in range(0, int(k)):
              params.append(self.results.x[i * 3])
              params.append(self.results.x[(i * 3 + 1)])
              params.append(pis[i])
          self.results.x = params
  
          # FIXME shall we multiply by -1 ??
          self.results.log_likelihood = self.model.log_likelihood(params, self.data)
          if self.results.log_likelihood and unstable is False:
              self.results.AIC = criteria.AIC(
                  self.results.log_likelihood, self.k, logL=True
              )
              self.results.AICc = criteria.AICc(
                  self.results.log_likelihood, self.k, self.data.size, logL=True
              )
              self.results.BIC = criteria.BIC(
                  self.results.log_likelihood, self.k, self.data.size, logL=True
              )
          else:
              self.results.AIC = 1000
              self.results.AICc = 1000
              self.results.BIC = 1000
  
          pis = np.array(self.results.x[2::3])
  
          self.results.pis = list(pis / pis.sum())
          self.results.sigmas = self.results.x[1::3]
          self.results.mus = self.results.x[0::3]
  
          return res
  
  
  class EM(Fitting):
      """Expectation minimization class to estimate parameters of GMM
  
      .. plot::
          :width: 50%
          :include-source:
  
          from biokit.stats.mixture import GaussianMixture, EM
          m = GaussianMixture(mu=[-1,1], sigma=[0.5,0.5], mixture=[0.2,0.8])
          em = EM(m.data)
          em.estimate(k=2)
          em.plot()
  
      """
  
      def __init__(self, data, model=None, max_iter=100):
          """.. rubric:: constructor
  
          :param data:
          :param model: not used. Model is the :class:`GaussianMixtureModel` but
              could be other model.
          :param int max_iter: max iteration for the minization
  
          """
          super(EM, self).__init__(data, k=2)  # default is k=2
          self.max_iter = max_iter
  
      # @do_profile()
      def estimate(self, guess=None, k=2):
          """
  
          :param list guess: a list to provide the initial guess. Order is mu1, sigma1,
              pi1, mu2, ...
          :param int k: number of models to be used.
          """
          # print("EM estimation")
          self.k = k
          # Initial guess of parameters and initializations
          if guess is None:
              # estimate the mu/sigma/pis from the data
              guess = self.get_guess()
  
          mu = np.array(guess[0::3])
          sig = np.array(guess[1::3])
          pi_ = np.array(guess[2::3])
          N_ = len(pi_)
  
          gamma = np.zeros((N_, int(self.size)))
          N_ = np.zeros(N_)
          p_new = guess
  
          # EM loop
          counter = 0
          converged = False
  
          self.mus = []
  
          while not converged:
              # Compute the responsibility func. and new parameters
              for k in range(0, self.k):
                  # unstable if eslf.model.pdf is made of zeros
  
                  # self.model.pdf(self.data, p_new,normalise=False).sum()!=0:
                  gamma[k, :] = pi_[k] * ss.norm.pdf(self.data, mu[k], sig[k])
                  gamma[k, :] /= self.model.pdf(self.data, p_new, normalise=False)
                  """else:
                      gamma[k, :] = pi_[k]*pylab.normpdf(self.data, mu[k],
                          sig[k])/(self.model.pdf(self.data, p_new,
                              normalise=False)+1e-6)
                  """
                  N_[k] = gamma[k].sum()
                  mu[k] = np.sum(gamma[k] * self.data) / N_[k]
                  sig[k] = pylab.sqrt(np.sum(gamma[k] * (self.data - mu[k]) ** 2) / N_[k])
                  pi_[k] = N_[k] / self.size
  
              self.results = {"x": p_new, "nfev": counter, "success": converged}
  
              p_new = []
              for this in range(self.k):
                  p_new.extend([mu[this], sig[this], pi_[this]])
  
              # p_new = [(mu[x], sig[x], pi_[x]) for x in range(0, self.k)]
              # p_new = list(pylab.flatten(p_new))
  
              self.status = True
              try:
                  assert abs(N_.sum() - self.size) / self.size < 1e-6
                  assert abs(pi_.sum() - 1) < 1e-6
              except:
                  print("issue arised at iteration %s" % counter)
                  self.debug = {"N": N_, "pis": pi_}
                  self.status = False
                  break
  
              self.mus.append(mu)
  
              # Convergence check
              counter += 1
              converged = counter >= self.max_iter
  
          self.gamma = gamma
  
          if self.status is True:
              self.results = {"x": p_new, "nfev": counter, "success": converged}
  
          self.results = AttrDict(**self.results)
          self.results.mus = self.results.x[0::3]
          self.results.sigmas = self.results.x[1::3]
          self.results.pis = self.results.x[2::3]
  
          log_likelihood = self.model.log_likelihood(self.results.x, self.data)
          self.results.AIC = criteria.AIC(log_likelihood, k, logL=True)
  
          self.results.log_likelihood = log_likelihood
          self.results.AIC = criteria.AIC(log_likelihood, self.k, logL=True)
          self.results.AICc = criteria.AICc(
              log_likelihood, self.k, self.data.size, logL=True
          )
          self.results.BIC = criteria.BIC(
              log_likelihood, self.k, self.data.size, logL=True
          )
  
      def plot(self, model_parameters=None, **kwargs):
          """Take a list of dictionnaries with models parameters to plot
          predicted models. If user doesn't provide parameters, the standard
          plot function from fitting is used.
  
          Example:
              model_parameters=[{"mu": 5, "sigma": 0.5, "pi": 1}]
          """
          if not model_parameters:
              return super(EM, self).plot(**kwargs)
          # Set parameters with the dictionnary
          self.k = len(model_parameters)
          self.results = AttrDict()
          self.results.mus = [model["mu"] for model in model_parameters]
          self.results.sigmas = [model["sigma"] for model in model_parameters]
          self.results.pis = [model["pi"] for model in model_parameters]
          parms_keys = ("mu", "sigma", "pi")
          self.results.x = [
              model[key] for model in model_parameters for key in parms_keys
          ]
          return super(EM, self).plot(**kwargs)
  
  
  class AdaptativeMixtureFitting(object):
      """Automatic Adaptative Mixture Fitting
  
      .. plot::
          :width: 80%
          :include-source:
  
          from biokit.stats.mixture import AdaptativeMixtureFitting, GaussianMixture
          m = GaussianMixture(mu=[-1,1], sigma=[0.5,0.5], mixture=[0.2,0.8])
          amf = AdaptativeMixtureFitting(m.data)
          amf.run(kmin=1, kmax=6)
          amf.diagnostic(k=amf.best_k)
  
  
      """
  
      # TODO: propose to use EM or Minimization.
      def __init__(self, data, method="Nelder-Mead"):
          self.fitting = GaussianMixtureFitting(data, method=method)
          self.verbose = True
  
      def run(self, kmin=1, kmax=6, criteria="AICc"):
          self.all_results = {}
          self.x = []
  
          for k in range(kmin, kmax + 1):
              self.fitting.estimate(k=k)
  
              # here criteria does not matter. if one fails, all fail
              if self.fitting.results["AIC"] != 1000:
                  self.x.append(k)
                  self.all_results[k] = self.fitting.results.copy()
  
          m = np.array([self.all_results[x][criteria] for x in self.x]).min()
          index = np.array([self.all_results[x][criteria] for x in self.x]).argmin()
          if self.verbose:
              print("Found min ", m, "for k  ", self.x[index])
          self.best_k = self.x[index]
          self.min_value = m
  
      def plot(self, criteria="AICc"):
  
          pylab.plot(
              self.x, [self.all_results[x]["BIC"] for x in self.x], "o-", label="BIC"
          )
          pylab.plot(
              self.x, [self.all_results[x]["AIC"] for x in self.x], "o-", label="AIC"
          )
          pylab.plot(
              self.x, [self.all_results[x]["AICc"] for x in self.x], "o-", label="AICc"
          )
  
          m = np.array([self.all_results[x][criteria] for x in self.x]).min()
          # [exp((m-this[criteria])/2) for this in amf.all_results]
  
          pylab.axhline(
              m - pylab.log(0.9) * 2, color="k", label="90% equi-probability", alpha=0.9
          )
          pylab.axhline(
              m - pylab.log(0.5) * 2, color="k", label="50% equi-probability", alpha=0.5
          )
          pylab.axhline(
              m - pylab.log(0.3) * 2, color="k", label="30% equi-probability", alpha=0.3
          )
  
          pylab.legend()
  
      def diagnostic(self, kmin=1, kmax=8, k=None, ymax=None):
          self.run(kmin=kmin, kmax=kmax)
          pylab.clf()
          pylab.subplot(3, 1, 2)
          self.plot()
          mf = GaussianMixtureFitting(self.fitting.data)
          if k is None:
              mf.estimate(k=self.best_k)
          else:
              mf.estimate(k=k)
          pylab.subplot(3, 1, 1)
          mf.plot()
          if ymax is not None:
              pylab.ylim([0, ymax])
  
          pylab.subplot(3, 1, 3)
          min_value = np.array([self.all_results[x]["AICc"] for x in self.x]).min()
          pylab.plot(
              self.x,
              [pylab.exp((min_value - self.all_results[k]["AICc"]) / 2) for k in self.x],
              "o-",
              label="AICc",
          )
          min_value = np.array([self.all_results[x]["AIC"] for x in self.x]).min()
          pylab.plot(
              self.x,
              [pylab.exp((min_value - self.all_results[k]["AIC"]) / 2) for k in self.x],
              "o-",
              label="AIC",
          )
  
          pylab.xlabel("probability of information loss (based on AICc")
          pylab.legend()
  