#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Spatial Factorization base class
@author: townesf
"""
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from os import path
from math import ceil
from tensorflow import linalg as tfl
from sklearn.decomposition import TruncatedSVD
from sklearn.linear_model import LinearRegression
from sklearn.neighbors import KNeighborsRegressor

from spatial_factorization import likelihoods, misc, nnfu

tfd = tfp.distributions
tfb = tfp.bijectors
tv = tfp.util.TransformedVariable
tfk = tfp.math.psd_kernels

dtp = "float32"
rng = np.random.default_rng()


class SpatialDEF(object):
    def __init__(
        self,
        J,
        L,
        Z,
        lik="poi",
        psd_kernel=tfk.MaternThreeHalves,
        nugget=1e-5,
        length_scale=0.1,
        disp="default",
        nonneg=False,
        isotropic=True,
        feature_means=None,
        **kwargs
    ):
        """
        Non-negative process factorization
        Parameters
        ----------
        J : integer scalar
            Number of features in the multivariate outcome.
        L : integer scalar
            Number of desired latent Gaussian processes.
        T : integer scalar
            Number of desired non-spatial latent factors
        Z : 2D numpy array
            Coordinates of inducing point locations in input space.
            First dimension: 'M' is number of inducing points.
            Second dimension: 'D' is dimensionality of input to GP.
            More inducing points= slower but more accurate inference.
        lik: likelihood (Poisson or Gaussian)
        disp: overdispersion parameter initialization.
        --for Gaussian likelihood, the scale (stdev) parameter
        --for negative binomial likelihood, the parameter phi such that
          var=mean+phi*mean^2, i.e. phi->0 corresponds to Poisson, phi=1 to geometric
        --for Poisson likelihood, this parameter is ignored and set to None
        psd_kernel : an object of class PositiveSemidefiniteKernel, must accept a
            length_scale parameter.
        feature_means : if input data is centered (for lik='gau', nonneg=False)
        """
        super().__init__(**kwargs)
        # self.is_spatial=True
        self.lik = lik
        self.isotropic = isotropic
        M, D = Z.shape
        self.Z = tf.Variable(Z, trainable=False, dtype=dtp, name="inducing_points")
        self.nonneg = tf.Variable(nonneg, trainable=False, name="is_non_negative")
        # variational parameters
        with tf.name_scope("variational"):
            self.delta = tf.Variable(
                rng.normal(size=(L, M)), dtype=dtp, name="mean"
            )  # LxM
            _Omega_tril = self._init_Omega_tril(L, M, nugget=nugget)
            # _Omega_tril = .01*tf.eye(M,batch_shape=[L],dtype=dtp)
            self.Omega_tril = tv(
                _Omega_tril, tfb.FillScaleTriL(), dtype=dtp, name="covar_tril"
            )  # LxMxM
        # GP hyperparameters
        with tf.name_scope("gp_mean"):
            if self.nonneg:
                prior_mu, prior_sigma = misc.lnormal_approx_dirichlet(max(L, 1.1))
                self.beta0 = tf.Variable(
                    prior_mu * tf.ones((L, 1)), dtype=dtp, name="intercepts"
                )
            else:
                self.beta0 = tf.Variable(
                    tf.zeros((L, 1)), dtype=dtp, name="intercepts"
                )  # L
            self.beta = tf.Variable(tf.zeros((L, D)), dtype=dtp, name="slopes")  # LxD
        with tf.name_scope("gp_kernel"):
            self.nugget = tf.Variable(nugget, dtype=dtp, trainable=False, name="nugget")
            self.amplitude = tv(
                np.tile(1.0, [L]), tfb.Softplus(), dtype=dtp, name="amplitude"
            )
            self._ls0 = length_scale  # store for .reset() method
            if self.isotropic:
                self.length_scale = tv(
                    np.tile(self._ls0, [L]),
                    tfb.Softplus(),
                    dtype=dtp,
                    name="length_scale",
                )
            else:
                self.scale_diag = tv(
                    np.tile(np.sqrt(self._ls0), [L, D]),
                    tfb.Softplus(),
                    dtype=dtp,
                    name="scale_diag",
                )
        # Loadings weights
        if self.nonneg:
            self.W = tf.Variable(
                rng.exponential(size=(J, L)),
                dtype=dtp,
                constraint=misc.make_nonneg,
                name="loadings",
            )  # JxL
        else:
            self.W = tf.Variable(
                rng.normal(size=(J, L)), dtype=dtp, name="loadings"
            )  # JxL
        self.psd_kernel = psd_kernel  # this is a class, not yet an object
        # likelihood parameters, set defaults
        self._disp0 = disp
        self._init_misc()
        self.Kuu_chol = tf.Variable(
            self.eval_Kuu_chol(self.get_kernel()), dtype=dtp, trainable=False
        )
        if self.lik == "gau" and not self.nonneg:
            self.feature_means = feature_means
        else:
            self.feature_means = None

    @staticmethod
    def _init_Omega_tril(L, M, nugget=None):
        """
        convenience function for initializing the batch of lower triangular
        cholesky factors of the variational covariance matrices.
        L: number of latent dimensions (factors)
        M: number of inducing points
        """
        # note most of these operations are in float64 by default
        # the 0.01 below is to make the elbo more numerically stable at initialization
        Omega_sqt = 0.01 * rng.normal(size=(L, M, M))  # LxMxM
        Omega = [
            Omega_sqt[l, :, :] @ Omega_sqt[l, :, :].T for l in range(L)
        ]  # list len L, elements MxM
        # Omega += nugget*np.eye(M)
        res = np.stack([np.linalg.cholesky(Omega[l]) for l in range(L)], axis=0)
        return res.astype(dtp)

    def _init_misc(self):
        """
        misc initialization shared between __init__ and reset
        """
        J = self.W.shape[0]
        self.disp = likelihoods.init_lik(self.lik, J, disp=self._disp0, dtp=dtp)
        # stuff to facilitate caching
        self.trvars_kernel = tuple(
            i for i in self.trainable_variables if i.name[:10] == "gp_kernel/"
        )
        self.trvars_nonkernel = tuple(
            i for i in self.trainable_variables if i.name[:10] != "gp_kernel/"
        )
        if self.isotropic:
            self.kernel = self.psd_kernel(
                amplitude=self.amplitude, length_scale=self.length_scale
            )
        else:
            self.kernel = tfk.FeatureScaled(
                self.psd_kernel(amplitude=self.amplitude), self.scale_diag
            )

    def get_dims(self):
        return self.W.shape[1]

    def get_loadings(self):
        return self.W.numpy()

    def set_loadings(self, Wnew):
        self.W.assign(Wnew, read_value=False)

    def init_loadings(self, Y, X=None, sz=1, **kwargs):
        """
        Use either PCA or NMF to initialize the loadings matrix from data Y
        """
        if self.nonneg:
            init_nsf_with_nmf(self, Y, X=X, sz=sz, **kwargs)
        else:  # real-valued factors
            if self.lik in ("poi", "nb"):
                pass  # use GLM-PCA?
            elif self.lik == "gau":
                L = self.W.shape[1]
                fit = TruncatedSVD(L).fit(Y)
                self.set_loadings(fit.components_.T)
            else:
                raise likelihoods.InvalidLikelihoodError

    def generate_pickle_path(self, sz, base=None):
        """
        sz : str
          Indicate what type of size factors are used (eg 'none' or 'scanpy').
        base : str, optional
          Parent directory for saving pickle files. The default is cwd.
        """
        pars = {
            "L": self.W.shape[1],
            "lik": self.lik,
            "sz": sz,
            "model": "NSF" if self.nonneg else "RSF",
            "kernel": self.psd_kernel.__name__,
            "M": self.Z.shape[0],
        }
        pth = misc.params2key(pars)
        if base:
            pth = path.join(base, pth)
        return pth

    def get_kernel(self):
        return self.kernel

    def eval_Kuu_chol(self, kernel=None):
        if kernel is None:
            kernel = self.get_kernel()
        M, D = self.Z.shape
        Kuu = kernel.matrix(self.Z, self.Z) + self.nugget * tf.eye(M)
        # Kuu_chol is LxMxM and lower triangular in last two dims
        return tfl.cholesky(Kuu)

    def get_Kuu_chol(self, kernel=None, from_cache=False):
        if not from_cache:
            Kuu_chol = self.eval_Kuu_chol(kernel=kernel)
            self.Kuu_chol.assign(Kuu_chol)  # update cache
            return Kuu_chol
        else:  # use cached variable, gradients cannot propagate to kernel hps
            return self.Kuu_chol

    def get_mu_z(self):
        return self.beta0 + tfl.matmul(self.beta, self.Z, transpose_b=True)  # LxM

    def sample_latent_GP_funcs(
        self, X, S=1, kernel=None, mu_z=None, Kuu_chol=None, chol=True
    ):
        """
        Draw random samples of the latent variational GP function values "F"
        based on spatial coordinates X.
        The sampling comes from the variational approximation to the posterior.
        This function is needed to compute the expected log-likelihood term of the
        ELBO.
        X is a numpy array with shape NxD
        N=number of observations
        D=number of spatial dimensions
        The first dimension can be observations from a minibatch.
        S is the number of random samples to draw from latent GPs
        logscale: if False the function vals are exponentiated before returning
        i.e. they are positive valued random functions.
        If logscale=True (default), the functions are real-valued
        """
        if kernel is None:
            kernel = self.get_kernel()
        if mu_z is None:
            mu_z = self.get_mu_z()
        if Kuu_chol is None:
            Kuu_chol = self.get_Kuu_chol(kernel=kernel, from_cache=(not chol))
        N = X.shape[0]
        L = self.W.shape[1]
        mu_x = self.beta0 + tfl.matmul(self.beta, X, transpose_b=True)  # LxN
        Kuf = kernel.matrix(self.Z, X)  # LxMxN
        Kff_diag = kernel.apply(X, X, example_ndims=1) + self.nugget  # LxN
        alpha_x = tfl.cholesky_solve(Kuu_chol, Kuf)  # LxMxN
        mu_tilde = mu_x + tfl.matvec(
            alpha_x, self.delta - mu_z, transpose_a=True
        )  # LxN
        # compute the alpha(x_i)'(K_uu-Omega)alpha(x_i) term
        a_t_Kchol = tfl.matmul(alpha_x, Kuu_chol, transpose_a=True)  # LxNxM
        aKa = tf.reduce_sum(tf.square(a_t_Kchol), axis=2)  # LxN
        a_t_Omega_tril = tfl.matmul(alpha_x, self.Omega_tril, transpose_a=True)  # LxNxM
        aOmega_a = tf.reduce_sum(tf.square(a_t_Omega_tril), axis=2)  # LxN
        Sigma_tilde = Kff_diag - aKa + aOmega_a  # LxN
        eps = tf.random.normal(
            (S, L, N)
        )  # note this is not the same random generator as self.rng!
        return mu_tilde + tf.math.sqrt(Sigma_tilde) * eps  # "F", dims: SxLxN

    def sample_predictive_mean(
        self, X, sz=1, S=1, kernel=None, mu_z=None, Kuu_chol=None, chol=True
    ):
        """
        See sample_latent_variational_GP_funcs for X,S definitions
        sz is a tensor of shape (N,1) of size factors.
        Typically sz would be the rowSums or rowMeans of the outcome matrix Y.
        """
        F = self.sample_latent_GP_funcs(
            X, S=S, kernel=kernel, mu_z=mu_z, Kuu_chol=Kuu_chol, chol=chol
        )  # SxLxN
        if self.nonneg:
            Lam = tfl.matrix_transpose(tfl.matmul(self.W, tf.exp(F)))  # SxNxJ
            if self.lik == "gau":
                return Lam
            else:
                return sz * Lam
        else:
            Lam = tfl.matrix_transpose(tfl.matmul(self.W, F))
            if self.lik == "gau":  # identity link
                return Lam
            else:  # log link (poi, nb)
                return tf.exp(tf.math.log(sz) + Lam)

    def eval_kl_term(self, mu_z, Kuu_chol):
        """
        KL divergence from the prior distribution to the variational distribution.
        This is one component of the ELBO:
        ELBO=expected log-likelihood - sum(kl_terms)
        qpars: a tuple containing (mu_z,Kuu_chol)
        qpars can be obtained by calling self.get_variational_params()
        mu_z is the GP mean function at all inducing points (dimension: LxM)
        Kuu_chol is the cholesky lower triangular of the kernel matrix of all inducing points.
        Its dimension is LxMxM
        where L = number of latent dimensions and M = number of inducing points.
        """
        qu = tfd.MultivariateNormalTriL(loc=self.delta, scale_tril=self.Omega_tril)
        pu = tfd.MultivariateNormalTriL(loc=mu_z, scale_tril=Kuu_chol)
        return qu.kl_divergence(pu)  # L-vector

    # @tf.function
    def elbo_avg(self, X, Y, sz=1, S=1, Ntot=None, chol=True):
        """
        Parameters
        ----------
        X : numpy array of spatial coordinates (NxD)
            **OR** a tuple of spatial coordinates, multivariate outcomes,
            and size factors (convenient for minibatches from tensor slices)
        Y : numpy array of multivariate outcomes (NxJ)
            If Y is None then X must be a tuple of length three
        sz : size factors, optional
            vector of length N, typically the rowSums or rowMeans of Y.
            If X is a tuple then this is ignored as sz is expected in the third
            element of the X tuple.
        S : integer, optional
            Number of random GP function evaluations to use. The default is 1.
            Larger S=more accurate approximation to true ELBO but slower
        Ntot : total number of observations in full dataset
            This is needed when X,Y,sz are a minibatch from the full data
            If Ntot is None, we assume X,Y,sz provided are the full data NOT a minibatch.
        Returns
        -------
        The numeric evidence lower bound value, divided by Ntot.
        """
        batch_size, J = Y.shape
        if Ntot is None:
            Ntot = batch_size  # no minibatch, all observations provided
        ker = self.get_kernel()
        mu_z = self.get_mu_z()
        Kuu_chol = self.get_Kuu_chol(kernel=ker, from_cache=(not chol))
        # kl_terms is not affected by minibatching so use reduce_sum
        kl_term = tf.reduce_sum(self.eval_kl_term(mu_z, Kuu_chol))
        # Sample z1
        z1 = self.sample_predictive_mean(
            X, sz=sz, S=S, kernel=ker, mu_z=mu_z, Kuu_chol=Kuu_chol
        )
        # Compute p(z0|z1) using sampled z1 and z0
        mean_top = jnp.exp(z1).dot(W1)
        # Sample z0
        z0 = gamma_sample(var_shape_z0, var_rate_z0)
        prior_z0 = gamma_logpdf(z0, alpha, alpha / mean_top)
        # And entropy of z0
        ent_z0 = -gamma_logpdf(z0, var_shape_z0, var_rate_z0)
        kl_term = kl_term + prior_z0 + ent_z0
        # Compute p(y|z0)
        mean_bottom = z0.dot(W0)
        eloglik = poisson_logpdf()
        eloglik = likelihoods.lik_to_distr(self.lik, Mu, self.disp).log_prob(Y)
        return J * tf.reduce_mean(eloglik) - kl_term / Ntot

    def train_step(self, D, optimizer, optimizer_k, S=1, Ntot=None, chol=True):
        """
        Executes one training step and returns the loss.
        D is training data: a tensorflow dataset (from slices) of (X,Y,sz)
        This function computes the loss and gradients, and uses the latter to
        update the model's parameters.
        """
        with tf.GradientTape(persistent=True) as tape:
            loss = -self.elbo_avg(D["X"], D["Y"], sz=D["sz"], S=S, Ntot=Ntot, chol=chol)
        try:
            gradients = tape.gradient(loss, self.trvars_nonkernel)
            if chol:
                gradients_k = tape.gradient(loss, self.trvars_kernel)
                optimizer_k.apply_gradients(zip(gradients_k, self.trvars_kernel))
            optimizer.apply_gradients(zip(gradients, self.trvars_nonkernel))
        finally:
            del tape
        return loss

    def validation_step(self, D, S=1, chol=False):
        """
        Compute the validation loss on held-out data D
        D is a tensorflow dataset (from slices) of (X,Y,sz)
        """
        return -self.elbo_avg(D["X"], D["Y"], sz=D["sz"], S=S, chol=chol)

    def predict(self, Dtr, Dval=None, S=10):
        """
        Here Dtr,Dval should be raw counts (not normalized or log-transformed)
        returns the predicted training data mean and validation data mean
        on the original count scale
        """
        Mu_tr = misc.t2np(self.sample_predictive_mean(Dtr["X"], sz=Dtr["sz"], S=S))
        if self.lik == "gau":
            sz_tr = Dtr["Y"].sum(axis=1)
            # note self.feature_means is None if self.nonneg=True
            misc.reverse_normalization(
                Mu_tr,
                feature_means=self.feature_means,
                transform=np.expm1,
                sz=sz_tr,
                inplace=True,
            )
        if Dval:
            Mu_val = misc.t2np(
                self.sample_predictive_mean(Dval["X"], sz=Dval["sz"], S=S)
            )
            if self.lik == "gau":
                sz_val = Dval["Y"].sum(axis=1)
                misc.reverse_normalization(
                    Mu_val,
                    feature_means=self.feature_means,
                    transform=np.expm1,
                    sz=sz_val,
                    inplace=True,
                )
        else:
            Mu_val = None
        return Mu_tr, Mu_val


def smooth_spatial_factors(F, Z, X=None):
    """
    F: real-valued factors (ie on the log scale for NSF)
    Z: inducing point locations
    X: spatial coordinates
    """
    M = Z.shape[0]
    if X is None:  # no spatial coordinates, just use the mean
        beta0 = F.mean(axis=0)
        U = np.tile(beta0, [M, 1])
        beta = None
    else:  # spatial coordinates
        lr = LinearRegression().fit(X, F)
        beta0 = lr.intercept_
        beta = lr.coef_
        nn = max(2, ceil(X.shape[0] / M))
        knn = KNeighborsRegressor(n_neighbors=nn).fit(X, F)
        U = knn.predict(Z)
    return U, beta0, beta


def init_nsf_with_nmf(
    fit, Y, X=None, sz=1, pseudocount=1e-2, factors=None, loadings=None, shrinkage=0.2
):
    L = fit.W.shape[1]
    # M = fit.Z.shape[0] #number of inducing points
    kw = likelihoods.choose_nmf_pars(fit.lik)
    F, W = nnfu.regularized_nmf(
        Y,
        L,
        sz=sz,
        pseudocount=pseudocount,
        factors=factors,
        loadings=loadings,
        shrinkage=shrinkage,
        **kw
    )
    # eF = factors
    # W = loadings
    # if eF is None or W is None:
    #   kw = likelihoods.choose_nmf_pars(fit.lik)
    #   nmf = NMF(L,**kw)
    #   eF = nmf.fit_transform(Y)#/sz
    #   W = nmf.components_.T
    # W = postprocess.shrink_loadings(W, shrinkage=shrinkage)
    # wsum = W.sum(axis=0)
    # eF = postprocess.shrink_factors(eF*wsum, shrinkage=shrinkage)
    # F = np.log(pseudocount+eF)-np.log(sz)
    # beta0 = fit.beta0.numpy().flatten()
    # wt_to_W = F.mean(axis=0)- beta0
    # F-= wt_to_W
    # W*= np.exp(wt_to_W-np.log(wsum))
    # W,wsum = normalize_cols(W)
    # eF *= wsum
    # eFm2 = eF.mean()/2
    # eF/=eFm2
    # W*=eFm2
    # F = np.log(pseudocount+eF)
    fit.set_loadings(W)
    U, beta0, beta = smooth_spatial_factors(F, fit.Z.numpy(), X=X)
    # if X is None: #no spatial coordinates, just use the mean
    #   beta0 = F.mean(axis=0)
    #   U = np.tile(beta0,[M,1])
    # else: #spatial coordinates
    #   lr = LinearRegression().fit(X,F)
    #   beta0 = lr.intercept_
    #   fit.beta.assign(lr.coef_,read_value=False)
    #   nn = max(2, ceil(X.shape[0]/M))
    #   knn = KNeighborsRegressor(n_neighbors=nn).fit(X,F)
    #   U = knn.predict(fit.Z.numpy())
    fit.beta0.assign(beta0[:, None], read_value=False)
    fit.delta.assign(U.T, read_value=False)
    if beta is not None:
        fit.beta.assign(beta, read_value=False)
