# -*- coding: utf-8 -*-

# Copyright © 2017-2021, Institut Pasteur
#   Contributor: François Laurent

# This file is part of the TRamWAy software available at
# "https://github.com/DecBayComp/TRamWAy" and is distributed under
# the terms of the CeCILL license as circulated at the following URL
# "http://www.cecill.info/licenses.en.html".

# The fact that you are presently reading this means that you have had
# knowledge of the CeCILL license and that you accept its terms.


from .base import *
from .gradient import *
from warnings import warn
from math import pi, log
import numpy as np
import pandas as pd
from scipy.optimize import minimize
from collections import OrderedDict, namedtuple


setup = {
    "name": ("standard.d", "smooth.d", "regularized.d"),
    "provides": "d",
    "arguments": OrderedDict(
        (
            (
                "localization_error",
                (
                    "-e",
                    dict(
                        type=float,
                        help="localization precision (see also sigma; default is 0.03)",
                    ),
                ),
            ),
            (
                "diffusivity_prior",
                ("-d", dict(type=float, help="prior on the diffusivity")),
            ),
            (
                "jeffreys_prior",
                ("-j", dict(action="store_true", help="Jeffreys' prior")),
            ),
            (
                "min_diffusivity",
                dict(type=float, help="minimum diffusivity value allowed"),
            ),
            ("max_iter", dict(type=int, help="maximum number of iterations")),
            (
                "rgrad",
                dict(
                    help="alternative gradient for the regularization; can be 'delta'/'delta0' or 'delta1'"
                ),
            ),
            ("tol", dict(type=float, help="tolerance for scipy minimizer")),
        )
    ),
    "cell_sampling": "group",
}
setup_with_grad_arguments(setup)


Result = namedtuple("Result", "result success message ncalls")


def interruptible_minimize(func, prms, args, **kwargs):
    callback_cache = dict(
        ncalls=0,
    )

    def caching_callback(x, *args):
        callback_cache["result"] = x
        callback_cache["ncalls"] += 1
        return False

    kwargs["callback"] = caching_callback

    try:
        return minimize(func, prms, args=args, **kwargs)

    except KeyboardInterrupt:
        if "result" in callback_cache:
            res = callback_cache["result"]
            ctr = callback_cache["ncalls"]
            ret = Result(res, False, "interrupted", ctr)
            return ret
        else:
            raise


def smooth_d_neg_posterior(
    diffusivity,
    cells,
    sigma2,
    diffusivity_prior,
    jeffreys_prior,
    dt_mean,
    min_diffusivity,
    index,
    reverse_index,
    grad_kwargs,
):
    """
    Adapted from InferenceMAP's *dDDPosterior* procedure:

    .. code-block:: c++

        for (int a = 0; a < NUMBER_OF_ZONES; a++) {
            ZONES[a].gradDx = dvGradDx(DD,a);
            ZONES[a].gradDy = dvGradDy(DD,a);
            ZONES[a].priorActive = true;
        }

        for (int z = 0; z < NUMBER_OF_ZONES; z++) {
            const double gradDx = ZONES[z].gradDx;
            const double gradDy = ZONES[z].gradDy;
            const double D = DD[z];

            for (int j = 0; j < ZONES[z].translocations; j++) {
                const double dt = ZONES[z].dt[j];
                const double dx = ZONES[z].dx[j];
                const double dy = ZONES[z].dy[j];
                const double Dnoise = LOCALIZATION_ERROR*LOCALIZATION_ERROR/dt;

                result += - log(4.0*PI*(D + Dnoise)*dt) - ( dx*dx + dy*dy)/(4.0*(D+Dnoise)*dt);
            }

            if (ZONES[z].priorActive == true) {
                result -= D_PRIOR*(gradDx*gradDx*ZONES[z].areaX + gradDy*gradDy*ZONES[z].areaY);
                if (JEFFREYS_PRIOR == 1) {
                    result += 2.0*log(D) - 2.0*log(D*ZONES[z].dtMean + LOCALIZATION_ERROR*LOCALIZATION_ERROR);
                }
            }
        }

        return -result;

    """
    if min_diffusivity is not None:
        observed_min = np.min(diffusivity)
        if observed_min < min_diffusivity and not np.isclose(
            observed_min, min_diffusivity
        ):
            warn(DiffusivityWarning(observed_min, min_diffusivity))
    noise_dt = sigma2
    result = 0.0
    for j, i in enumerate(index):
        cell = cells[i]
        n = len(cell)
        # posterior calculations
        if cell.cache is None:
            cell.cache = dict(dr2=None)
        if cell.cache["dr2"] is None:
            cell.cache["dr2"] = np.sum(cell.dr * cell.dr, axis=1)  # dx**2 + dy**2 + ..
        D_dt = 4.0 * (diffusivity[j] * cell.dt + noise_dt)  # 4*(D+Dnoise)*dt
        result += n * log(pi) + np.sum(np.log(D_dt))  # sum(log(4*pi*Dtot*dt))
        result += np.sum(cell.cache["dr2"] / D_dt)  # sum((dx**2+dy**2+..)/(4*Dtot*dt))
        # prior
        if diffusivity_prior:
            # gradient of diffusivity
            gradD = cells.grad(i, diffusivity, reverse_index, **grad_kwargs)
            if gradD is not None:
                result += diffusivity_prior * cells.grad_sum(i, gradD * gradD)
    if jeffreys_prior:
        result += 2.0 * np.sum(np.log(diffusivity * dt_mean + sigma2))
    return result


def d_neg_posterior1(
    diffusivity,
    cells,
    sigma2,
    diffusivity_prior,
    jeffreys_prior,
    dt_mean,
    min_diffusivity,
    index,
    reverse_index,
    grad_kwargs,
):
    """
    Similar to :func:`smooth_d_neg_posterior`.
    The smoothing prior features an alternative spatial "gradient" implemented using
    :meth:`~tramway.inference.base.Distributed.local_variation` instead of
    :meth:`~tramway.inference.base.Distributed.grad`.
    """
    if min_diffusivity is not None:
        observed_min = np.min(diffusivity)
        if (
            observed_min < min_diffusivity
        ):  # and not np.isclose(observed_min, min_diffusivity):
            warn(DiffusivityWarning(observed_min, min_diffusivity))
    noise_dt = sigma2
    result = 0.0
    for j, i in enumerate(index):
        cell = cells[i]
        n = len(cell)
        # posterior calculations
        if cell.cache is None:
            cell.cache = dict(dr2=None)
        if cell.cache["dr2"] is None:
            cell.cache["dr2"] = np.sum(cell.dr * cell.dr, axis=1)  # dx**2 + dy**2 + ..
        D_dt = 4.0 * (diffusivity[j] * cell.dt + noise_dt)  # 4*(D+Dnoise)*dt
        result += n * log(pi) + np.sum(np.log(D_dt))  # sum(log(4*pi*Dtot*dt))
        result += np.sum(cell.cache["dr2"] / D_dt)  # sum((dx**2+dy**2+..)/(4*Dtot*dt))
        # prior
        if diffusivity_prior:
            # gradient of diffusivity
            deltaD = cells.local_variation(i, diffusivity, reverse_index, **grad_kwargs)
            if deltaD is not None:
                result += diffusivity_prior * cells.grad_sum(i, deltaD * deltaD)
    if jeffreys_prior:
        result += 2.0 * np.sum(np.log(diffusivity * dt_mean + sigma2))
    return result


def infer_smooth_D(
    cells,
    diffusivity_prior=None,
    jeffreys_prior=None,
    min_diffusivity=None,
    max_iter=None,
    epsilon=None,
    rgrad=None,
    verbose=False,
    **kwargs
):

    # initial values
    localization_error = cells.get_localization_error(kwargs, 0.03, True)
    (
        index,
        reverse_index,
        n,
        dt_mean,
        D_initial,
        min_diffusivity,
        D_bounds,
        _,
    ) = smooth_infer_init(
        cells,
        min_diffusivity=min_diffusivity,
        jeffreys_prior=jeffreys_prior,
        sigma2=localization_error,
    )

    # gradient options
    grad_kwargs = get_grad_kwargs(kwargs, epsilon=epsilon)

    # parametrize the optimization algorithm
    default_lBFGSb_options = dict(maxiter=1e3, maxfun=1e10, ftol=1e-6)
    # in L-BFGS-B the number of iterations is usually very low (~10-100) while the number of
    # function evaluations is much higher (~1e4-1e5);
    # with maxfun defined, an iteration can stop anytime and the optimization may terminate
    # with an error message
    if min_diffusivity is None:
        options = {}
    else:
        kwargs["bounds"] = D_bounds
        options = dict(default_lBFGSb_options)
    options.update(kwargs.pop("options", {}))
    if max_iter:
        options["maxiter"] = max_iter
    if verbose:
        options["disp"] = verbose
    if options:
        kwargs["options"] = options

    # posterior function
    if rgrad in ("delta", "delta0", "delta1"):
        fun = d_neg_posterior1
    else:
        if rgrad not in (None, "grad", "grad1", "gradn"):
            warn("unsupported rgrad: {}".format(rgrad), RuntimeWarning)
        fun = smooth_d_neg_posterior

    args = (
        cells,
        localization_error,
        diffusivity_prior,
        jeffreys_prior,
        dt_mean,
        min_diffusivity,
        index,
        reverse_index,
        grad_kwargs,
    )

    # run the optimization
    result = interruptible_minimize(fun, D_initial, args=args, **kwargs)
    if not (result.success or verbose):
        warn("{}".format(result.message), OptimizationWarning)

    # format the result
    D = result.x
    D = pd.DataFrame(D, index=index, columns=["diffusivity"])

    return D
