"""Maximum likelihood estimation of a ZINB model of single cell RNA-sequencing
data
We model scRNA-Seq count data for cells i from multiple individuals/conditions
k as being generated by a zero-inflated negative binomial distribution per gene
j:
r_ijk ~ Poisson(lambda_ijk)
lambda_ijk ~ pi_ik delta_0(.) + (1 - pi_ik) Gamma(mu_ik, phi_ik)
The full specification of the model is given in Sarkar et al. Biorxiv 2018.
To fit this model in a reasonable amount of time for high capture rate, high
depth scRNA-Seq data (>5000 cells, >10000 genes, >100000 molecules per cell),
we optimize all of the parameters simultaneously using gradient descent. 
The key idea of the implementation is to use one-hot encoding to map samples to
paramaters, allowing us to compute a matrix of per-observation log likelihoods
using matrix multiplications and elementwise operations. This allows us to
exploit automatic differentiation and GPU code generation in tensorflow to
efficiently optimize the log likelihood.
Author: Abhishek Sarkar <aksarkar@alum.mit.edu>
"""
import numpy as np
import tensorflow as tf

def nb_llik(x, mean, inv_disp):
  """Return the log likelihood of x distributed as NB
  See Hilbe 2012, eq. 8.10
  mean - mean (> 0)
  inv_disp - inverse dispersion (> 0)
  """
  assert x.shape == mean.shape == inv_disp.shape
  return (x * tf.log(mean / inv_disp) -
          x * tf.log(1 + mean / inv_disp) -
          inv_disp * tf.log(1 + mean / inv_disp) +
          tf.lgamma(x + inv_disp) -
          tf.lgamma(inv_disp) -
          tf.lgamma(x + 1))

def zinb_llik(x, mean, inv_disp, logodds):
  """Log likelihood of x distributed as ZINB
  See Hilbe 2012, eq. 11.12, 11.13
  mean - mean (> 0)
  inv_disp - inverse dispersion (> 0)
  logodds - logit proportion of excess zeros
  """
  # Important identities:
  # log(x + y) = log(x) + softplus(log(y) - log(x))
  # log(sigmoid(x)) = -softplus(-x)
  case_zero = -tf.nn.softplus(-logodds) + tf.nn.softplus(nb_llik(x, mean, inv_disp) - logodds)
  case_non_zero = -tf.nn.softplus(logodds) + nb_llik(x, mean, inv_disp)
  return tf.where(tf.less(x, 1), case_zero, case_non_zero)

def fit(umi, onehot, size_factor, design=None, learning_rate=1e-2,
        max_epochs=100000, return_beta=False, warm_start=None, verbose=False):
  """Return the maximum likelihood solution of the model.
  If warm_start is not given, optimize the negative binomial likelihood only to
  find a good initialization. That problem is convex, so the initialization
  does not matter.
  umi - count matrix (n x p; float32)
  onehot - mapping of individuals to cells (m x n; float32)
  size_factor - size factor vector (n x 1; float32)
  design - confounder matrix (n x q; float32)
  learning_rate - step size for gradient descent
  max_epochs - maximum number of gradient descent steps
  return_beta - if True, returns the estimated confounding effect size matrix
  verbose - if True, outputs the log likelihood every 500 epochs
  warm_start - tuple of (log_mean, log_disp, logodds)
  Returns:
  log_mean - log mean parameter (m x p)
  log_disp - log dispersion parameter (m x p)
  logodds - logit proportion of excess zeros (m x p)
  nb_nll - negative binomial negative log likelihood at the returned solution
  zinb_nll - zero-inflated negative binomial negative log likelihood at the returned solution
  beta - if return_beta, confounding effect size matrix (q x p)
  """
  assert onehot.shape[0] == umi.shape[0]
  if design is not None:
    assert onehot.shape[0] == design.shape[0]
  n, p = umi.shape
  _, m = onehot.shape

  graph = tf.Graph()
  with graph.as_default(), graph.device('/gpu:*'):
    size_factor = tf.Variable(size_factor, trainable=False)
    umi = tf.Variable(umi, trainable=False)
    onehot = tf.Variable(onehot, trainable=False)
    if design is not None:
      _, k = design.shape
      design = tf.Variable(design, trainable=False)
      beta = tf.Variable(tf.zeros([k, p]))

    if warm_start is not None:
      log_mean, log_disp, logodds = warm_start
      assert log_mean.shape == (m, p)
      assert log_disp.shape == (m, p)
      assert logodds.shape == (m, p)
      mean = tf.exp(tf.Variable(log_mean))
      inv_disp = tf.exp(tf.Variable(-log_disp))
      logodds = tf.Variable(logodds)
    else:
      mean = tf.exp(tf.Variable(tf.zeros([m, p])))
      inv_disp = tf.exp(tf.Variable(tf.zeros([m, p])))
      logodds = tf.Variable(tf.constant(-8.0, shape=[m, p]))

    lam = size_factor * tf.matmul(onehot, mean)
    if design is not None:
      lam *= tf.exp(tf.matmul(design, beta))

    nb_nll = -tf.reduce_sum(nb_llik(umi, lam, tf.matmul(onehot, inv_disp)))
    zinb_nll = -tf.reduce_sum(zinb_llik(umi, lam, tf.matmul(onehot, inv_disp), tf.matmul(onehot, logodds)))

    train_nb = tf.train.RMSPropOptimizer(learning_rate=learning_rate).minimize(nb_nll)
    train_zinb = tf.train.RMSPropOptimizer(learning_rate=learning_rate).minimize(zinb_nll)
    print(train_nb)

    opt = [tf.log(mean), -tf.log(inv_disp), logodds, nb_nll, zinb_nll]
    if return_beta:
      opt.append(beta)

    obj = float('-inf')
    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      for i in range(max_epochs):
        if warm_start is None:
          _, update = sess.run([train_nb, nb_nll])
        else:
          _, update = sess.run([train_zinb, zinb_nll])          
        if not np.isfinite(update):
          raise tf.train.NanLossDuringTrainingError
        if verbose and not i % 500:
          print(i, update, end='\r')
      if verbose:
        print(i, update)
      return sess.run(opt)