# -*- coding: utf-8 -*-
"""
Regularized Unbalanced OT solvers
"""

# Author: Hicham Janati <hicham.janati@inria.fr>
#         Laetitia Chapel <laetitia.chapel@univ-ubs.fr>
#         Quang Huy Tran <quang-huy.tran@univ-ubs.fr>
#
# License: MIT License

import warnings
import numpy as np
from scipy.optimize import minimize, Bounds

from ..backend import get_backend
from ..utils import list_to_array, get_parameter_pair


def _get_loss_unbalanced(a, b, c, M, reg, reg_m1, reg_m2, reg_div="kl", regm_div="kl"):
    """
    Return loss function for the L-BFGS-B solver

    .. note:: This function will be fed into scipy.optimize, so all input arrays must be Numpy arrays.

    Parameters
    ----------
    a : array-like (dim_a,)
        Unnormalized histogram of dimension `dim_a`
    b : array-like (dim_b,)
        Unnormalized histogram of dimension `dim_b`
    M : array-like (dim_a, dim_b)
        loss matrix
    reg: float
        regularization term >=0
    c : array-like (dim_a, dim_b), optional (default = None)
        Reference measure for the regularization.
        If None, then use :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T`.
    reg_m1: float
        Marginal relaxation term with respect to the first marginal:
        nonnegative (including 0) but cannot be infinity.
    reg_m2: float
        Marginal relaxation term with respect to the second marginal:
        nonnegative (including 0) but cannot be infinity.
    reg_div: string, optional
        Divergence used for regularization.
        Can take three values: 'entropy' (negative entropy), or
        'kl' (Kullback-Leibler) or 'l2' (half-squared) or a tuple
        of two calable functions returning the reg term and its derivative.
        Note that the callable functions should be able to handle Numpy arrays
        and not tesors from the backend
    regm_div: string, optional
        Divergence to quantify the difference between the marginals.
        Can take three values: 'kl' (Kullback-Leibler) or 'l2' (half-squared) or 'tv' (Total Variation)

    Returns
    -------
    Loss function (scipy.optimize compatible) for regularized unbalanced OT
    """

    m, n = M.shape
    nx_numpy = get_backend(M, a, b)

    def reg_l2(G):
        return np.sum((G - c) ** 2) / 2

    def grad_l2(G):
        return G - c

    def reg_kl(G):
        return nx_numpy.kl_div(G, c, mass=True)

    def grad_kl(G):
        return np.log(G / c + 1e-16)

    def reg_entropy(G):
        return np.sum(G * np.log(G + 1e-16)) - np.sum(G)

    def grad_entropy(G):
        return np.log(G + 1e-16)

    if reg_div == "kl":
        reg_fun = reg_kl
        grad_reg_fun = grad_kl
    elif reg_div == "entropy":
        reg_fun = reg_entropy
        grad_reg_fun = grad_entropy
    elif isinstance(reg_div, tuple):
        reg_fun = reg_div[0]
        grad_reg_fun = reg_div[1]
    else:
        reg_fun = reg_l2
        grad_reg_fun = grad_l2

    def marg_l2(G):
        return reg_m1 * 0.5 * np.sum((G.sum(1) - a) ** 2) + reg_m2 * 0.5 * np.sum(
            (G.sum(0) - b) ** 2
        )

    def grad_marg_l2(G):
        return reg_m1 * np.outer((G.sum(1) - a), np.ones(n)) + reg_m2 * np.outer(
            np.ones(m), (G.sum(0) - b)
        )

    def marg_kl(G):
        return reg_m1 * nx_numpy.kl_div(
            G.sum(1), a, mass=True
        ) + reg_m2 * nx_numpy.kl_div(G.sum(0), b, mass=True)

    def grad_marg_kl(G):
        return reg_m1 * np.outer(
            np.log(G.sum(1) / a + 1e-16), np.ones(n)
        ) + reg_m2 * np.outer(np.ones(m), np.log(G.sum(0) / b + 1e-16))

    def marg_tv(G):
        return reg_m1 * np.sum(np.abs(G.sum(1) - a)) + reg_m2 * np.sum(
            np.abs(G.sum(0) - b)
        )

    def grad_marg_tv(G):
        return reg_m1 * np.outer(np.sign(G.sum(1) - a), np.ones(n)) + reg_m2 * np.outer(
            np.ones(m), np.sign(G.sum(0) - b)
        )

    if regm_div == "kl":
        regm_fun = marg_kl
        grad_regm_fun = grad_marg_kl
    elif regm_div == "tv":
        regm_fun = marg_tv
        grad_regm_fun = grad_marg_tv
    else:
        regm_fun = marg_l2
        grad_regm_fun = grad_marg_l2

    def _func(G):
        G = G.reshape((m, n))

        # compute loss
        val = np.sum(G * M) + regm_fun(G)
        if reg > 0:
            val = val + reg * reg_fun(G)
        # compute gradient
        grad = M + grad_regm_fun(G)
        if reg > 0:
            grad = grad + reg * grad_reg_fun(G)

        return val, grad.ravel()

    return _func


def lbfgsb_unbalanced(
    a,
    b,
    M,
    reg,
    reg_m,
    c=None,
    reg_div="kl",
    regm_div="kl",
    G0=None,
    numItermax=1000,
    stopThr=1e-15,
    method="L-BFGS-B",
    verbose=False,
    log=False,
):
    r"""
    Solve the unbalanced optimal transport problem and return the OT plan using L-BFGS-B algorithm.
    The function solves the following optimization problem:

    .. math::
        W = \arg \min_\gamma \quad \langle \gamma, \mathbf{M} \rangle_F +
        \mathrm{reg} \mathrm{div}(\gamma, \mathbf{c}) +
        \mathrm{reg_{m1}} \cdot \mathrm{div_m}(\gamma \mathbf{1}, \mathbf{a}) +
        \mathrm{reg_{m2}} \cdot \mathrm{div_m}(\gamma^T \mathbf{1}, \mathbf{b})

        s.t.
             \gamma \geq 0

    where:

    - :math:`\mathbf{M}` is the (`dim_a`, `dim_b`) metric cost matrix
    - :math:`\mathbf{a}` and :math:`\mathbf{b}` are source and target unbalanced distributions
    - :math:`\mathbf{c}` is a reference distribution for the regularization
    - :math:`\mathrm{div_m}` is a divergence, either Kullback-Leibler divergence,
    or half-squared :math:`\ell_2` divergence, or Total variation
    - :math:`\mathrm{div}` is a divergence, either Kullback-Leibler divergence,
    or half-squared :math:`\ell_2` divergence

    .. note:: This function is backend-compatible and will work on arrays
        from all compatible backends. First, it converts all arrays into Numpy arrays,
        then uses the L-BFGS-B algorithm from scipy.optimize to solve the optimization problem.

    Parameters
    ----------
    a : array-like (dim_a,)
        Unnormalized histogram of dimension `dim_a`
        If `a` is an empty list or array ([]),
        then `a` is set to uniform distribution.
    b : array-like (dim_b,)
        Unnormalized histogram of dimension `dim_b`
        If `b` is an empty list or array ([]),
        then `b` is set to uniform distribution.
    M : array-like (dim_a, dim_b)
        loss matrix
    reg: float
        regularization term >=0
    c : array-like (dim_a, dim_b), optional (default = None)
        Reference measure for the regularization.
        If None, then use :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T`.
    reg_m: float or indexable object of length 1 or 2
        Marginal relaxation term: nonnegative (including 0) but cannot be infinity.
        If :math:`\mathrm{reg_{m}}` is a scalar or an indexable object of length 1,
        then the same :math:`\mathrm{reg_{m}}` is applied to both marginal relaxations.
        If :math:`\mathrm{reg_{m}}` is an array, it must be a Numpy array.
    reg_div: string, optional
        Divergence used for regularization.
        Can take three values: 'entropy' (negative entropy), or
        'kl' (Kullback-Leibler) or 'l2' (half-squared) or a tuple
        of two calable functions returning the reg term and its derivative.
        Note that the callable functions should be able to handle Numpy arrays
        and not tesors from the backend
    regm_div: string, optional
        Divergence to quantify the difference between the marginals.
        Can take three values: 'kl' (Kullback-Leibler) or 'l2' (half-squared) or 'tv' (Total Variation)
    G0: array-like (dim_a, dim_b)
        Initialization of the transport matrix
    numItermax : int, optional
        Max number of iterations
    stopThr : float, optional
        Stop threshold on error (> 0)
    verbose : bool, optional
        Print information along iterations
    log : bool, optional
        record log if True

    Returns
    -------
    gamma : (dim_a, dim_b) array-like
            Optimal transportation matrix for the given parameters
    log : dict
        log dictionary returned only if `log` is `True`

    Examples
    --------
    >>> import ot
    >>> import numpy as np
    >>> a=[.5, .5]
    >>> b=[.5, .5]
    >>> M=[[1., 36.],[9., 4.]]
    >>> np.round(ot.unbalanced.lbfgsb_unbalanced(a, b, M, reg=0, reg_m=5, reg_div='kl', regm_div='kl'), 2)
    array([[0.45, 0.  ],
           [0.  , 0.34]])
    >>> np.round(ot.unbalanced.lbfgsb_unbalanced(a, b, M, reg=0, reg_m=5, reg_div='l2', regm_div='l2'), 2)
    array([[0.4, 0. ],
           [0. , 0.1]])

    References
    ----------
    .. [41] Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021).
        Unbalanced optimal transport through non-negative penalized
        linear regression. NeurIPS.
    See Also
    --------
    ot.lp.emd2 : Unregularized OT loss
    ot.unbalanced.sinkhorn_unbalanced2 : Entropic regularized OT loss
    """

    # wrap the callable function to handle numpy arrays
    if isinstance(reg_div, tuple):
        f0, df0 = reg_div
        try:
            f0(G0)
            df0(G0)
        except BaseException:
            warnings.warn(
                "The callable functions should be able to handle numpy arrays, wrapper ar added to handle this which comes with overhead"
            )

            def f(x):
                return nx.to_numpy(f0(nx.from_numpy(x, type_as=M0)))

            def df(x):
                return nx.to_numpy(df0(nx.from_numpy(x, type_as=M0)))

            reg_div = (f, df)

    else:
        reg_div = reg_div.lower()
        if reg_div not in ["entropy", "kl", "l2"]:
            raise ValueError(
                "Unknown reg_div = {}. Must be either 'entropy', 'kl' or 'l2', or a tuple".format(
                    reg_div
                )
            )

    regm_div = regm_div.lower()
    if regm_div not in ["kl", "l2", "tv"]:
        raise ValueError(
            "Unknown regm_div = {}. Must be either 'kl', 'l2' or 'tv'".format(regm_div)
        )

    reg_m1, reg_m2 = get_parameter_pair(reg_m)

    M, a, b = list_to_array(M, a, b)
    nx = get_backend(M, a, b)
    M0 = M

    dim_a, dim_b = M.shape

    if len(a) == 0:
        a = nx.ones(dim_a, type_as=M) / dim_a
    if len(b) == 0:
        b = nx.ones(dim_b, type_as=M) / dim_b

    # convert to numpy
    a, b, M, reg_m1, reg_m2, reg = nx.to_numpy(a, b, M, reg_m1, reg_m2, reg)
    G0 = a[:, None] * b[None, :] if G0 is None else nx.to_numpy(G0)
    c = a[:, None] * b[None, :] if c is None else nx.to_numpy(c)

    _func = _get_loss_unbalanced(a, b, c, M, reg, reg_m1, reg_m2, reg_div, regm_div)

    res = minimize(
        _func,
        G0.ravel(),
        method=method,
        jac=True,
        bounds=Bounds(0, np.inf),
        tol=stopThr,
        options=dict(maxiter=numItermax, disp=verbose),
    )

    G = nx.from_numpy(res.x.reshape(M.shape), type_as=M0)

    if log:
        log = {"cost": nx.sum(G * M), "res": res}
        log["total_cost"] = nx.from_numpy(res.fun, type_as=M0)
        return G, log
    else:
        return G


def lbfgsb_unbalanced2(
    a,
    b,
    M,
    reg,
    reg_m,
    c=None,
    reg_div="kl",
    regm_div="kl",
    G0=None,
    returnCost="linear",
    numItermax=1000,
    stopThr=1e-15,
    method="L-BFGS-B",
    verbose=False,
    log=False,
):
    r"""
    Solve the unbalanced optimal transport problem and return the OT cost using L-BFGS-B.
    The function solves the following optimization problem:

    .. math::
        \min_\gamma \quad \langle \gamma, \mathbf{M} \rangle_F +
        \mathrm{reg} \mathrm{div}(\gamma, \mathbf{c}) +
        \mathrm{reg_{m1}} \cdot \mathrm{div_m}(\gamma \mathbf{1}, \mathbf{a}) +
        \mathrm{reg_{m2}} \cdot \mathrm{div_m}(\gamma^T \mathbf{1}, \mathbf{b})

        s.t.
             \gamma \geq 0

    where:

    - :math:`\mathbf{M}` is the (`dim_a`, `dim_b`) metric cost matrix
    - :math:`\mathbf{a}` and :math:`\mathbf{b}` are source and target unbalanced distributions
    - :math:`\mathbf{c}` is a reference distribution for the regularization
    - :math:`\mathrm{div_m}` is a divergence, either Kullback-Leibler divergence,
    or half-squared :math:`\ell_2` divergence, or Total variation
    - :math:`\mathrm{div}` is a divergence, either Kullback-Leibler divergence,
    or half-squared :math:`\ell_2` divergence

    .. note:: This function is backend-compatible and will work on arrays
        from all compatible backends. First, it converts all arrays into Numpy arrays,
        then uses the L-BFGS-B algorithm from scipy.optimize to solve the optimization problem.

    Parameters
    ----------
    a : array-like (dim_a,)
        Unnormalized histogram of dimension `dim_a`
        If `a` is an empty list or array ([]),
        then `a` is set to uniform distribution.
    b : array-like (dim_b,)
        Unnormalized histogram of dimension `dim_b`
        If `b` is an empty list or array ([]),
        then `b` is set to uniform distribution.
    M : array-like (dim_a, dim_b)
        loss matrix
    reg: float
        regularization term >=0
    c : array-like (dim_a, dim_b), optional (default = None)
        Reference measure for the regularization.
        If None, then use :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T`.
    reg_m: float or indexable object of length 1 or 2
        Marginal relaxation term: nonnegative (including 0) but cannot be infinity.
        If :math:`\mathrm{reg_{m}}` is a scalar or an indexable object of length 1,
        then the same :math:`\mathrm{reg_{m}}` is applied to both marginal relaxations.
        If :math:`\mathrm{reg_{m}}` is an array, it must be a Numpy array.
    reg_div: string, optional
        Divergence used for regularization.
        Can take three values: 'entropy' (negative entropy), or
        'kl' (Kullback-Leibler) or 'l2' (half-squared) or a tuple
        of two calable functions returning the reg term and its derivative.
        Note that the callable functions should be able to handle Numpy arrays
        and not tesors from the backend
    regm_div: string, optional
        Divergence to quantify the difference between the marginals.
        Can take three values: 'kl' (Kullback-Leibler) or 'l2' (half-squared) or 'tv' (Total Variation)
    G0: array-like (dim_a, dim_b)
        Initialization of the transport matrix
    returnCost: string, optional (default = "linear")
        If `returnCost` = "linear", then return the linear part of the unbalanced OT loss.
        If `returnCost` = "total", then return the total unbalanced OT loss.
    numItermax : int, optional
        Max number of iterations
    stopThr : float, optional
        Stop threshold on error (> 0)
    verbose : bool, optional
        Print information along iterations
    log : bool, optional
        record log if True

    Returns
    -------
    ot_cost : array-like
        the OT cost between :math:`\mathbf{a}` and :math:`\mathbf{b}`
    log : dict
        log dictionary returned only if `log` is `True`

    Examples
    --------
    >>> import ot
    >>> import numpy as np
    >>> a=[.5, .5]
    >>> b=[.5, .5]
    >>> M=[[1., 36.],[9., 4.]]
    >>> np.round(ot.unbalanced.lbfgsb_unbalanced2(a, b, M, reg=0, reg_m=5, reg_div='kl', regm_div='kl'), 2)
    1.79
    >>> np.round(ot.unbalanced.lbfgsb_unbalanced2(a, b, M, reg=0, reg_m=5, reg_div='l2', regm_div='l2'), 2)
    0.8

    References
    ----------
    .. [41] Chapel, L., Flamary, R., Wu, H., Févotte, C., and Gasso, G. (2021).
        Unbalanced optimal transport through non-negative penalized
        linear regression. NeurIPS.
    See Also
    --------
    ot.lp.emd2 : Unregularized OT loss
    ot.unbalanced.sinkhorn_unbalanced2 : Entropic regularized OT loss
    """

    _, log_lbfgs = lbfgsb_unbalanced(
        a=a,
        b=b,
        M=M,
        reg=reg,
        reg_m=reg_m,
        c=c,
        reg_div=reg_div,
        regm_div=regm_div,
        G0=G0,
        numItermax=numItermax,
        stopThr=stopThr,
        method=method,
        verbose=verbose,
        log=True,
    )

    if returnCost == "linear":
        cost = log_lbfgs["cost"]
    elif returnCost == "total":
        cost = log_lbfgs["total_cost"]
    else:
        raise ValueError("Unknown returnCost = {}".format(returnCost))

    if log:
        return cost, log_lbfgs
    else:
        return cost