#Define the input variable functions
import numpy as np

from climada.entity import ImpactFunc, ImpactFuncSet, Exposures
from climada.util.constants import EXP_DEMO_H5, HAZ_DEMO_H5
from climada.hazard import Hazard
from functools import partial
import scipy as sp
from climada.engine.unsequa import InputVar
from climada.engine.unsequa import CalcImpact
import time

def impf_func(G=1, v_half=84.7, vmin=25.7, k=3, _id=1):

    def xhi(v, v_half, vmin):
        return max([(v - vmin), 0]) / (v_half - vmin)

    def sigmoid_func(v, G, v_half, vmin, k):
        return G * xhi(v, v_half, vmin)**k / (1 + xhi(v, v_half, vmin)**k)

    #In-function imports needed only for parallel computing on Windows
    import numpy as np
    from climada.entity import ImpactFunc, ImpactFuncSet
    intensity_unit = 'm/s'
    intensity = np.linspace(0, 150, num=100)
    mdd = np.repeat(1, len(intensity))
    paa = np.array([sigmoid_func(v, G, v_half, vmin, k) for v in intensity])
    imp_fun = ImpactFunc("TC", _id, intensity, mdd, paa, intensity_unit)
    imp_fun.check()
    impf_set = ImpactFuncSet([imp_fun])
    return impf_set

haz = Hazard.from_hdf5(HAZ_DEMO_H5)
exp_base = Exposures.from_hdf5(EXP_DEMO_H5)

def prep_full_calc(haz, exp_base, impf_func):
    #It is a good idea to assign the centroids to the base exposures in order to avoid repeating this
    # potentially costly operation for each sample.
    exp_base.assign_centroids(haz)
    def exp_base_func(x_exp, exp_base):
        exp = exp_base.copy()
        exp.gdf['value'] *= x_exp
        return exp
    exp_func = partial(exp_base_func, exp_base=exp_base)
    #Define the InputVars


    exp_distr = {"x_exp": sp.stats.beta(10, 1.1)} #This is not really a reasonable distribution but is used
                                                #here to show that you can use any scipy distribution.

    exp_iv = InputVar(exp_func, exp_distr)

    impf_distr = {
        "G": sp.stats.truncnorm(0.5, 1.5),
        "v_half": sp.stats.uniform(35, 65),
        "vmin": sp.stats.uniform(0, 15),
        "k": sp.stats.uniform(1, 4)
        }
    impf_iv = InputVar(impf_func, impf_distr)
    calc_imp = CalcImpact(exp_iv, impf_iv, haz)
    output_imp = calc_imp.make_sample(N=2**7, sampling_kwargs={'skip_values': 2**8})
    return calc_imp, output_imp

def run_unc_calc(calc_imp, output_imp):
    calc_imp.uncertainty(output_imp, rp = [50, 100, 250])

calc_imp, output_imp = prep_full_calc(haz, exp_base, impf_func)

%timeit -r 20 run_unc_calc(calc_imp, output_imp)