import numpy as np

import pandas as pd

import numpy.random as npr

import torch

from scipy.special import xlogy

def rbf_kernel(

    x1, x2, lengthscale_unconstrained, output_variance_unconstrained, diag=False

):

    lengthscale = torch.exp(lengthscale_unconstrained)

    output_variance = torch.exp(output_variance_unconstrained)

    if diag:

        diffs = x1 - x2

    else:

        diffs = x1.unsqueeze(-2) - x2.unsqueeze(-3)

    K = output_variance * torch.exp(

        -0.5 * torch.sum(torch.square(diffs / lengthscale), dim=-1)

    )

    return K

def rbf_kernel_numpy(x, xp, kernel_params):

    output_scale = np.exp(kernel_params[0])

    lengthscales = np.exp(kernel_params[1:])

    diffs = np.expand_dims(x / lengthscales, 1) - np.expand_dims(xp / lengthscales, 0)

    return output_scale * np.exp(-0.5 * np.sum(diffs**2, axis=2))

def matern12_kernel(

    x1, x2, lengthscale_unconstrained, output_variance_unconstrained, diag=False

):

    lengthscale = torch.exp(lengthscale_unconstrained)

    output_variance = torch.exp(output_variance_unconstrained)

    if diag:

        diffs = x1 - x2

    else:

        diffs = x1.unsqueeze(-2) - x2.unsqueeze(-3)

    eps = 1e-10

    dists = torch.sqrt(torch.sum(torch.square(diffs), dim=-1) + eps)

    return output_variance * torch.exp(-0.5 * dists / lengthscale)

def matern32_kernel(

    x1, x2, lengthscale_unconstrained, output_variance_unconstrained, diag=False

):

    lengthscale = torch.exp(lengthscale_unconstrained)

    output_variance = torch.exp(output_variance_unconstrained)

    if diag:

        diffs = x1 - x2

    else:

        diffs = x1.unsqueeze(-2) - x2.unsqueeze(-3)

    eps = 1e-10

    dists = torch.sqrt(torch.sum(torch.square(diffs), dim=-1) + eps)

    inner_term = np.sqrt(3.0) * dists / lengthscale

    K = output_variance * (1 + inner_term) * torch.exp(-inner_term)

    return K

def polar_warp(X, r, theta):

    return np.array([X[:, 0] + r * np.cos(theta), X[:, 1] + r * np.sin(theta)]).T

def get_st_coordinates(df):

    """

    Extracts spatial coordinates from ST data with index in 'AxB' type format.

    Return: pandas dataframe of coordinates

    """

    coor = []

    for spot in df.index:

        coordinates = spot.split("x")

        coordinates = [float(i) for i in coordinates]

        coor.append(coordinates)

    return np.array(coor)

def compute_distance(X1, X2):

    return np.mean(np.sqrt(np.sum((X1 - X2) ** 2, axis=1)))

def make_pinwheel(

    radial_std, tangential_std, num_classes, num_per_class, rate, rs=npr.RandomState(0)

):

    """Based on code by Ryan P. Adams."""

    rads = np.linspace(0, 2 * np.pi, num_classes, endpoint=False)

    features = rs.randn(num_classes * num_per_class, 2) * np.array(

        [radial_std, tangential_std]

    )

    features[:, 0] += 1

    labels = np.repeat(np.arange(num_classes), num_per_class)

    angles = rads[labels] + rate * np.exp(features[:, 0])

    rotations = np.stack(

        [np.cos(angles), -np.sin(angles), np.sin(angles), np.cos(angles)]

    )

    rotations = np.reshape(rotations.T, (-1, 2, 2))

    return np.einsum("ti,tij->tj", features, rotations)

class ConvergenceChecker(object):

    def __init__(self, span, dtp="float64"):

        self.span = span

        x = np.arange(span, dtype=dtp)

        x -= x.mean()

        X = np.column_stack((np.ones(shape=x.shape), x, x**2, x**3))

        self.U = np.linalg.svd(X, full_matrices=False)[0]

    def smooth(self, y):

        return self.U @ (self.U.T @ y)

    def subset(self, y, idx=-1):

        span = self.U.shape[0]

        lo = idx - span + 1

        if idx == -1:

            return y[lo:]

        else:

            return y[lo : (idx + 1)]

    def relative_change(self, y, idx=-1, smooth=True):

        y = self.subset(y, idx=idx)

        if smooth:

            y = self.smooth(y)

        prev = y[-2]

        return (y[-1] - prev) / (0.1 + abs(prev))

    def converged(self, y, tol=1e-4, **kwargs):

        return abs(self.relative_change(y, **kwargs)) < tol

    def relative_change_all(self, y, smooth=True):

        n = len(y)

        span = self.U.shape[0]

        cc = np.tile([np.nan], n)

        for i in range(span, n):

            cc[i] = self.relative_change(y, idx=i, smooth=smooth)

        return cc

    def converged_all(self, y, tol=1e-4, smooth=True):

        cc = self.relative_change_all(y, smooth=smooth)

        return np.abs(cc) < tol

# Function for computing size factors

def compute_size_factors(m):

    # given matrix m with samples in the columns

    # compute size factors

    sz = np.sum(m.values, axis=0)  # column sums (sum of counts in each cell)

    lsz = np.log(sz)

    # make geometric mean of sz be 1 for poisson

    sz_poisson = np.exp(lsz - np.mean(lsz))

    return sz_poisson

def poisson_deviance(X, sz):

    LP = X.values / sz  # recycling

    # import ipdb; ipdb.set_trace()

    LP[LP > 0] = np.log(LP[LP > 0])  # log transform nonzero elements only

    # Transpose to make features in cols, observations in rows

    X = X.T

    ll_sat = np.sum(np.multiply(X, LP.T), axis=0)

    feature_sums = np.sum(X, axis=0)

    ll_null = feature_sums * np.log(feature_sums / np.sum(sz))

    return 2 * (ll_sat - ll_null)

def deviance_feature_selection(X):

    # Remove cells without any counts

    X = X[np.sum(X, axis=1) > 0]

    # Compute size factors

    sz = compute_size_factors(X)

    # Compute deviances

    devs = poisson_deviance(X, sz)

    # Get associated gene names

    gene_names = X.index.values

    assert gene_names.shape[0] == devs.values.shape[0]

    return devs.values, gene_names

def deviance_residuals(x, theta, mu=None):

    """Computes deviance residuals for NB model with a fixed theta"""

    if mu is None:

        counts_sum0 = np.sum(x, axis=0, keepdims=True)

        counts_sum1 = np.sum(x, axis=1, keepdims=True)

        counts_sum = np.sum(x)

        # get residuals

        mu = counts_sum1 @ counts_sum0 / counts_sum

    def remove_negatives(sqrt_term):

        negatives_idx = sqrt_term < 0

        if np.any(negatives_idx):

            n_negatives = np.sum(negatives_idx)

            print(

                "Setting %u negative sqrt term values to 0 (%f%%)"

                % (n_negatives, n_negatives / np.product(sqrt_term.shape))

            )

            sqrt_term[negatives_idx] = 0

    if np.isinf(theta):  ### POISSON

        x_minus_mu = x - mu

        sqrt_term = 2 * (

            xlogy(x, x / mu) - x_minus_mu

        )  # xlogy(x,x/mu) computes xlog(x/mu) and returns 0 if x=0

        remove_negatives(sqrt_term)

        dev = np.sign(x_minus_mu) * np.sqrt(sqrt_term)

    else:  ### NEG BIN

        x_plus_theta = x + theta

        sqrt_term = 2 * (

            xlogy(x, x / mu) - (x_plus_theta) * np.log(x_plus_theta / (mu + theta))

        )  # xlogy(x,x/mu) computes xlog(x/mu) and returns 0 if x=0

        remove_negatives(sqrt_term)

        dev = np.sign(x - mu) * np.sqrt(sqrt_term)

    return dev

def pearson_residuals(counts, theta, clipping=True):

    """Computes analytical residuals for NB model with a fixed theta, clipping outlier residuals to sqrt(N)"""

    counts_sum0 = np.sum(counts, axis=0, keepdims=True)

    counts_sum1 = np.sum(counts, axis=1, keepdims=True)

    counts_sum = np.sum(counts)

    # get residuals

    mu = counts_sum1 @ counts_sum0 / counts_sum

    z = (counts - mu) / np.sqrt(mu + mu**2 / theta)

    # clip to sqrt(n)

    if clipping:

        n = counts.shape[0]

        z[z > np.sqrt(n)] = np.sqrt(n)

        z[z < -np.sqrt(n)] = -np.sqrt(n)

    return z

class LossNotDecreasingChecker:

    def __init__(self, max_epochs, atol=1e-2, window_size=10):

        self.max_epochs = max_epochs

        self.atol = atol

        self.window_size = window_size

        self.decrease_in_loss = np.zeros(max_epochs)

        self.average_decrease_in_loss = np.zeros(max_epochs)

    def check_loss(self, iternum, loss_trace):

        if iternum >= 1:

            self.decrease_in_loss[iternum] = (

                loss_trace[iternum - 1] - loss_trace[iternum]

            )

            if iternum >= self.window_size:

                self.average_decrease_in_loss[iternum] = np.mean(

                    self.decrease_in_loss[iternum - self.window_size + 1 : iternum]

                )

                has_converged = self.average_decrease_in_loss[iternum] < self.atol

                return has_converged

        return False