# -*- 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 tramway.core import ChainArray
from .base import *
from .gradient import *
from .standard_d import interruptible_minimize
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
import time


setup = {
    "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")),
            ),
            (
                "potential_prior",
                ("-v", dict(type=float, help="prior on the potential")),
            ),
            (
                "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 (~100)")),
            (
                "compatibility",
                (
                    "-c",
                    "--inferencemap",
                    "--compatible",
                    dict(action="store_true", help="InferenceMAP compatible"),
                ),
            ),
            (
                "rgrad",
                dict(
                    help="alternative gradient for the regularization; can be 'delta0' or 'delta1'"
                ),
            ),
            ("export_centers", dict(action="store_true")),
            ("verbose", ()),
        )
    ),
    "cell_sampling": "connected",
}
setup_with_grad_arguments(setup)


class DV(ChainArray):
    """
    Convenience class for handling the D+V parameter vector and related priors.

    Objects of this class are not supposed to be serialized in .rwa files.
    """

    __slots__ = (
        "_diffusivity_prior",
        "_potential_prior",
        "minimum_diffusivity",
        "prior_include",
    )

    def __init__(
        self,
        diffusivity,
        potential,
        diffusivity_prior=None,
        potential_prior=None,
        minimum_diffusivity=None,
        positive_diffusivity=None,
        prior_include=None,
    ):
        # positive_diffusivity is for backward compatibility
        ChainArray.__init__(self, "D", diffusivity, "V", potential)
        self._diffusivity_prior = diffusivity_prior
        self._potential_prior = potential_prior
        self.minimum_diffusivity = minimum_diffusivity
        if minimum_diffusivity is None and positive_diffusivity is True:
            self.minimum_diffusivity = 0
        self.prior_include = prior_include

    @property
    def D(self):
        return self["D"]

    @property
    def V(self):
        return self["V"]

    @D.setter
    def D(self, diffusivity):
        self["D"] = diffusivity

    @V.setter
    def V(self, potential):
        self["V"] = potential

    def diffusivity_prior(self, j):
        if self._diffusivity_prior and (
            self.prior_include is None or self.prior_include[j]
        ):
            return self._diffusivity_prior
        else:
            return None

    def potential_prior(self, j):
        if self._potential_prior and (
            self.prior_include is None or self.prior_include[j]
        ):
            return self._potential_prior
        else:
            return None


def dv_neg_posterior(
    x,
    dv,
    cells,
    sigma2,
    jeffreys_prior,
    dt_mean,
    index,
    reverse_index,
    grad_kwargs,
    y0,
    verbose,
    posteriors,
):
    """
    Adapted from InferenceMAP's *dvPosterior* procedure modified:

    .. code-block:: c++

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


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

            const double D = DV[2*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-D*gradVx*dt)*(dx-D*gradVx*dt) + (dy-D*gradVy*dt)*(dy-D*gradVy*dt))/(4.0*(D+Dnoise)*dt);
            }

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

    with ``dx-D*gradVx*dt`` and ``dy-D*gradVy*dt`` modified as ``dx+D*gradVx*dt`` and ``dy+D*gradVy*dt`` respectively.

    See also :ref:`DV <inference_dv>`.
    """
    if verbose:
        t = time.time()

    # extract `D` and `V`
    dv.update(x)
    D = dv.D
    V = dv.V
    #

    if dv.minimum_diffusivity is not None:
        observed_min = np.min(D)
        if observed_min < dv.minimum_diffusivity and not np.isclose(
            observed_min, dv.minimum_diffusivity
        ):
            warn(DiffusivityWarning(observed_min, dv.minimum_diffusivity))
    noise_dt = sigma2

    # for all cell
    raw_posterior = priors = 0.0
    for j, i in enumerate(index):
        cell = cells[i]
        n = len(cell)  # number of translocations

        # spatial gradient of the local potential energy
        gradV = cells.grad(i, V, reverse_index, **grad_kwargs)
        # print('{}\t{}\t{}\t{}\t{}\t{}'.format(i+1,D[j], V[j], -gradV[0], -gradV[1], result))
        # print('{}\t{}\t{}'.format(i+1, *gradV))
        if gradV is None:
            continue

        # various posterior terms
        # print(cell.dt)
        D_dt = D[j] * cell.dt
        denominator = 4.0 * (D_dt + noise_dt)
        dr_minus_drift = cell.dr + np.outer(D_dt, gradV)
        # non-directional squared displacement
        ndsd = np.sum(dr_minus_drift * dr_minus_drift, axis=1)
        res = n * log(pi) + np.sum(np.log(denominator)) + np.sum(ndsd / denominator)
        if np.isnan(res):
            # print('isnan')
            continue
        raw_posterior += res

        # priors
        potential_prior = dv.potential_prior(j)
        if potential_prior:
            priors += potential_prior * cells.grad_sum(i, gradV * gradV, reverse_index)
        diffusivity_prior = dv.diffusivity_prior(j)
        if diffusivity_prior:
            # spatial gradient of the local diffusivity
            gradD = cells.grad(i, D, reverse_index, **grad_kwargs)
            if gradD is not None:
                # `grad_sum` memoizes and can be called several times at no extra cost
                priors += diffusivity_prior * cells.grad_sum(
                    i, gradD * gradD, reverse_index
                )
        # print('{}\t{}\t{}'.format(i+1, D[j], result))
    if jeffreys_prior:
        priors += 2.0 * np.sum(np.log(D * dt_mean + sigma2) - np.log(D))

    result = raw_posterior + priors
    posteriors.append([raw_posterior, result])

    if verbose:
        print(
            "objective: {}\t time: {}ms".format(
                result, int(round((time.time() - t) * 1e3))
            )
        )

    return result - y0


def dv_neg_posterior1(
    x,
    dv,
    cells,
    sigma2,
    jeffreys_prior,
    dt_mean,
    index,
    reverse_index,
    grad_kwargs,
    y0,
    verbose,
    posteriors,
):
    """
    Similar to :func:`dv_neg_posterior`.
    The smoothing priors feature an alternative spatial "gradient" implemented using
    :meth:`~tramway.inference.base.Distributed.local_variation` instead of
    :meth:`~tramway.inference.base.Distributed.grad`.
    """
    if verbose:
        t = time.time()

    # extract `D` and `V`
    dv.update(x)
    D = dv.D
    V = dv.V
    #

    if dv.minimum_diffusivity is not None:
        observed_min = np.min(D)
        if observed_min < dv.minimum_diffusivity and not np.isclose(
            observed_min, dv.minimum_diffusivity
        ):
            warn(DiffusivityWarning(observed_min, dv.minimum_diffusivity))
    noise_dt = sigma2

    # for all cell
    raw_posterior = priors = 0.0
    for j, i in enumerate(index):
        cell = cells[i]
        n = len(cell)  # number of translocations

        # spatial gradient of the local potential energy
        gradV = cells.grad(i, V, reverse_index, **grad_kwargs)
        # print('{}\t{}\t{}\t{}\t{}\t{}'.format(i+1,D[j], V[j], -gradV[0], -gradV[1], result))
        # print('{}\t{}\t{}'.format(i+1, *gradV))
        if gradV is None:
            continue

        # various posterior terms
        # print(cell.dt)
        D_dt = D[j] * cell.dt
        denominator = 4.0 * (D_dt + noise_dt)
        dr_minus_drift = cell.dr + np.outer(D_dt, gradV)
        # non-directional squared displacement
        ndsd = np.sum(dr_minus_drift * dr_minus_drift, axis=1)
        res = n * log(pi) + np.sum(np.log(denominator)) + np.sum(ndsd / denominator)
        if np.isnan(res):
            # print('isnan')
            continue
        raw_posterior += res

        # priors
        potential_prior = dv.potential_prior(j)
        if potential_prior:
            # spatial variation of the local potential
            deltaV = cells.local_variation(i, V, reverse_index, **grad_kwargs)
            if deltaV is not None:
                priors += potential_prior * cells.grad_sum(
                    i, deltaV * deltaV, reverse_index
                )
        diffusivity_prior = dv.diffusivity_prior(j)
        if diffusivity_prior:
            # spatial variation of the local diffusivity
            deltaD = cells.local_variation(i, D, reverse_index, **grad_kwargs)
            if deltaD is not None:
                # `grad_sum` memoizes and can be called several times at no extra cost
                priors += diffusivity_prior * cells.grad_sum(
                    i, deltaD * deltaD, reverse_index
                )
        # print('{}\t{}\t{}'.format(i+1, D[j], result))
    if jeffreys_prior:
        priors += 2.0 * np.sum(np.log(D * dt_mean + sigma2) - np.log(D))

    result = raw_posterior + priors
    if posteriors is not None:
        posteriors.append([raw_posterior, result])

    if verbose:
        print(
            "objective: {}\t time: {}ms".format(
                result, int(round((time.time() - t) * 1e3))
            )
        )

    return result - y0


def inferDV(
    cells,
    diffusivity_prior=None,
    potential_prior=None,
    jeffreys_prior=False,
    min_diffusivity=None,
    max_iter=None,
    epsilon=None,
    export_centers=False,
    verbose=True,
    compatibility=False,
    D0=None,
    V0=None,
    rgrad=None,
    **kwargs
):

    localization_error = cells.get_localization_error(kwargs, 0.03, True)

    # initial values
    (
        index,
        reverse_index,
        n,
        dt_mean,
        D_initial,
        min_diffusivity,
        D_bounds,
        border,
    ) = smooth_infer_init(
        cells,
        min_diffusivity=min_diffusivity,
        jeffreys_prior=jeffreys_prior,
        sigma2=localization_error,
    )
    # V initial values
    if V0 is None:
        try:
            if compatibility:
                raise Exception  # skip to the except block
            volume = [cells[i].volume for i in index]
        except:
            V_initial = -np.log(n / np.max(n))
        else:
            density = n / np.array([np.inf if v is None else v for v in volume])
            density[density == 0] = np.min(density[0 < density])
            V_initial = np.log(np.max(density)) - np.log(density)
    else:
        if np.isscalar(V0):
            V_initial = np.full(D_initial.size, V0)
        elif V0.size == D_initial.size:
            V_initial = V0
        else:
            raise ValueError("wrong size for V0")
    if D0 is not None:
        if np.isscalar(D0):
            D_initial[...] = D0
        elif D0.size == D_initial.size:
            D_initial = D0
        else:
            raise ValueError("wrong size for D0")

    dv = DV(D_initial, V_initial, diffusivity_prior, potential_prior, min_diffusivity)
    posteriors = []

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

    # parametrize the optimization algorithm
    default_lBFGSb_options = dict(
        maxiter=1e3, maxfun=1e10, maxcor=dv.combined.size, ftol=1e-8
    )
    # 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:
        bounds = None
        options = {}
    else:
        V_bounds = [(None, None)] * V_initial.size
        bounds = D_bounds + V_bounds
        options = dict(default_lBFGSb_options)
    _kwargs = {}
    for k in ("jac",):
        try:
            _kwargs[k] = kwargs.pop(k)
        except KeyError:
            pass
    options.update(kwargs.pop("options", {}))
    options.update(**kwargs)  # for backward compatibility
    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 = dv_neg_posterior1
    else:
        if rgrad not in (None, "grad", "grad1", "gradn"):
            warn("unsupported rgrad: {}".format(rgrad), RuntimeWarning)
        fun = dv_neg_posterior

    # posterior function input arguments
    args = (
        dv,
        cells,
        localization_error,
        jeffreys_prior,
        dt_mean,
        index,
        reverse_index,
        grad_kwargs,
    )

    # get the initial posterior value so that it is subtracted from the further evaluations
    y0 = fun(dv.combined, *(args + (0.0, False, [])))
    if verbose:
        print("At X0\tactual posterior= {}\n".format(y0))
    # y0 = 0.
    args = args + (y0, 1 < int(verbose), posteriors)

    # run the optimization routine
    result = interruptible_minimize(
        fun, dv.combined, args=args, bounds=bounds, **_kwargs
    )
    if not (result.success or verbose):
        warn("{}".format(result.message), OptimizationWarning)

    y = np.array(result.x)

    # collect the result
    dv.update(y)
    D, V = dv.D, dv.V
    if np.any(V < 0):
        V -= np.min(V)
    DVF = pd.DataFrame(
        np.stack((D, V), axis=1), index=index, columns=["diffusivity", "potential"]
    )

    # derivate the forces
    index_, F = [], []
    for i in index:
        gradV = cells.grad(i, V, reverse_index, **grad_kwargs)
        if gradV is not None:
            index_.append(i)
            F.append(-gradV)
    if F:
        F = pd.DataFrame(
            np.stack(F, axis=0),
            index=index_,
            columns=["force " + col for col in cells.space_cols],
        )
    else:
        warn("not any cell is suitable for evaluating the local force", RuntimeWarning)
        F = pd.DataFrame(
            np.zeros((0, len(cells.space_cols)), dtype=V.dtype),
            columns=["force " + col for col in cells.space_cols],
        )
    DVF = DVF.join(F)

    # add extra information if required
    if export_centers:
        xy = np.vstack([cells[i].center for i in index])
        DVF = DVF.join(pd.DataFrame(xy, index=index, columns=cells.space_cols))
        # DVF.to_csv('results.csv', sep='\t')

    # format the posteriors
    posteriors = pd.DataFrame(np.array(posteriors), columns=["fit", "total"])

    return DVF, posteriors
