from matplotlib import pyplot as plt
import torch
import numpy as np


def groundmetric(width, height=None, p=2, normed=False):
    """Compute ground metric matrix on the 1D grid 0:`n_features`.

    Parameters
    ----------
    width : int
        width of the matrix
    height : int (optional, default equal to width)
        height of the matrix
    p: int > 0.
        Power to raise the pairwise distance metrix. Quadratic by default.
    normed: boolean (default True)
        If True, the matrix is divided by its median.

    Returns
    -------
    M: 2D array (width, height).

    """
    if height is None:
        height = width
    x = np.arange(0, width).reshape(-1, 1).astype(float)
    y = np.arange(0, height).reshape(-1, 1).astype(float)
    xx, yy = np.meshgrid(x, y)
    M = abs(xx - yy) ** p
    if normed:
        M /= np.median(M)
    return M


def groundmetric_img(width, height=None, p=2, normed=False):
    """Compute ground metric for convolutional Wasserstein.

    Parameters
    ----------

    width, height : int,
        shape of images
    p: int, optional (default 2)
        Power to raise the pairwise distance metrix. Quadratic by default.
    normed: boolean (default True)
        If True, the matrix is divided by its median.
    Returns
    -------
    M: 2D array (width, height).
    """
    if height is None:
        height = width
    M = groundmetric(width, height, p=2, normed=False)
    if normed:
        Mlarge = groundmetric2d(width, height, p=2, normed=False)
        median = np.median(Mlarge)
        M /= median
    return M


def groundmetric2d(width, height=None, p=1, normed=False):
    """Compute ground metric matrix on the 2D grid (width, height).

    Parameters
    ----------
    n_features: int > 0.
    p: int > 0.
        Power to raise the pairwise distance metrix. Quadratic by default.
    normed: boolean (default True)
        If True, the matrix is divided by its median.

    Returns
    -------
    M: 2D array (n_features, n_features).

    """
    if height is None:
        height = width
    n_features = width * height
    M = groundmetric(width, height, p=2, normed=False)
    M = M[:, np.newaxis, :, np.newaxis] + M[np.newaxis, :, np.newaxis, :]
    M = M.reshape(n_features, n_features) ** (p / 2)

    if normed:
        M /= np.median(M)
    return M


def compute_gamma(tau, M):
    """Compute the sufficient KL weight for a full mass minimum."""
    xp = get_module(M)
    return max(0., - M.max() / (2 * xp.log(tau)))


def show_ts(X, axes=None, title="", cmap=None, show_time=False,
            normalize=True):
    T, p, p = X.shape
    vmin = X.min()
    vmax = X.max()
    if axes is None:
        f, axes = plt.subplots(T, 1)
    for i, (img, ax) in enumerate(zip(X, axes)):
        ax.imshow(img, cmap=cmap, vmin=vmin, vmax=vmax)
        ax.set_xticks([])
        ax.set_yticks([])
        if show_time:
            ax.set_title(f"t = {i}")
    axes[0].set_ylabel(title)


def prox_simplex_all(x):
    """Projection on the simplex."""
    shape = x.shape
    if x.ndimension() == 2:
        x = x.reshape(1, -1)
    else:
        x = x.reshape(len(x), -1)
    size = x[0].nelement()
    u = torch.sort(x, dim=1, descending=True)[0]
    cu = torch.cumsum(u, dim=1) - 1
    ucu = u - cu / torch.arange(1., size + 1., dtype=x.dtype,
                                device=x.device)[None, :]
    cucu = torch.cumsum(ucu, dim=1)
    ss = torch.argmax(cucu, dim=1)
    vv = cu[torch.arange(len(x)), ss] / (ss.type(x.dtype).to(x.device) + 1.)
    p = torch.clamp(x - vv[:, None], 0., None)
    return p.reshape(shape)


def get_module(x):
    if type(x) == np.ndarray:
        return np
    else:
        return torch


def kl(p, q):
    """Compute Kullback-Leibler divergence.

    Compute the element-wise sum of:
    `p * log(p / q) + p - q`.

    Parameters
    ----------
    p: array-like.
        must be positive.
    q: array-like.
        must be positive, same shape and dimension of `p`.

    """
    xp = get_module(p)
    logpq = xp.zeros_like(p)
    logpq[p > 0] = xp.log(p[p > 0] / (q[p > 0] + 1e-300))
    kl = (p * logpq + q - p)

    return kl


def wklobjective_plan(plan, p, q, K, epsilon, gamma):
    """Compute unbalanced ot objective function naively."""
    f = epsilon * kl(plan, K).sum()
    margs = kl(plan.sum(1), p) + kl(plan.sum(0), q)
    f += gamma * margs.sum()

    return f


def wklobjective(a, Kb, p, q, f0, epsilon, gamma, u=0):
    """Compute unbalanced ot objective function for solver monitoring."""
    xp = get_module(a)
    aKb = a * Kb
    f = gamma * kl(aKb, p)
    f += (aKb * (epsilon * xp.log(a + 1e-300) + u - epsilon - gamma))
    f += f0
    f += gamma * q

    return f


def tonumpy(x):
    if isinstance(x, np.ndarray):
        return x
    return x.detach().cpu().numpy()


def fancy_imshow(matrix, ax=None):
    """imshow with values"""
    # The normal figure
    if ax is None:
        f, ax = plt.subplots(1, 1)
    ax.imshow(matrix, interpolation='None',
              cmap='viridis', alpha=0.6)
    n_cols, n_rows = matrix.shape
    for j in range(n_cols):
        for i in range(n_rows):
            label = np.round(matrix.T[i, j], 2)
            ax.text(i, j, label, color='black', ha='center', va='center')
    return ax


def generate_time_series(dimension, n_times, seed=42, device="cpu",
                         dtype=torch.float64, required_grad_X=True):
    """Fake time series."""
    rng = np.random.RandomState(seed)
    x = np.zeros((n_times, dimension, dimension)) + 0.001
    y = np.zeros((n_times, dimension, dimension)) + 0.001
    ts = np.arange(n_times)
    index_x, index_y = rng.randint(dimension, size=(2, n_times))
    x[ts, index_x, index_y] = rng.rand(n_times) + 1.
    y[ts, index_y, index_x] = rng.rand(n_times) + 1.

    X = torch.tensor(x, dtype=dtype, device=device,
                     requires_grad=required_grad_X)
    Y = torch.tensor(y, dtype=dtype, device=device, requires_grad=False)
    return X, Y


def cost_matrix(dimension, device="cpu", dtype=torch.float64):
    Mimg = groundmetric(dimension, p=2, normed=False)
    M = groundmetric2d(dimension, p=2, normed=False)
    m = np.median(M)
    Mimg /= m
    M /= m
    Mimg = torch.tensor(Mimg, dtype=dtype, device=device, requires_grad=False)
    M = torch.tensor(M, dtype=dtype, device=device, requires_grad=False)

    return M, Mimg
