from typing import Union
import numpy as np
from numpy import log, sqrt
from scipy.optimize import brentq


FLOAT_OR_ARRAY = Union[float, np.array]


def lake_problem(
    X: FLOAT_OR_ARRAY,
    a: FLOAT_OR_ARRAY = 0.1,
    q: FLOAT_OR_ARRAY = 2.0,
    b: FLOAT_OR_ARRAY = 0.42,
    eps: FLOAT_OR_ARRAY = 0.02,
) -> float:
    """Lake Problem as given in Hadka et al., (2015) and Kwakkel (2017) modified for use as a test function.

    The `mean` and `stdev` parameters control the log normal distribution of
    natural inflows (`epsilon` in [1] and [2]).

    References
    ----------
    .. [1] Hadka, D., Herman, J., Reed, P., Keller, K., (2015).
           "An open source framework for many-objective robust decision making."
           Environmental Modelling & Software 74, 114-129.
           doi:10.1016/j.envsoft.2015.07.014

    .. [2] Kwakkel, J.H, (2017). "The Exploratory Modeling Workbench: An open
           source toolkit for exploratory modeling, scenario discovery, and
           (multi-objective) robust decision making."
           Environmental Modelling & Software 96, 239-250.
           doi:10.1016/j.envsoft.2017.06.054

    .. [3] Singh, R., Reed, P., Keller, K., (2015). "Many-objective robust
           decision making for managing an ecosystem with a deeply uncertain
           threshold response."
           Ecology and Society 20.
           doi:10.5751/ES-07687-200312


    Parameters
    ----------
    X : float or np.ndarray
        normalized concentration of Phosphorus at point in time
    a : float or np.ndarray
        rate of anthropogenic pollution (0.0 to 0.1)
    q : float or np.ndarray
        exponent controlling recycling rate (2.0 to 4.5).
    b : float or np.ndarray
        decay rate for phosphorus
        (0.1 to 0.45, where default 0.42 is irreversible, as described in [1])
    eps : float or np.ndarray
        natural inflows of phosphorus (pollution), see [3]



    Returns
    -------
    float, phosphorus pollution for a point in time
    """
    Xq = X**q
    X_t1 = X + a + (Xq / (1.0 + Xq)) - (b * X) + eps

    return X_t1


def evaluate_lake(values: np.ndarray, seed=101) -> np.ndarray:
    """Evaluate the Lake Problem with an array of parameter values.

    References
    ----------
    .. [1] Hadka, D., Herman, J., Reed, P., Keller, K., (2015).
           "An open source framework for many-objective robust decision
                making."
           Environmental Modelling & Software 74, 114–129.
           doi:10.1016/j.envsoft.2015.07.014

    .. [2] Singh, R., Reed, P., Keller, K., (2015). "Many-objective robust
           decision making for managing an ecosystem with a deeply uncertain
           threshold response."
           Ecology and Society 20.
           doi:10.5751/ES-07687-200312

    Parameters
    ----------
    values : np.ndarray,
             model inputs in the (column) order of
             a, q, b, mean, stdev

             where
             * `a` is rate of anthropogenic pollution
             * `q` is exponent controlling recycling rate
             * `b` is decay rate for phosphorus
             * `mean` and
             * `stdev` set the log normal distribution of `eps`, see [2]

    Returns
    -------
    np.ndarray, of Phosphorus pollution over time `t`
    """
    rng = np.random.default_rng(seed)

    nvars = values.shape[0]

    a, q, b, mean, stdev = values.T

    sq_mean = mean**2
    sq_std = stdev**2
    eps = rng.lognormal(
        log(sq_mean / sqrt(sq_std + sq_mean)),
        sqrt(log(1.0 + sq_std / sq_mean)),
        size=nvars,
    )

    Y = np.zeros((nvars, nvars))
    for t in range(nvars):
        # First X value will be last Y value (should be 0 as we are filling
        # an array of zeros)
        Y[t] = lake_problem(Y[t - 1], a, q, b, eps)

    return Y


def evaluate(values: np.ndarray, nvars: int = 100, seed=101):
    """Evaluate the Lake Problem with an array of parameter values.

    Parameters
    ----------
    values : np.ndarray,
             model inputs in the (column) order of
             a, q, b, mean, stdev, delta, alpha

    nvars : int,
            number of decision variables to simulate (default: 100)


    Returns
    -------
    np.ndarray : max_P, utility, inertia, reliability
    """
    a, q, b, _, __, delta, alpha = values.T

    nsamples = len(a)
    Y = np.empty((nsamples, 4))
    for i in range(nsamples):
        tmp = evaluate_lake(values[i, :5], seed=seed)

        a_i = a[i]
        q_i = q[i]

        Pcrit = brentq(lambda x: x**q_i / (1.0 + x**q_i) - b[i] * x, 0.01, 1.5)

        reliability = len(tmp[tmp < Pcrit]) / nvars

        max_P = np.max(tmp)
        utility = np.sum(alpha[i] * a_i * np.power(delta[i], np.arange(nvars)))

        # In practice, `a` will be set by a separate decision model
        # See [2] in docstring for `lake_problem()`
        # Here, it is a constant for a given scenario.
        # The value for `tau` (0.02) is taken from [2].
        inertia = len(a_i[a_i < 0.02]) / nvars

        Y[i, :] = max_P, utility, inertia, reliability

    return Y


if __name__ == "__main__":
    from SALib.sample import latin
    from SALib.analyze import delta

    SEED_VAL = 101

    LAKE_SPEC = {
        "num_vars": 7,
        "names": ["a", "q", "b", "mean", "stdev", "delta", "alpha"],
        "bounds": [
            [0.0, 0.1],
            [2.0, 4.5],
            [0.1, 0.45],
            [0.01, 0.05],
            [0.001, 0.005],
            [0.93, 0.99],
            [0.2, 0.5],
        ],
        "outputs": ["max_P", "Utility", "Inertia", "Reliability"],
    }

    latin_samples = latin.sample(LAKE_SPEC, 1000, seed=SEED_VAL)
    Y = evaluate(latin_samples)

    for i, name in enumerate(LAKE_SPEC["outputs"]):
        Si = delta.analyze(LAKE_SPEC, latin_samples, Y[:, i])
        print(name)
        print(Si.to_df())
