# -*- coding: utf-8 -*-
"""
pysteps.nowcasts.linear_blending
================================

Linear blending method to blend (ensemble) nowcasts and (ensemble) Numerical
Weather Prediction (NWP) forecasts. The method uses a predefined start and end
time for the blending. Before this start time, the resulting blended forecasts only
consists of the nowcast(s). In between the start and end time, the nowcast(s)
weight decreases and NWP forecasts weight increases linearly from 1(0) to
0(1). After the end time, the blended forecast entirely consists of the NWP
forecasts. The saliency-based blending method also takes into account the pixel
intensities and preserves them if they are strong enough based on their ranked salience.

Implementation of the linear blending and saliency-based blending between nowcast and NWP data.

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

    forecast
"""

import numpy as np
from pysteps import nowcasts
from pysteps.utils import conversion
from scipy.stats import rankdata


def forecast(
    precip,
    precip_metadata,
    velocity,
    timesteps,
    timestep,
    nowcast_method,
    precip_nwp=None,
    precip_nwp_metadata=None,
    start_blending=120,
    end_blending=240,
    fill_nwp=True,
    saliency=False,
    nowcast_kwargs=None,
):
    """Generate a forecast by linearly or saliency-based blending of nowcasts with NWP data

    Parameters
    ----------
    precip: array_like
      Array containing the input precipitation field(s) ordered by timestamp
      from oldest to newest. The time steps between the inputs are assumed
      to be regular.
    precip_metadata: dict
        Metadata dictionary containing (at least) the transform, unit and threshold
        attributes as described in the documentation of :py:mod:`pysteps.io.importers`.
    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. All values are required to be finite.
    timesteps: int
      Number of time steps to forecast.
    timestep: int or float
      The time difference (in minutes) between consecutive forecast fields.
    nowcast_method: str
      Name of the nowcasting method. See :py:mod:`pysteps.nowcasts.interface`
      for the list of available methods.
    precip_nwp: array_like or NoneType, optional
      Array of shape (timesteps, m, n) in the case of no ensemble or
      of shape (n_ens_members, timesteps, m, n) in the case of an ensemble
      containing the NWP precipitation fields ordered by timestamp from oldest
      to newest. The time steps between the inputs are assumed to be regular
      (and identical to the time step between the nowcasts). If no NWP
      data is given the value of precip_nwp is None and no blending will be performed.
    precip_nwp_metadata: dict or NoneType, optional
        NWP metadata dictionary containing (at least) the transform, unit and threshold
        attributes as described in the documentation of :py:mod:`pysteps.io.importers`.
    start_blending: int, optional
      Time stamp (in minutes) after which the blending should start. Before this
      only the nowcast data is used.
    end_blending: int, optional
      Time stamp (in minutes) after which the blending should end. Between
      start_blending and end_blending the nowcasts and NWP data are linearly
      merged with each other. After end_blending only the NWP data is used.
    fill_nwp: bool, optional
      Standard value is True. If True, the NWP data will be used to fill in the
      no data mask of the nowcast.
    saliency: bool, optional
      Default value is False. If True, saliency will be used for blending. The blending
      is based on intensities and forecast times as described in :cite:`Hwang2015`. The blended
      product preserves pixel intensities with time if they are strong enough based on their ranked
      salience.
    nowcast_kwargs: dict, optional
      Dictionary containing keyword arguments for the nowcast method.


    Returns
    -------
    precip_blended: ndarray
      Array of shape (timesteps, m, n) in the case of no ensemble or
      of shape (n_ens_members, timesteps, m, n) in the case of an ensemble
      containing the precipation forecast generated by linearly blending
      the nowcasts and the NWP data. n_ens_members equals the maximum no. of
      ensemble members in either the nowcast or nwp model(s).
    """

    if nowcast_kwargs is None:
        nowcast_kwargs = dict()

    # Calculate the nowcasts
    nowcast_method_func = nowcasts.get_method(nowcast_method)
    precip_nowcast = nowcast_method_func(
        precip,
        velocity,
        timesteps,
        **nowcast_kwargs,
    )

    # Make sure that precip_nowcast and precip_nwp are in mm/h
    precip_nowcast, _ = conversion.to_rainrate(precip_nowcast, metadata=precip_metadata)

    # Check if NWP data is given as input
    if precip_nwp is not None:
        precip_nwp, _ = conversion.to_rainrate(precip_nwp, metadata=precip_nwp_metadata)

        if len(precip_nowcast.shape) == 4:
            n_ens_members_nowcast = precip_nowcast.shape[0]
            if n_ens_members_nowcast == 1:
                precip_nowcast = np.squeeze(precip_nowcast)
        else:
            n_ens_members_nowcast = 1

        if len(precip_nwp.shape) == 4:
            n_ens_members_nwp = precip_nwp.shape[0]
            if n_ens_members_nwp == 1:
                precip_nwp = np.squeeze(precip_nwp)
        else:
            n_ens_members_nwp = 1

        # Now, repeat the nowcast ensemble members or the nwp models/members until
        # it has the same amount of members as n_ens_members_max. For instance, if
        # you have 10 ensemble nowcasts members and 3 NWP members, the output will
        # be an ensemble of 10 members. Hence, the three NWP members are blended
        # with the first three members of the nowcast (member one with member one,
        # two with two, etc.), subsequently, the same NWP members are blended with
        # the next three members (NWP member one with member 4, NWP member 2 with
        # member 5, etc.), until 10 is reached.
        n_ens_members_max = max(n_ens_members_nowcast, n_ens_members_nwp)
        n_ens_members_min = min(n_ens_members_nowcast, n_ens_members_nwp)

        if n_ens_members_min != n_ens_members_max:
            if n_ens_members_nwp == 1:
                precip_nwp = np.repeat(
                    precip_nwp[np.newaxis, :, :], n_ens_members_max, axis=0
                )
            elif n_ens_members_nowcast == 1:
                precip_nowcast = np.repeat(
                    precip_nowcast[np.newaxis, :, :], n_ens_members_max, axis=0
                )
            else:
                repeats = [
                    (n_ens_members_max + i) // n_ens_members_min
                    for i in range(n_ens_members_min)
                ]

                if n_ens_members_nwp == n_ens_members_min:
                    precip_nwp = np.repeat(precip_nwp, repeats, axis=0)
                elif n_ens_members_nowcast == n_ens_members_min:
                    precip_nowcast = np.repeat(precip_nowcast, repeats, axis=0)

        # Check if dimensions are correct
        assert (
            precip_nwp.shape == precip_nowcast.shape
        ), "The dimensions of precip_nowcast and precip_nwp need to be identical: dimension of precip_nwp = {} and dimension of precip_nowcast = {}".format(
            precip_nwp.shape, precip_nowcast.shape
        )

        # Initialise output
        precip_blended = np.zeros_like(precip_nowcast)

        # Calculate the weights
        for i in range(timesteps):
            # Calculate what time we are at
            t = (i + 1) * timestep

            if n_ens_members_max == 1:
                ref_dim = 0
            else:
                ref_dim = 1

            # apply blending
            # compute the slice indices
            slc_id = _get_slice(precip_blended.ndim, ref_dim, i)

            # Calculate the weight with a linear relation (weight_nwp at start_blending = 0.0)
            # and (weight_nwp at end_blending = 1.0)
            weight_nwp = (t - start_blending) / (end_blending - start_blending)

            # Set weights at times before start_blending and after end_blending
            if weight_nwp <= 0.0:
                weight_nwp = 0.0
                precip_blended[slc_id] = precip_nowcast[slc_id]

            elif weight_nwp >= 1.0:
                weight_nwp = 1.0
                precip_blended[slc_id] = precip_nwp[slc_id]

            else:
                # Calculate weight_nowcast
                weight_nowcast = 1.0 - weight_nwp

                # Calculate output by combining precip_nwp and precip_nowcast,
                # while distinguishing between ensemble and non-ensemble methods
                if saliency:
                    ranked_salience = _get_ranked_salience(
                        precip_nowcast[slc_id], precip_nwp[slc_id]
                    )
                    ws = _get_ws(weight_nowcast, ranked_salience)
                    precip_blended[slc_id] = (
                        ws * precip_nowcast[slc_id] + (1 - ws) * precip_nwp[slc_id]
                    )

                else:
                    precip_blended[slc_id] = (
                        weight_nwp * precip_nwp[slc_id]
                        + weight_nowcast * precip_nowcast[slc_id]
                    )

            # Find where the NaN values are and replace them with NWP data
            if fill_nwp:
                nan_indices = np.isnan(precip_blended)
                precip_blended[nan_indices] = precip_nwp[nan_indices]
    else:
        # If no NWP data is given, the blended field is simply equal to the nowcast field
        precip_blended = precip_nowcast

    return precip_blended


def _get_slice(n_dims, ref_dim, ref_id):
    """source: https://stackoverflow.com/a/24399139/4222370"""
    slc = [slice(None)] * n_dims
    slc[ref_dim] = ref_id
    return tuple(slc)


def _get_ranked_salience(precip_nowcast, precip_nwp):
    """Calculate ranked salience, which show how close the pixel is to the maximum intensity difference [r(x,y)=1]
      or the minimum intensity difference [r(x,y)=0]

    Parameters
    ----------
    precip_nowcast: array_like
      Array of shape (m,n) containing the extrapolated precipitation field at a specified timestep
    precip_nwp: array_like
      Array of shape (m,n) containing the NWP fields at a specified timestep

    Returns
    -------
    ranked_salience:
      Array of shape (m,n) containing ranked salience
    """

    # calcutate normalized intensity
    if np.max(precip_nowcast) == 0:
        norm_nowcast = np.zeros_like(precip_nowcast)
    else:
        norm_nowcast = precip_nowcast / np.max(precip_nowcast)

    if np.max(precip_nwp) == 0:
        norm_nwp = np.zeros_like(precip_nwp)
    else:
        norm_nwp = precip_nwp / np.max(precip_nwp)

    diff = norm_nowcast - norm_nwp

    # Calculate ranked salience, based on dense ranking method, in which equally comparable values receive the same ranking number
    ranked_salience = rankdata(diff, method="dense").reshape(diff.shape).astype("float")
    ranked_salience /= ranked_salience.max()

    return ranked_salience


def _get_ws(weight, ranked_salience):
    """Calculate salience weight based on linear weight and ranked salience as described in :cite:`Hwang2015`.
      Cells with higher intensities result in larger weights.

    Parameters
    ----------
    weight: int
      Varying between 0 and 1
    ranked_salience: array_like
      Array of shape (m,n) containing ranked salience

    Returns
    -------
    ws: array_like
      Array of shape (m,n) containing salience weight, which preserves pixel intensties with time if they are strong
      enough based on the ranked salience.
    """

    # Calculate salience weighte
    ws = 0.5 * (
        (weight * ranked_salience)
        / (weight * ranked_salience + (1 - weight) * (1 - ranked_salience))
        + (
            np.sqrt(ranked_salience**2 + weight**2)
            / (
                np.sqrt(ranked_salience**2 + weight**2)
                + np.sqrt((1 - ranked_salience) ** 2 + (1 - weight) ** 2)
            )
        )
    )
    return ws
