"""
pysteps.nowcasts.linda
======================

This module implements the Lagrangian INtegro-Difference equation model with
Autoregression (LINDA). The model combines extrapolation, S-PROG, STEPS, ANVIL,
integro-difference equation (IDE) and cell tracking methods. It can produce
both deterministic and probabilistic nowcasts. LINDA is specifically designed
for nowcasting intense localized rainfall. For this purpose, it is expected to
give better forecast skill than S-PROG or STEPS.

The model consists of the following components:

1. feature detection to identify rain cells
2. advection-based extrapolation
3. autoregressive integrated ARI(p,1) process for growth and decay of rainfall
4. convolution to account for loss of predictability
5. stochastic perturbations to simulate forecast errors

LINDA utilizes a sparse feature-based representation of the input rain rate
fields. This allows localization to cells containing intense rainfall. Building
on extrapolation nowcast, the temporal evolution of rainfall is modeled in the
Lagrangian coordinates. Using the ARI process is adapted from ANVIL
:cite:`PCLH2020`, and the convolution is adapted from the integro-difference
equation (IDE) models proposed in :cite:`FW2005` and :cite:`XWF2005`. The
combination of these two approaches essentially replaces the cascade-based
autoregressive process used in S-PROG and STEPS. Using the convolution gives
several advantages such as the ability to handle anisotropic structure, domain
boundaries and missing data. Based on the marginal distribution and covariance
structure of forecast errors, localized perturbations are generated by adapting
the short-space Fourier transform (SSFT) methodology developed in
:cite:`NBSG2017`.

.. autosummary::
    :toctree: ../generated/

    forecast
"""

import time
import warnings

try:
    import dask

    DASK_IMPORTED = True
except ImportError:
    DASK_IMPORTED = False
import numpy as np
from scipy.integrate import nquad
from scipy.interpolate import interp1d
from scipy import optimize as opt
from scipy.signal import convolve
from scipy import stats

from pysteps import extrapolation, feature, noise
from pysteps.decorators import deprecate_args
from pysteps.nowcasts.utils import nowcast_main_loop


@deprecate_args(
    {
        "precip_fields": "precip",
        "advection_field": "velocity",
        "num_ens_members": "n_ens_members",
    },
    "1.8.0",
)
def forecast(
    precip,
    velocity,
    timesteps,
    feature_method="blob",
    max_num_features=25,
    feature_kwargs=None,
    ari_order=1,
    kernel_type="anisotropic",
    localization_window_radius=None,
    errdist_window_radius=None,
    acf_window_radius=None,
    extrap_method="semilagrangian",
    extrap_kwargs=None,
    add_perturbations=True,
    pert_thrs=(0.5, 1.0),
    n_ens_members=10,
    vel_pert_method="bps",
    vel_pert_kwargs=None,
    kmperpixel=None,
    timestep=None,
    seed=None,
    num_workers=1,
    use_multiprocessing=False,
    measure_time=False,
    callback=None,
    return_output=True,
):
    """
    Generate a deterministic or ensemble nowcast by using the Lagrangian
    INtegro-Difference equation model with Autoregression (LINDA) model.

    Parameters
    ----------
    precip: array_like
        Array of shape (ari_order + 2, m, n) containing the input rain rate
        or reflectivity fields (in linear scale) ordered by timestamp from
        oldest to newest. The time steps between the inputs are assumed to be
        regular.
    velocity: array_like
        Array of shape (2, m, n) containing the x- and y-components of the
        advection field. The velocities are assumed to represent one time step
        between the inputs.
    timesteps: int
        Number of time steps to forecast.
    feature_method: {'blob', 'domain' 'shitomasi'}
        Feature detection method:

        +-------------------+-----------------------------------------------------+
        |    Method name    |                  Description                        |
        +===================+=====================================================+
        |  blob             | Laplacian of Gaussian (LoG) blob detector           |
        |                   | implemented in scikit-image                         |
        +-------------------+-----------------------------------------------------+
        |  domain           | no feature detection, the model is applied over the |
        |                   | whole domain without localization                   |
        +-------------------+-----------------------------------------------------+
        |  shitomasi        | Shi-Tomasi corner detector implemented in OpenCV    |
        +-------------------+-----------------------------------------------------+

        Default: 'blob'
    max_num_features: int, optional
        Maximum number of features to use. It is recommended to set this between
        20 and 50, which gives a good tradeoff between localization and
        computation time. Default: 25
    feature_kwargs: dict, optional
        Keyword arguments that are passed as ``**kwargs`` for the feature detector.
        See :py:mod:`pysteps.feature.blob` and :py:mod:`pysteps.feature.shitomasi`.
    ari_order: {1, 2}, optional
        The order of the ARI(p, 1) model. Default: 1
    kernel_type: {"anisotropic", "isotropic"}, optional
        The type of the kernel. Default: 'anisotropic'
    localization_window_radius: float, optional
        The standard deviation of the Gaussian localization window.
        Default: 0.2 * min(m, n)
    errdist_window_radius: float, optional
        The standard deviation of the Gaussian window for estimating the
        forecast error distribution. Default: 0.15 * min(m, n)
    acf_window_radius: float, optional
        The standard deviation of the Gaussian window for estimating the
        forecast error ACF. Default: 0.25 * min(m, n)
    extrap_method: str, optional
        The extrapolation method to use. See the documentation of
        :py:mod:`pysteps.extrapolation.interface`. Default: 'semilagrangian'
    extrap_kwargs: dict, optional
        Optional dictionary containing keyword arguments for the extrapolation
        method. See :py:mod:`pysteps.extrapolation.interface`.
    add_perturbations: bool
        Set to False to disable perturbations and generate a single
        deterministic nowcast. Default: True
    pert_thrs: float
        Two-element tuple containing the threshold values for estimating the
        perturbation parameters (mm/h). Default: (0.5, 1.0)
    n_ens_members: int, optional
        The number of ensemble members to generate. Default: 10
    vel_pert_method: {'bps', None}, optional
        Name of the generator to use for perturbing the advection field. See
        :py:mod:`pysteps.noise.interface`. Default: 'bps'
    vel_pert_kwargs: dict, optional
        Optional dictionary containing keyword arguments 'p_par' and 'p_perp'
        for the initializer of the velocity perturbator. The choice of the
        optimal parameters depends on the domain and the used optical flow
        method. For the default values and parameters optimized for different
        domains, see :py:func:`pysteps.nowcasts.steps.forecast`.
    kmperpixel: float, optional
        Spatial resolution of the input data (kilometers/pixel). Required if
        vel_pert_method is not None.
    timestep: float, optional
        Time step of the motion vectors (minutes). Required if vel_pert_method
        is not None.
    seed: int, optional
        Optional seed for the random generators.
    num_workers: int, optional
        The number of workers to use for parallel computations. Applicable if
        dask is installed. Default: 1
    use_multiprocessing: bool, optional
        Set to True to improve the performance of certain parallelized parts of
        the code. If set to True, the main script calling linda.forecast must
        be enclosed within the 'if __name__ == "__main__":' block.
        Default: False
    measure_time: bool, optional
        If set to True, measure, print and return the computation time.
        Default: False
    callback: function, optional
        Optional function that is called after computation of each time step of
        the nowcast. The function takes one argument: a three-dimensional array
        of shape (n_ens_members,h,w), where h and w are the height and width
        of the input precipitation fields, respectively. This can be used, for
        instance, writing the outputs into files. Default: None
    return_output: bool, optional
        Set to False to disable returning the outputs as numpy arrays. This can
        save memory if the intermediate results are written to output files
        using the callback function. Default: True

    Returns
    -------
    out: numpy.ndarray
        A four-dimensional array of shape (n_ens_members, timesteps, m, n)
        containing a time series of forecast precipitation fields for each
        ensemble member. If add_perturbations is False, the first dimension is
        dropped. The time series starts from t0 + timestep, where timestep is
        taken from the input fields. If measure_time is True, the return value
        is a three-element tuple containing the nowcast array, the initialization
        time of the nowcast generator and the time used in the main loop
        (seconds). If return_output is set to False, a single None value is
        returned instead.

    Notes
    -----
    It is recommended to choose the feature detector parameters so that the
    number of features is around 20-40. This gives a good tradeoff between
    localization and computation time.

    It is highly recommented to set num_workers>1 to reduce computation time.
    In this case, it is advisable to disable OpenMP by setting the environment
    variable OMP_NUM_THREADS to 1. This avoids slowdown caused by too many
    simultaneous threads.
    """
    _check_inputs(precip, velocity, timesteps, ari_order)

    if feature_kwargs is None:
        feature_kwargs = dict()
    if extrap_kwargs is None:
        extrap_kwargs = dict()
    else:
        extrap_kwargs = extrap_kwargs.copy()

    if localization_window_radius is None:
        localization_window_radius = 0.2 * np.min(precip.shape[1:])
    if add_perturbations:
        if errdist_window_radius is None:
            errdist_window_radius = 0.15 * min(precip.shape[1], precip.shape[2])
        if acf_window_radius is None:
            acf_window_radius = 0.25 * min(precip.shape[1], precip.shape[2])
        if vel_pert_method is not None:
            if kmperpixel is None:
                raise ValueError("vel_pert_method is set but kmperpixel is None")
            if timestep is None:
                raise ValueError("vel_pert_method is set but timestep is None")
        if vel_pert_kwargs is None:
            vel_pert_kwargs = dict()

    print("Computing LINDA nowcast")
    print("-----------------------")
    print("")

    print("Inputs")
    print("------")
    print(f"dimensions:           {precip.shape[1]}x{precip.shape[2]}")
    print(f"number of time steps: {precip.shape[0]}")
    print("")

    print("Methods")
    print("-------")
    nowcast_type = "ensemble" if add_perturbations else "deterministic"
    print(f"nowcast type:         {nowcast_type}")
    print(f"feature detector:     {feature_method}")
    print(f"extrapolator:         {extrap_method}")
    print(f"kernel type:          {kernel_type}")
    if add_perturbations and vel_pert_method is not None:
        print(f"velocity perturbator: {vel_pert_method}")
    print("")

    print("Parameters")
    print("----------")
    print(f"number of time steps:       {timesteps}")
    print(f"ARI model order:            {ari_order}")
    print(f"localization window radius: {localization_window_radius}")
    if add_perturbations:
        print(f"error dist. window radius:  {errdist_window_radius}")
        print(f"error ACF window radius:    {acf_window_radius}")
        print(f"ensemble size:              {n_ens_members}")
        print(f"parallel workers:           {num_workers}")
        print(f"seed:                       {seed}")
        if vel_pert_method == "bps":
            vp_par = vel_pert_kwargs.get(
                "p_par", noise.motion.get_default_params_bps_par()
            )
            vp_perp = vel_pert_kwargs.get(
                "p_perp", noise.motion.get_default_params_bps_perp()
            )
            print(
                f"velocity perturbations, parallel:      {vp_par[0]:.2f}, {vp_par[1]:.2f}, {vp_par[2]:.2f}"
            )
            print(
                f"velocity perturbations, perpendicular: {vp_perp[0]:.2f}, {vp_perp[1]:.2f}, {vp_perp[2]:.2f}"
            )
            vel_pert_kwargs = vel_pert_kwargs.copy()
            vel_pert_kwargs["vp_par"] = vp_par
            vel_pert_kwargs["vp_perp"] = vp_perp

    extrap_kwargs["allow_nonfinite_values"] = (
        True if np.any(~np.isfinite(precip)) else False
    )

    forecast_gen = _linda_deterministic_init(
        precip,
        velocity,
        feature_method,
        max_num_features,
        feature_kwargs,
        ari_order,
        kernel_type,
        localization_window_radius,
        extrap_method,
        extrap_kwargs,
        add_perturbations,
        num_workers,
        measure_time,
    )
    if measure_time:
        forecast_gen, precip_lagr_diff, init_time = forecast_gen
    else:
        forecast_gen, precip_lagr_diff = forecast_gen

    if add_perturbations:
        pert_gen = _linda_perturbation_init(
            precip,
            precip_lagr_diff,
            velocity,
            forecast_gen,
            pert_thrs,
            localization_window_radius,
            errdist_window_radius,
            acf_window_radius,
            vel_pert_method,
            vel_pert_kwargs,
            kmperpixel,
            timestep,
            num_workers,
            use_multiprocessing,
            measure_time,
        )
        if measure_time:
            precip_pert_gen, velocity_pert_gen, pert_init_time = pert_gen
            init_time += pert_init_time
        else:
            precip_pert_gen, velocity_pert_gen = pert_gen
    else:
        precip_pert_gen = None
        velocity_pert_gen = None

    precip_forecast = _linda_forecast(
        precip,
        precip_lagr_diff[1:],
        timesteps,
        forecast_gen,
        precip_pert_gen,
        velocity_pert_gen,
        n_ens_members,
        seed,
        measure_time,
        True,
        return_output,
        callback,
    )

    if return_output:
        if measure_time:
            return precip_forecast[0], init_time, precip_forecast[1]
        else:
            return precip_forecast
    else:
        return None


def _check_inputs(precip, velocity, timesteps, ari_order):
    if ari_order not in [1, 2]:
        raise ValueError(f"ari_order {ari_order} given, 1 or 2 required")
    if len(precip.shape) != 3:
        raise ValueError("precip must be a three-dimensional array")
    if precip.shape[0] < ari_order + 2:
        raise ValueError("precip.shape[0] < ari_order+2")
    if len(velocity.shape) != 3:
        raise ValueError("velocity must be a three-dimensional array")
    if precip.shape[1:3] != velocity.shape[1:3]:
        raise ValueError(
            f"dimension mismatch between precip and velocity: precip.shape={precip.shape}, velocity.shape={velocity.shape}"
        )
    if not isinstance(timesteps, int):
        raise ValueError("timesteps is not an integer")


def _composite_convolution(field, kernels, weights):
    """
    Compute a localized convolution by applying a set of kernels with the
    given spatial weights. The weights are assumed to be normalized.
    """
    n = len(kernels)
    field_c = 0.0

    for i in range(n):
        field_c += weights[i] * _masked_convolution(field, kernels[i])

    return field_c


def _compute_ellipse_bbox(phi, sigma1, sigma2, cutoff):
    """Compute the bounding box of an ellipse."""
    r1 = cutoff * sigma1
    r2 = cutoff * sigma2
    phi_r = phi / 180.0 * np.pi

    if np.abs(phi_r - np.pi / 2) > 1e-6 and np.abs(phi_r - 3 * np.pi / 2) > 1e-6:
        alpha = np.arctan(-r2 * np.sin(phi_r) / (r1 * np.cos(phi_r)))
        w = r1 * np.cos(alpha) * np.cos(phi_r) - r2 * np.sin(alpha) * np.sin(phi_r)

        alpha = np.arctan(r2 * np.cos(phi_r) / (r1 * np.sin(phi_r)))
        h = r1 * np.cos(alpha) * np.sin(phi_r) + r2 * np.sin(alpha) * np.cos(phi_r)
    else:
        w = sigma2 * cutoff
        h = sigma1 * cutoff

    return -abs(h), -abs(w), abs(h), abs(w)


def _compute_inverse_acf_mapping(target_dist, target_dist_params, n_intervals=10):
    """Compute the inverse ACF mapping between two distributions."""
    phi = (
        lambda x1, x2, rho: 1.0
        / (2 * np.pi * np.sqrt(1 - rho**2))
        * np.exp(-(x1**2 + x2**2 - 2 * rho * x1 * x2) / (2 * (1 - rho**2)))
    )

    rho_1 = np.linspace(-0.9, 0.9, n_intervals)
    rho_2 = np.empty(len(rho_1))

    mu = target_dist.mean(*target_dist_params)
    sigma = target_dist.std(*target_dist_params)

    cdf_trans = lambda x: target_dist.ppf(stats.norm.cdf(x), *target_dist_params)
    int_range = (-6, 6)

    for i, rho_1_ in enumerate(rho_1):
        f = (
            lambda x1, x2: (cdf_trans(x1) - mu)
            * (cdf_trans(x2) - mu)
            * phi(x1, x2, rho_1_)
        )
        opts = {"epsabs": 1e-8, "epsrel": 1e-8, "limit": 1}
        rho_2[i] = nquad(f, (int_range, int_range), opts=opts)[0] / (sigma * sigma)

    return interp1d(rho_2, rho_1, fill_value="extrapolate")


def _compute_kernel_anisotropic(params, cutoff=6.0):
    """Compute anisotropic Gaussian convolution kernel."""
    phi, sigma1, sigma2 = params

    phi_r = phi / 180.0 * np.pi
    rot_inv = np.array(
        [[np.cos(phi_r), np.sin(phi_r)], [-np.sin(phi_r), np.cos(phi_r)]]
    )

    bb_y1, bb_x1, bb_y2, bb_x2 = _compute_ellipse_bbox(phi, sigma1, sigma2, cutoff)

    x = np.arange(int(bb_x1), int(bb_x2) + 1).astype(float)
    if len(x) % 2 == 0:
        x = np.arange(int(bb_x1) - 1, int(bb_x2) + 1).astype(float)
    y = np.arange(int(bb_y1), int(bb_y2) + 1).astype(float)
    if len(y) % 2 == 0:
        y = np.arange(int(bb_y1) - 1, int(bb_y2) + 1).astype(float)

    x_grid, y_grid = np.meshgrid(x, y)
    xy_grid = np.vstack([x_grid.flatten(), y_grid.flatten()])
    xy_grid = np.dot(rot_inv, xy_grid)

    x2 = xy_grid[0, :] * xy_grid[0, :]
    y2 = xy_grid[1, :] * xy_grid[1, :]
    result = np.exp(-(x2 / sigma1**2 + y2 / sigma2**2))

    return np.reshape(result / np.sum(result), x_grid.shape)


def _compute_kernel_isotropic(sigma, cutoff=6.0):
    """Compute isotropic Gaussian convolution kernel."""
    bb_y1, bb_x1, bb_y2, bb_x2 = (
        -sigma * cutoff,
        -sigma * cutoff,
        sigma * cutoff,
        sigma * cutoff,
    )

    x = np.arange(int(bb_x1), int(bb_x2) + 1).astype(float)
    if len(x) % 2 == 0:
        x = np.arange(int(bb_x1) - 1, int(bb_x2) + 1).astype(float)
    y = np.arange(int(bb_y1), int(bb_y2) + 1).astype(float)
    if len(y) % 2 == 0:
        y = np.arange(int(bb_y1) - 1, int(bb_y2) + 1).astype(float)

    x_grid, y_grid = np.meshgrid(x / sigma, y / sigma)

    r2 = x_grid * x_grid + y_grid * y_grid
    result = np.exp(-0.5 * r2)

    return result / np.sum(result)


def _compute_parametric_acf(params, m, n):
    """Compute parametric ACF."""
    c, phi, sigma1, sigma2 = params

    phi_r = phi / 180.0 * np.pi
    rot_inv = np.array(
        [[np.cos(phi_r), np.sin(phi_r)], [-np.sin(phi_r), np.cos(phi_r)]]
    )

    if n % 2 == 0:
        n_max = int(n / 2)
    else:
        n_max = int(n / 2) + 1
    x = np.fft.ifftshift(np.arange(-int(n / 2), n_max))
    if m % 2 == 0:
        m_max = int(m / 2)
    else:
        m_max = int(m / 2) + 1
    y = np.fft.ifftshift(np.arange(-int(m / 2), m_max))

    grid_x, grid_y = np.meshgrid(x, y)
    grid_xy = np.vstack([grid_x.flatten(), grid_y.flatten()])
    grid_xy = np.dot(rot_inv, grid_xy)

    grid_xy[0, :] = grid_xy[0, :] / sigma1
    grid_xy[1, :] = grid_xy[1, :] / sigma2

    r2 = np.reshape(
        grid_xy[0, :] * grid_xy[0, :] + grid_xy[1, :] * grid_xy[1, :], grid_x.shape
    )
    result = np.exp(-np.sqrt(r2))

    return c * result


def _compute_sample_acf(field):
    """Compute sample ACF from FFT."""
    # TODO: let user choose the FFT method
    field_fft = np.fft.rfft2((field - np.mean(field)) / np.std(field))
    fft_abs = np.abs(field_fft * np.conj(field_fft))

    return np.fft.irfft2(fft_abs, s=field.shape) / (field.shape[0] * field.shape[1])


def _compute_window_weights(coords, grid_height, grid_width, window_radius):
    """Compute interpolation weights."""
    coords = coords.astype(float).copy()
    num_features = coords.shape[0]

    coords[:, 0] /= grid_height
    coords[:, 1] /= grid_width

    window_radius_1 = window_radius / grid_height
    window_radius_2 = window_radius / grid_width

    grid_x = (np.arange(grid_width) + 0.5) / grid_width
    grid_y = (np.arange(grid_height) + 0.5) / grid_height

    grid_x, grid_y = np.meshgrid(grid_x, grid_y)

    w = np.empty((num_features, grid_x.shape[0], grid_x.shape[1]))

    if coords.shape[0] > 1:
        for i, c in enumerate(coords):
            dy = c[0] - grid_y
            dx = c[1] - grid_x

            w[i, :] = np.exp(
                -dy * dy / (2 * window_radius_1**2)
                - dx * dx / (2 * window_radius_2**2)
            )
    else:
        w[0, :] = np.ones((grid_height, grid_width))

    return w


def _estimate_ar1_params(
    field_src, field_dst, estim_weights, interp_weights, num_workers=1
):
    """Constrained optimization of AR(1) parameters."""

    def objf(p, *args):
        i = args[0]
        field_ar = p * field_src
        return np.nansum(estim_weights[i] * (field_dst - field_ar) ** 2.0)

    bounds = (-0.98, 0.98)

    def worker(i):
        return opt.minimize_scalar(objf, method="bounded", bounds=bounds, args=(i,)).x

    if DASK_IMPORTED and num_workers > 1:
        res = []
        for i in range(len(estim_weights)):
            res.append(dask.delayed(worker)(i))

        psi = dask.compute(*res, num_workers=num_workers, scheduler="threads")
    else:
        psi = []
        for i in range(len(estim_weights)):
            psi.append(worker(i))

    return [np.sum([psi_ * interp_weights[i] for i, psi_ in enumerate(psi)], axis=0)]


def _estimate_ar2_params(
    field_src, field_dst, estim_weights, interp_weights, num_workers=1
):
    """Constrained optimization of AR(2) parameters."""

    def objf(p, *args):
        i = args[0]
        field_ar = p[0] * field_src[1] + p[1] * field_src[0]
        return np.nansum(estim_weights[i] * (field_dst - field_ar) ** 2.0)

    bounds = [(-1.98, 1.98), (-0.98, 0.98)]
    constraints = [
        opt.LinearConstraint(
            np.array([(1, 1), (-1, 1)]),
            (-np.inf, -np.inf),
            (0.98, 0.98),
            keep_feasible=True,
        )
    ]

    def worker(i):
        return opt.minimize(
            objf,
            (0.8, 0.0),
            method="trust-constr",
            bounds=bounds,
            constraints=constraints,
            args=(i,),
        ).x

    if DASK_IMPORTED and num_workers > 1:
        res = []
        for i in range(len(estim_weights)):
            res.append(dask.delayed(worker)(i))

        psi = dask.compute(*res, num_workers=num_workers, scheduler="threads")
    else:
        psi = []
        for i in range(len(estim_weights)):
            psi.append(worker(i))

    psi_out = []
    for i in range(2):
        psi_out.append(
            np.sum([psi[j][i] * interp_weights[j] for j in range(len(psi))], axis=0)
        )

    return psi_out


def _estimate_convol_params(
    field_src,
    field_dst,
    weights,
    mask,
    kernel_type="anisotropic",
    kernel_params=None,
    num_workers=1,
):
    """Estimation of convolution kernel."""
    if kernel_params is None:
        kernel_params = {}
    masks = []
    for weight in weights:
        masks.append(np.logical_and(mask, weight > 1e-3))

    def objf_aniso(p, *args):
        i = args[0]
        p = _get_anisotropic_kernel_params(p)
        kernel = _compute_kernel_anisotropic(p, **kernel_params)

        field_src_c = _masked_convolution(field_src, kernel)
        fval = np.sqrt(weights[i][masks[i]]) * (
            field_dst[masks[i]] - field_src_c[masks[i]]
        )

        return fval

    def objf_iso(p, *args):
        i = args[0]
        kernel = _compute_kernel_isotropic(p, **kernel_params)

        field_src_c = _masked_convolution(field_src, kernel)
        fval = np.sum(
            weights[i][masks[i]] * (field_dst[masks[i]] - field_src_c[masks[i]]) ** 2
        )

        return fval

    def worker(i):
        if kernel_type == "anisotropic":
            bounds = ((-np.inf, 0.1, 0.2), (np.inf, 10.0, 5.0))
            p_opt = opt.least_squares(
                objf_aniso,
                np.array((0.0, 1.0, 1.0)),
                bounds=bounds,
                method="trf",
                ftol=1e-6,
                xtol=1e-4,
                gtol=1e-6,
                args=(i,),
            )
            p_opt = _get_anisotropic_kernel_params(p_opt.x)

            return _compute_kernel_anisotropic(p_opt, **kernel_params)
        else:
            p_opt = opt.minimize_scalar(
                objf_iso, bounds=[0.01, 10.0], method="bounded", args=(i,)
            )
            p_opt = p_opt.x

            return _compute_kernel_isotropic(p_opt, **kernel_params)

    if DASK_IMPORTED and num_workers > 1:
        res = []
        for i in range(len(weights)):
            res.append(dask.delayed(worker)(i))
        kernels = dask.compute(*res, num_workers=num_workers, scheduler="threads")
    else:
        kernels = []
        for i in range(len(weights)):
            kernels.append(worker(i))

    return kernels


def _estimate_perturbation_params(
    forecast_err,
    forecast_gen,
    errdist_window_radius,
    acf_window_radius,
    interp_window_radius,
    measure_time,
    num_workers,
    use_multiprocessing,
):
    """
    Estimate perturbation generator parameters from forecast errors."""
    pert_gen = {}
    pert_gen["m"] = forecast_err.shape[0]
    pert_gen["n"] = forecast_err.shape[1]

    feature_coords = forecast_gen["feature_coords"]

    print("Estimating perturbation parameters... ", end="", flush=True)

    if measure_time:
        starttime = time.time()

    mask_finite = np.isfinite(forecast_err)

    forecast_err = forecast_err.copy()
    forecast_err[~mask_finite] = 1.0

    weights_dist = _compute_window_weights(
        feature_coords,
        forecast_err.shape[0],
        forecast_err.shape[1],
        errdist_window_radius,
    )

    acf_winfunc = _window_tukey if feature_coords.shape[0] > 1 else _window_uniform

    def worker(i):
        weights_acf = acf_winfunc(
            forecast_err.shape[0],
            forecast_err.shape[1],
            feature_coords[i, 0],
            feature_coords[i, 1],
            acf_window_radius,
            acf_window_radius,
        )

        mask = np.logical_and(mask_finite, weights_dist[i] > 0.1)
        if np.sum(mask) > 10 and np.sum(np.abs(forecast_err[mask] - 1.0) >= 1e-3) > 10:
            distpar = _fit_dist(forecast_err, stats.lognorm, weights_dist[i], mask)
            inv_acf_mapping = _compute_inverse_acf_mapping(stats.lognorm, distpar)
            mask_acf = weights_acf > 1e-4
            std = _weighted_std(forecast_err[mask_acf], weights_dist[i][mask_acf])
            if np.isfinite(std):
                acf = inv_acf_mapping(
                    _compute_sample_acf(weights_acf * (forecast_err - 1.0) / std)
                )
                acf = _fit_acf(acf)
            else:
                distpar = None
                std = None
                acf = None
        else:
            distpar = None
            std = None
            acf = None

        return distpar, std, np.sqrt(np.abs(np.fft.rfft2(acf)))

    dist_params = []
    stds = []
    acf_fft_ampl = []

    if DASK_IMPORTED and num_workers > 1:
        res = []
        for i in range(feature_coords.shape[0]):
            res.append(dask.delayed(worker)(i))
        scheduler = "threads" if not use_multiprocessing else "multiprocessing"
        res = dask.compute(*res, num_workers=num_workers, scheduler=scheduler)
        for r in res:
            dist_params.append(r[0])
            stds.append(r[1])
            acf_fft_ampl.append(r[2])
    else:
        for i in range(feature_coords.shape[0]):
            r = worker(i)
            dist_params.append(r[0])
            stds.append(r[1])
            acf_fft_ampl.append(r[2])

    pert_gen["dist_param"] = dist_params
    pert_gen["std"] = stds
    pert_gen["acf_fft_ampl"] = acf_fft_ampl

    weights = _compute_window_weights(
        feature_coords,
        forecast_err.shape[0],
        forecast_err.shape[1],
        interp_window_radius,
    )
    pert_gen["weights"] = weights / np.sum(weights, axis=0)

    if measure_time:
        print(f"{time.time() - starttime:.2f} seconds.")
    else:
        print("done.")

    return pert_gen


def _fit_acf(acf):
    """
    Fit a parametric ACF to the given sample estimate."""

    def objf(p, *args):
        p = _get_acf_params(p)
        fitted_acf = _compute_parametric_acf(p, acf.shape[0], acf.shape[1])

        return (acf - fitted_acf).flatten()

    bounds = ((0.01, -np.inf, 0.1, 0.2), (10.0, np.inf, 10.0, 5.0))
    p_opt = opt.least_squares(
        objf,
        np.array((1.0, 0.0, 1.0, 1.0)),
        bounds=bounds,
        method="trf",
        ftol=1e-6,
        xtol=1e-4,
        gtol=1e-6,
    )

    return _compute_parametric_acf(_get_acf_params(p_opt.x), acf.shape[0], acf.shape[1])


def _fit_dist(err, dist, wf, mask):
    """
    Fit a lognormal distribution by maximizing the log-likelihood function
    with the constraint that the mean value is one."""
    func = lambda p: -np.sum(np.log(stats.lognorm.pdf(err[mask], p, -0.5 * p**2)))
    p_opt = opt.minimize_scalar(func, bounds=(1e-3, 20.0), method="Bounded")

    return (p_opt.x, -0.5 * p_opt.x**2)


# TODO: restrict the perturbation generation inside the radar mask
def _generate_perturbations(pert_gen, num_workers, seed):
    """Generate perturbations based on the estimated forecast error statistics."""
    m, n = pert_gen["m"], pert_gen["n"]
    dist_param = pert_gen["dist_param"]
    std = pert_gen["std"]
    acf_fft_ampl = pert_gen["acf_fft_ampl"]
    weights = pert_gen["weights"]

    perturb = stats.norm.rvs(size=(m, n), random_state=seed)
    perturb_fft = np.fft.rfft2(perturb)

    out = np.zeros((m, n))

    def worker(i):
        if std[i] > 0.0:
            filtered_noise = np.fft.irfft2(acf_fft_ampl[i] * perturb_fft, s=(m, n))
            filtered_noise /= np.std(filtered_noise)
            filtered_noise = stats.lognorm.ppf(
                stats.norm.cdf(filtered_noise), *dist_param[i]
            )
        else:
            filtered_noise = np.ones(weights[i].shape)

        return weights[i] * filtered_noise

    if DASK_IMPORTED and num_workers > 1:
        res = []
        for i in range(weights.shape[0]):
            res.append(dask.delayed(worker)(i))
        res = dask.compute(*res, num_workers=num_workers, scheduler="threads")
        for r in res:
            out += r
    else:
        for i in range(weights.shape[0]):
            out += worker(i)

    return out


def _get_acf_params(p):
    """Get ACF parameters from the given parameter vector."""
    return p[0], p[1], p[2], p[3] * p[2]


def _get_anisotropic_kernel_params(p):
    """Get anisotropic convolution kernel parameters from the given parameter
    vector."""
    return p[0], p[1], p[2] * p[1]


# TODO: use the method implemented in pysteps.timeseries.autoregression
def _iterate_ar_model(input_fields, psi):
    """Iterate autoregressive process."""
    input_field_new = 0.0

    for i, psi_ in enumerate(psi):
        input_field_new += psi_ * input_fields[-(i + 1), :]

    return np.concatenate([input_fields[1:, :], input_field_new[np.newaxis, :]])


def _linda_forecast(
    precip,
    precip_lagr_diff,
    timesteps,
    forecast_gen,
    precip_pert_gen,
    velocity_pert_gen,
    n_ensemble_members,
    seed,
    measure_time,
    print_info,
    return_output,
    callback,
):
    """Compute LINDA nowcast."""
    # compute convolved difference fields
    precip_lagr_diff = precip_lagr_diff.copy()

    for i in range(precip_lagr_diff.shape[0]):
        for _ in range(forecast_gen["ari_order"] - i):
            precip_lagr_diff[i] = _composite_convolution(
                precip_lagr_diff[i],
                forecast_gen["kernels_1"],
                forecast_gen["interp_weights"],
            )

    # initialize the random generators
    if precip_pert_gen is not None:
        rs_precip_pert = []
        np.random.seed(seed)
        for _ in range(n_ensemble_members):
            rs = np.random.RandomState(seed)
            rs_precip_pert.append(rs)
            seed = rs.randint(0, high=1e9)
    else:
        rs_precip_pert = None

    if velocity_pert_gen is not None:
        velocity_perturbators = []
        np.random.seed(seed)
        for _ in range(n_ensemble_members):
            vp = velocity_pert_gen["init_func"](seed)
            velocity_perturbators.append(
                lambda t, vp=vp: velocity_pert_gen["gen_func"](
                    vp, t * velocity_pert_gen["timestep"]
                )
            )
            seed = np.random.RandomState(seed).randint(0, high=1e9)
    else:
        velocity_perturbators = None

    state = {
        "precip_forecast": [precip[-1].copy() for _ in range(n_ensemble_members)],
        "precip_lagr_diff": [
            precip_lagr_diff.copy() for _ in range(n_ensemble_members)
        ],
        "rs_precip_pert": rs_precip_pert,
    }
    params = {
        "interp_weights": forecast_gen["interp_weights"],
        "kernels_1": forecast_gen["kernels_1"],
        "kernels_2": forecast_gen["kernels_2"],
        "mask_adv": forecast_gen["mask_adv"],
        "num_ens_members": n_ensemble_members,
        "num_workers": forecast_gen["num_workers"],
        "num_ensemble_workers": min(n_ensemble_members, forecast_gen["num_workers"]),
        "precip_pert_gen": precip_pert_gen,
        "psi": forecast_gen["psi"],
    }

    precip_forecast = nowcast_main_loop(
        precip[-1],
        forecast_gen["velocity"],
        state,
        timesteps,
        forecast_gen["extrap_method"],
        _update,
        extrap_kwargs=forecast_gen["extrap_kwargs"],
        velocity_pert_gen=velocity_perturbators,
        params=params,
        ensemble=True,
        num_ensemble_members=n_ensemble_members,
        callback=callback,
        return_output=return_output,
        num_workers=forecast_gen["num_workers"],
        measure_time=measure_time,
    )
    if measure_time:
        precip_forecast, mainloop_time = precip_forecast

    if return_output:
        if not forecast_gen["add_perturbations"]:
            precip_forecast = precip_forecast[0]
        if measure_time:
            return precip_forecast, mainloop_time
        else:
            return precip_forecast
    else:
        return None


def _linda_deterministic_init(
    precip,
    velocity,
    feature_method,
    max_num_features,
    feature_kwargs,
    ari_order,
    kernel_type,
    localization_window_radius,
    extrap_method,
    extrap_kwargs,
    add_perturbations,
    num_workers,
    measure_time,
):
    """Initialize the deterministic LINDA nowcast model."""
    forecast_gen = {}
    forecast_gen["velocity"] = velocity
    forecast_gen["extrap_method"] = extrap_method
    forecast_gen["ari_order"] = ari_order
    forecast_gen["add_perturbations"] = add_perturbations
    forecast_gen["num_workers"] = num_workers
    forecast_gen["measure_time"] = measure_time

    precip = precip[-(ari_order + 2) :]
    input_length = precip.shape[0]

    starttime_init = time.time()

    extrapolator = extrapolation.get_method(extrap_method)
    extrap_kwargs = extrap_kwargs.copy()
    extrap_kwargs["allow_nonfinite_values"] = True
    forecast_gen["extrapolator"] = extrapolator
    forecast_gen["extrap_kwargs"] = extrap_kwargs

    # detect features from the most recent input field
    if feature_method in {"blob", "shitomasi"}:
        precip_ = precip[-1].copy()
        precip_[~np.isfinite(precip_)] = 0.0
        feature_detector = feature.get_method(feature_method)

        if measure_time:
            starttime = time.time()

        feature_kwargs = feature_kwargs.copy()
        feature_kwargs["max_num_features"] = max_num_features

        feature_coords = np.fliplr(feature_detector(precip_, **feature_kwargs)[:, :2])

        feature_type = "blobs" if feature_method == "blob" else "corners"
        print("")
        print("Detecting features... ", end="", flush=True)
        if measure_time:
            print(
                f"found {feature_coords.shape[0]} {feature_type} in {time.time() - starttime:.2f} seconds."
            )
        else:
            print(f"found {feature_coords.shape[0]} {feature_type}.")

        if len(feature_coords) == 0:
            raise ValueError(
                "no features found, check input data and feature detector configuration"
            )
    elif feature_method == "domain":
        feature_coords = np.zeros((1, 2), dtype=int)
    else:
        raise NotImplementedError(
            "feature detector '%s' not implemented" % feature_method
        )
    forecast_gen["feature_coords"] = feature_coords

    # compute interpolation weights
    interp_weights = _compute_window_weights(
        feature_coords,
        precip.shape[1],
        precip.shape[2],
        localization_window_radius,
    )
    interp_weights /= np.sum(interp_weights, axis=0)
    forecast_gen["interp_weights"] = interp_weights

    # transform the input fields to the Lagrangian coordinates
    precip_lagr = np.empty(precip.shape)

    def worker(i):
        precip_lagr[i, :] = extrapolator(
            precip[i, :],
            velocity,
            input_length - 1 - i,
            "min",
            **extrap_kwargs,
        )[-1]

    if DASK_IMPORTED and num_workers > 1:
        res = []

    print("Transforming to Lagrangian coordinates... ", end="", flush=True)

    if measure_time:
        starttime = time.time()

    for i in range(precip.shape[0] - 1):
        if DASK_IMPORTED and num_workers > 1:
            res.append(dask.delayed(worker)(i))
        else:
            worker(i)

    if DASK_IMPORTED and num_workers > 1:
        dask.compute(*res, num_workers=min(num_workers, len(res)), scheduler="threads")
    precip_lagr[-1] = precip[-1]

    if measure_time:
        print(f"{time.time() - starttime:.2f} seconds.")
    else:
        print("done.")

    # compute advection mask and set nan to pixels, where one or more of the
    # advected input fields has a nan value
    mask_adv = np.all(np.isfinite(precip_lagr), axis=0)
    forecast_gen["mask_adv"] = mask_adv
    for i in range(precip_lagr.shape[0]):
        precip_lagr[i, ~mask_adv] = np.nan

    # compute differenced input fields in the Lagrangian coordinates
    precip_lagr_diff = np.diff(precip_lagr, axis=0)

    # estimate parameters of the deterministic model (i.e. the convolution and
    # the ARI process)

    print("Estimating the first convolution kernel... ", end="", flush=True)

    if measure_time:
        starttime = time.time()

    # estimate convolution kernel for the differenced component
    convol_weights = _compute_window_weights(
        feature_coords,
        precip.shape[1],
        precip.shape[2],
        localization_window_radius,
    )

    kernels_1 = _estimate_convol_params(
        precip_lagr_diff[-2],
        precip_lagr_diff[-1],
        convol_weights,
        mask_adv,
        kernel_type=kernel_type,
        num_workers=num_workers,
    )
    forecast_gen["kernels_1"] = kernels_1

    if measure_time:
        print(f"{time.time() - starttime:.2f} seconds.")
    else:
        print("done.")

    # compute convolved difference fields
    precip_lagr_diff_c = precip_lagr_diff[:-1].copy()
    for i in range(precip_lagr_diff_c.shape[0]):
        for _ in range(ari_order - i):
            precip_lagr_diff_c[i] = _composite_convolution(
                precip_lagr_diff_c[i],
                kernels_1,
                interp_weights,
            )

    print("Estimating the ARI(p,1) parameters... ", end="", flush=True)

    if measure_time:
        starttime = time.time()

    # estimate ARI(p,1) parameters
    weights = _compute_window_weights(
        feature_coords,
        precip.shape[1],
        precip.shape[2],
        localization_window_radius,
    )

    if ari_order == 1:
        psi = _estimate_ar1_params(
            precip_lagr_diff_c[-1],
            precip_lagr_diff[-1],
            weights,
            interp_weights,
            num_workers=num_workers,
        )
    else:
        psi = _estimate_ar2_params(
            precip_lagr_diff_c[-2:],
            precip_lagr_diff[-1],
            weights,
            interp_weights,
            num_workers=num_workers,
        )
    forecast_gen["psi"] = psi

    if measure_time:
        print(f"{time.time() - starttime:.2f} seconds.")
    else:
        print("done.")

    # apply the ARI(p,1) model and integrate the differences
    precip_lagr_diff_c = _iterate_ar_model(precip_lagr_diff_c, psi)
    precip_forecast = precip_lagr[-2] + precip_lagr_diff_c[-1]
    precip_forecast[precip_forecast < 0.0] = 0.0

    print("Estimating the second convolution kernel... ", end="", flush=True)

    if measure_time:
        starttime = time.time()

    # estimate the second convolution kernels based on the forecast field
    # computed above
    kernels_2 = _estimate_convol_params(
        precip_forecast,
        precip[-1],
        convol_weights,
        mask_adv,
        kernel_type=kernel_type,
        num_workers=num_workers,
    )
    forecast_gen["kernels_2"] = kernels_2

    if measure_time:
        print(f"{time.time() - starttime:.2f} seconds.")
    else:
        print("done.")

    if measure_time:
        return forecast_gen, precip_lagr_diff, time.time() - starttime_init
    else:
        return forecast_gen, precip_lagr_diff


def _linda_perturbation_init(
    precip,
    precip_lagr_diff,
    velocity,
    forecast_gen,
    pert_thrs,
    localization_window_radius,
    errdist_window_radius,
    acf_window_radius,
    vel_pert_method,
    vel_pert_kwargs,
    kmperpixel,
    timestep,
    num_workers,
    use_multiprocessing,
    measure_time,
):
    """Initialize the LINDA perturbation generator."""
    if measure_time:
        starttime = time.time()

    print("Estimating forecast errors... ", end="", flush=True)

    forecast_gen = forecast_gen.copy()
    forecast_gen["add_perturbations"] = False
    forecast_gen["num_ens_members"] = 1

    precip_forecast_det = _linda_forecast(
        precip[:-1],
        precip_lagr_diff[:-1],
        1,
        forecast_gen,
        None,
        None,
        1,
        None,
        False,
        False,
        True,
        None,
    )

    # compute multiplicative forecast errors
    err = precip_forecast_det[-1] / precip[-1]

    # mask small precipitation intensities
    mask = np.logical_or(
        np.logical_and(
            precip_forecast_det[-1] >= pert_thrs[1], precip[-1] >= pert_thrs[0]
        ),
        np.logical_and(
            precip_forecast_det[-1] >= pert_thrs[0], precip[-1] >= pert_thrs[1]
        ),
    )
    err[~mask] = np.nan

    if measure_time:
        print(f"{time.time() - starttime:.2f} seconds.")
    else:
        print("done.")

    pert_gen = _estimate_perturbation_params(
        err,
        forecast_gen,
        errdist_window_radius,
        acf_window_radius,
        localization_window_radius,
        measure_time,
        num_workers,
        use_multiprocessing,
    )

    if vel_pert_method == "bps":
        init_vel_noise, generate_vel_noise = noise.get_method("bps")

        vp_par = vel_pert_kwargs["vp_par"]
        vp_perp = vel_pert_kwargs["vp_perp"]

        kwargs = {
            "p_par": vp_par,
            "p_perp": vp_perp,
        }
        velocity_pert_gen = {
            "gen_func": generate_vel_noise,
            "init_func": lambda seed: init_vel_noise(
                velocity, 1.0 / kmperpixel, timestep, seed=seed, **kwargs
            ),
            "timestep": timestep,
        }
    else:
        velocity_pert_gen = None

    if measure_time:
        return pert_gen, velocity_pert_gen, time.time() - starttime
    else:
        return pert_gen, velocity_pert_gen


def _masked_convolution(field, kernel):
    """Compute "masked" convolution where non-finite values are ignored."""
    mask = np.isfinite(field)

    field = field.copy()
    field[~mask] = 0.0

    field_c = np.ones(field.shape) * np.nan
    field_c[mask] = convolve(field, kernel, mode="same")[mask]
    field_c[mask] /= convolve(mask.astype(float), kernel, mode="same")[mask]

    return field_c


def _update(state, params):
    def worker(j):
        state["precip_lagr_diff"][j] = _iterate_ar_model(
            state["precip_lagr_diff"][j], params["psi"]
        )
        state["precip_forecast"][j] += state["precip_lagr_diff"][j][-1]
        for i in range(state["precip_lagr_diff"][j].shape[0]):
            state["precip_lagr_diff"][j][i] = _composite_convolution(
                state["precip_lagr_diff"][j][i],
                params["kernels_1"],
                params["interp_weights"],
            )
        state["precip_forecast"][j] = _composite_convolution(
            state["precip_forecast"][j], params["kernels_2"], params["interp_weights"]
        )

        out = state["precip_forecast"][j].copy()
        out[out < 0.0] = 0.0
        out[~params["mask_adv"]] = np.nan

        # apply perturbations
        if params["precip_pert_gen"] is not None:
            seed = state["rs_precip_pert"][j].randint(0, high=1e9)
            perturb = _generate_perturbations(
                params["precip_pert_gen"], params["num_workers"], seed
            )
            out *= perturb

        return out

    out = []
    if DASK_IMPORTED and params["num_workers"] > 1 and params["num_ens_members"] > 1:
        res = []
        for j in range(params["num_ens_members"]):
            res.append(dask.delayed(worker)(j))
        out = dask.compute(
            *res, num_workers=params["num_ensemble_workers"], scheduler="threads"
        )
    else:
        for j in range(params["num_ens_members"]):
            out.append(worker(j))

    return np.stack(out), state


def _weighted_std(f, w):
    """
    Compute standard deviation of forecast errors with spatially varying weights.
    Values close to zero are omitted.
    """
    mask = np.abs(f - 1.0) > 1e-4
    n = np.count_nonzero(mask)
    if n > 0:
        c = (w[mask].size - 1.0) / n
        return np.sqrt(np.sum(w[mask] * (f[mask] - 1.0) ** 2.0) / (c * np.sum(w[mask])))
    else:
        return np.nan


def _window_tukey(m, n, ci, cj, ri, rj, alpha=0.5):
    """Tukey window function centered at the given coordinates."""
    j, i = np.meshgrid(np.arange(n), np.arange(m))

    di = np.abs(i - ci)
    dj = np.abs(j - cj)

    mask1 = np.logical_and(di <= ri, dj <= rj)

    w1 = np.zeros(di.shape)
    mask2 = di <= alpha * ri
    mask12 = np.logical_and(mask1, ~mask2)
    w1[mask12] = 0.5 * (
        1.0 + np.cos(np.pi * (di[mask12] - alpha * ri) / ((1.0 - alpha) * ri))
    )
    w1[np.logical_and(mask1, mask2)] = 1.0

    w2 = np.zeros(dj.shape)
    mask2 = dj <= alpha * rj
    mask12 = np.logical_and(mask1, ~mask2)
    w2[mask12] = 0.5 * (
        1.0 + np.cos(np.pi * (dj[mask12] - alpha * rj) / ((1.0 - alpha) * rj))
    )
    w2[np.logical_and(mask1, mask2)] = 1.0

    weights = np.zeros((m, n))
    weights[mask1] = w1[mask1] * w2[mask1]

    return weights


def _window_uniform(m, n, ci, cj, ri, rj):
    """Uniform window function with all values set to one."""
    return np.ones((m, n))
