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

# Copyright © 2017-2018, 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 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


setup = {
    "name": ("unsmooth.dd", "unsmooth.ddrift", "degraded.dd", "degraded.ddrift"),
    "provides": ("dd", "ddrift"),
    "arguments": OrderedDict(
        (
            (
                "localization_error",
                (
                    "-e",
                    dict(
                        type=float,
                        help="localization precision (see also sigma; default is 0.03)",
                    ),
                ),
            ),
            (
                "jeffreys_prior",
                ("-j", dict(action="store_true", help="Jeffreys' prior")),
            ),
            (
                "min_diffusivity",
                dict(type=float, help="minimum diffusivity value allowed"),
            ),
        )
    ),
    "cell_sampling": "individual",
}


def dd_neg_posterior(x, dd, cell, sigma2, jeffreys_prior, dt_mean, min_diffusivity):
    dd.update(x)
    D, drift = dd["D"], dd["drift"]
    if D < min_diffusivity:
        warn(DiffusivityWarning(D, min_diffusivity))
    noise_dt = sigma2
    n = len(cell)  # number of translocations
    denominator = 4.0 * (D * cell.dt + noise_dt)  # 4*(D+Dnoise)*dt
    dr_minus_drift_dt = cell.dr - np.outer(cell.dt, drift)
    # non-directional squared displacement
    ndsd = np.sum(dr_minus_drift_dt * dr_minus_drift_dt, axis=1)
    neg_posterior = (
        n * log(pi) + np.sum(np.log(denominator)) + np.sum(ndsd / denominator)
    )
    if jeffreys_prior:
        neg_posterior += 2.0 * log(D * dt_mean + sigma2)
    return neg_posterior


def infer_DD(
    cells, localization_error=None, jeffreys_prior=False, min_diffusivity=None, **kwargs
):
    if isinstance(cells, Distributed):  # multiple cells
        localization_error = cells.get_localization_error(
            kwargs, 0.03, True, localization_error=localization_error
        )
        args = (localization_error, jeffreys_prior, min_diffusivity)
        index, inferred = [], []
        for i in cells:
            cell = cells[i]
            index.append(i)
            inferred.append(infer_DD(cell, *args, **kwargs))
        any_cell = cell
        inferred = pd.DataFrame(
            np.stack(inferred, axis=0),
            index=index,
            columns=["diffusivity"] + ["drift " + col for col in any_cell.space_cols],
        )
        return inferred
    else:  # single cell
        cell = cells
        # sanity checks
        if not bool(cell):
            raise ValueError("empty cells")
        if cell.dr.shape[1] == 0:
            raise ValueError("translocation array has no column")
        if cell.dt.shape[1:]:
            raise ValueError("time deltas are structured in multiple dimensions")
        # ensure that translocations are properly oriented in time
        if not np.all(0 < cell.dt):
            warn("translocation dts are non-positive", RuntimeWarning)
            cell.dr[cell.dt < 0] *= -1.0
            cell.dt[cell.dt < 0] *= -1.0
        #
        dt_mean = np.mean(cell.dt)
        D_initial = np.mean(cell.dr * cell.dr) / (2.0 * dt_mean)
        initial_drift = np.zeros(cell.dim, dtype=D_initial.dtype)
        dd = ChainArray("D", D_initial, "drift", initial_drift)
        if min_diffusivity is not False:
            if min_diffusivity is None:
                noise_dt = localization_error
                min_diffusivity = (1e-16 - noise_dt) / np.max(cell.dt)
            kwargs["bounds"] = [(min_diffusivity, None)] + [(None, None)] * cell.dim
        # cell.cache = None # no cache needed
        result = minimize(
            dd_neg_posterior,
            dd.combined,
            args=(
                dd,
                cell,
                localization_error,
                jeffreys_prior,
                dt_mean,
                min_diffusivity,
            ),
            **kwargs
        )
        # dd.update(result.x)
        # return (dd['D'], dd['drift'])
        return result.x  # needless to split dd.combined into D and drift
