"""
Use the Ratio of Uniforms algorithm to sample a distribution
============================================================
"""

# %%
# In this example, we illustrate in dimension :math:`\inputDim = 1` the ratio of uniforms algorithm  to sample a distribution which
# probability density function is known up to a normalization factor.
#
# Consider a scalar distribution whose
# probability density function :math:`p` is known up to a normalization factor:
# let :math:`f: \Rset \rightarrow \Rset^+` be a function such that
# :math:`p(x) = cf(x)` for all :math:`x \in \Rset` with :math:`c \in \Rset^+_*`.
#
# Let :math:`r \in \Rset^+_*`. Let  :math:`A_{f,r}` be the domain defined by:
#
# .. math::
#
#    A_{f,r} = \left\{ (u, v) \in \Rset^2 \, |\, 0 \leq u \leq f\left(\dfrac{v}{u^r}\right)^{\frac{1}{1+r}}\right \}
#
# The Lebesgue measure of :math:`A_{f,r}`  is equal to :math:`\dfrac{1}{c(1+r)}`.
#
# Let :math:`(U, V)` be a random variable uniformly distributed on :math:`A_{f,r}`.
# Then, :math:`X = \dfrac{V}{U^r}` is a random variable on
# :math:`\Rset` distributed according to :math:`p`.
#
# Under the condition that :math:`f(x)^{\frac{1}{1+r}}` and
# :math:`x f(x)^{\frac{1}{1+r}}` are bounded, then the domain
# :math:`A_{f,r}` is bounded and can be included in a rectangular bounding box :math:`\tilde{A}_{f,r}`
# defined by:
#
# .. math::
#
#    \tilde{A}_{f,r} & = \left[0, \sup_{x} f(x)^{\frac{1}{1+r}} \right] \times
#    \left[ \inf_{x} x f(x)^{\frac{1}{1+r}} ,  \sup_{x} x f(x)^{\frac{1}{1+r}}\right]\\
#                    & = \left[0, u_{sup} \right] \times \left[v_{inf}, v_{sup}\right]
#
# Note that :math:`u_{sup} > 0`, :math:`v_{sup} \geq 0` and  :math:`v_{inf} \leq 0`.
#
# This allows one to sample uniformly the domain :math:`A_{f,r}` by rejection sampling inside
# :math:`\tilde{A}_{f,r}`.
#
# Study of the :math:`(u,v)`-space.
# ---------------------------------
# The boundary of :math:`\tilde{A}_{f,r}` is defined by:
#
# .. math::
#
#     \partial \tilde{A}_{f,r} = \cF_1 \cup \cF_2 \cup \cF_3 \cup \cF_4
#
# where the boundaries :math:`\cF_i` are defined by:
#
# .. math::
#
#    \cF_1 & = \left \{ (u, v_{sup}), 0 \leq u \leq u_{sup} \right \} \\
#    \cF_2 & = \left \{ (u, v_{inf}), 0 \leq u \leq u_{sup} \right \}\\
#    \cF_3 & = \left \{ (u_{sup}, v), v_{inf} \leq v \leq v_{sup} \right \}\\
#    \cF_4 & = \left \{ (0, v), v_{inf} \leq v \leq v_{sup} \right \}
#
# The boundary of :math:`A_{f,r}` is defined by:
#
# .. math::
#
#    \partial A_{f,r}  = \left\{ (u,v) = \left( f(x)^{\dfrac{1}{1+r}}, x f(x)^{\dfrac{r}{1+r}} \right),
#    x \in \Rset \right\}
#
# The following function draws the domain :math:`A_{f,r}` and the bounding box :math:`\partial \tilde{A}_{f,r}`
# in the :math:`(u,v)`-space.
import openturns as ot
import openturns.viewer as otv
import openturns.experimental as otexp
from math import pi, exp


def draw_RatioOfUniforms_UV(ratio_algo):
    r = ratio_algo.getR()
    range_logF = ratio_algo.getRange()
    xInf = range_logF.getLowerBound()[0]
    xSup = range_logF.getUpperBound()[0]
    uSup = ratio_algo.getSupU()
    vInf = ratio_algo.getInfV()[0]
    vSup = ratio_algo.getSupV()[0]
    logF = ratio_algo.getLogUnscaledPDF()

    # \partial A_{f,r} domain
    x = xInf
    u = 0.0
    v = 0.0
    xPrec = None
    uPrec = None
    vPrec = None
    first = True
    step = 0.01 * (uSup + vSup - vInf)
    delta = step
    data_UV = ot.Sample(0, 2)
    while x < xSup:
        val_F = exp(logF([x])[0])
        u = val_F ** (1.0 / (1.0 + r))
        v = x * val_F ** (r / (1.0 + r))
        # Step adaptation after the first point
        if not first:
            # If the step is too large
            while abs(u - uPrec) + abs(v - vPrec) > 2 * delta:
                step *= 0.5
                x = xPrec + step
                if x < xSup:
                    val_F = exp(logF([x])[0])
                    u = val_F ** (1.0 / (1.0 + r))
                    v = x * val_F ** (r / (1.0 + r))
            # If the step is too small
            while abs(u - uPrec) + abs(v - vPrec) < 0.5 * delta and x < xSup:
                step *= 2.0
                x = xPrec + step
                if x < xSup:
                    val_F = exp(logF([x])[0])
                    u = val_F ** (1.0 / (1.0 + r))
                    v = x * val_F ** (r / (1.0 + r))
        xPrec = x
        uPrec = u
        vPrec = v
        x += step
        first = False
        data_UV.add([u, v])
    # to get a closed boundary
    data_UV.add(data_UV[0])
    g = ot.Graph("", "u", "v")
    g.setLegendPosition("bottomright")
    data_A = ot.Curve(data_UV)
    data_A.setColor("purple")
    g.add(data_A)

    # \partial \tilde{A}_{f,r} bounding box
    boundary_1 = ot.Curve([[0, vSup], [uSup, vSup]])
    boundary_1.setColor("blue")
    boundary_2 = ot.Curve([[0, vInf], [uSup, vInf]])
    boundary_2.setColor("red")
    boundary_3 = ot.Curve([[uSup, vInf], [uSup, vSup]])
    boundary_3.setColor("green")
    boundary_4 = ot.Curve([[0, vInf], [0, vSup]])
    boundary_4.setColor("black")
    g.add(boundary_1)
    g.add(boundary_2)
    g.add(boundary_3)
    g.add(boundary_4)

    leg = g.getLegends()
    leg[0] = r"boundary of $A_{f,r}$"
    leg[1] = r"$\mathcal{F}_1$"
    leg[2] = r"$\mathcal{F}_2$"
    leg[3] = r"$\mathcal{F}_3$"
    leg[4] = r"$\mathcal{F}_4$"
    g.setLegends(leg)

    return g


# %%
# Study of the :math:`(x,y)`-space.
# ---------------------------------
# We use the mapping function :math:`\varphi` defined by:
#
# .. math::
#
#    \varphi: & \Rset^2 \rightarrow \Rset^2 \\
#             & (u,v)  \rightarrow \left( \dfrac{v}{u^r}, u \right)
#
# Then we have:
#
# .. math::
#
#    \varphi(A_{f,r}) = \left\{ (x,y) \, |\, x \in \Rset, 0 \leq y \leq f(x)^{\dfrac{1}{1+r}} \right\}
#
# and:
#
# .. math::
#
#     \partial \varphi(\tilde{A}_{f,r}) = \varphi(\cF_1) \cup \varphi(\cF_2) \cup \varphi(\cF_3)
#
# where the boundaries :math:`\varphi(\cF_i)` are defined by:
#
# .. math::
#
#    \varphi(\cF_1) & = \left\{ \varphi(u, v_{sup}), 0 \leq u \leq u_{sup} \right\}\\
#          & = \left\{ \left(x,\left(\dfrac{v_{sup}}{x}\right)^{1/r} \right),
#              x \geq \dfrac{v_{sup}}{u_{sup}^{r}} \right \}
#
# .. math::
#
#    \varphi(\cF_2) & = \left \{ \varphi(u, v_{inf}), 0 \leq u \leq u_{sup} \right\}\\
#          & = \left\{ \left(x,\left(\dfrac{v_{inf}}{x}\right)^{1/r}\right),
#              x \leq \dfrac{v_{inf}}{u_{sup}^{r}} \right \}
#
# .. math::
#
#    \varphi(\cF_3) & = \left\{ \varphi(u_{sup}, v), v_{inf} \leq v \leq v_{sup} \right\}\\
#          & = \left\{ (x, u_{sup}),
#              \dfrac{v_{inf}}{u_{sup}^{r}} \leq x \leq \dfrac{v_{sup}}{u_{sup}^{r}}  \right\}
#
# The following function draws the boundary :math:`\partial \varphi(A_{f,r})` and the bounding box
# :math:`\partial \tilde{A}_{f,r}` in the :math:`(x,y)`-space.


def draw_RatioOfUniforms_XY(ratio_algo, nPoints):
    r = ratio_algo.getR()
    range_logF = ratio_algo.getRange()
    xInf = range_logF.getLowerBound()[0]
    xSup = range_logF.getUpperBound()[0]
    uSup = ratio_algo.getSupU()
    vInf = ratio_algo.getInfV()[0]
    vSup = ratio_algo.getSupV()[0]
    logF = ratio_algo.getLogUnscaledPDF()

    # \partial varpĥi(A_{f,r}) domain
    eps = 1e-10
    step = (xSup - xInf - 2.0 * eps) / (nPoints - 1)
    data_XY = ot.Sample(nPoints, 2)
    for i in range(nPoints):
        x = xInf + eps + i * step
        val_F = exp(logF([x])[0])
        data_XY[i, 0] = x
        data_XY[i, 1] = val_F ** (1.0 / (1.0 + r))
    g = ot.Graph("", "x", "y")
    g.setLegendPosition("lower center")
    data_A = ot.Curve(data_XY)
    data_A.setColor("purple")
    g.add(data_A)

    # \partial varphi(\tilde{A}_{f,r}) bounding box
    leg = ot.Description()
    leg.add(r"boundary of $\varphi(A_{f,r})$")
    boundary_1_fct = ot.SymbolicFunction(
        "x", "(" + str(vSup) + "/x)^(1/" + str(r) + ")"
    )
    if vSup / (uSup**r) < xSup:
        boundary_1 = boundary_1_fct.draw(vSup / (uSup**r), xSup, 101)
        boundary_1.setColors(["blue"])
        leg.add(r"$\varphi(\mathcal{F}_1)$")
        g.add(boundary_1)
    if vInf / (uSup**r) < uSup:
        boundary_2_fct = ot.SymbolicFunction(
            "x", "(" + str(vInf) + "/x)^(1/" + str(r) + ")"
        )
        boundary_2 = boundary_2_fct.draw(xInf, vInf / (uSup**r), 101)
        boundary_2.setColors(["red"])
        leg.add(r"$\varphi(\mathcal{F}_2)$")
        g.add(boundary_2)
    boundary_3 = ot.Curve([[vInf / (uSup**r), uSup], [vSup / (uSup**r), uSup]])
    boundary_3.setColor("green")
    g.add(boundary_3)
    leg.add(r"$\varphi(\mathcal{F}_3)$")
    g.setLegends(leg)

    return g


# %%
# Illustration.
# -------------
# We create the function :math:`\log f` defined by:
#
# .. math::
#    :label: robertEx
#
#    \log f(x) = \log(\cos(x)) + x,  \quad  x \in \left[-\dfrac{\pi}{2}, \dfrac{\pi}{2}\right]
#
# and the :class:`~openturns.experimental.RatioOfUniforms`.


log_UnscaledPDF = ot.SymbolicFunction("x", "log(cos(x)) + x")
eps = 1e-5
range_PDF = ot.Interval(-pi / 2.0 + eps, pi / 2.0 - eps)
ratio_Algo = otexp.RatioOfUniforms(log_UnscaledPDF, range_PDF)

# %%
# We draw the boundary :math:`\partial A_{f,r}` and the bounding box :math:`\partial \tilde{A}_{f,r}`
# in the :math:`(u,v)`-space.
g_UV = draw_RatioOfUniforms_UV(ratio_Algo)
g_UV.setTitle(
    rf"$r = ${ratio_Algo.getR():.3g}, $\tau = ${ratio_Algo.getAcceptanceRatio():.3g}"
)

# %%
# We draw the boundary :math:`\partial \varphi(A_{f,r})` and the bounding box :math:`\partial \varphi(\tilde{A}_{f,r})`
# in the :math:`(x,y)`-space.
# sphinx_gallery_thumbnail_number = 1
g_XY = draw_RatioOfUniforms_XY(ratio_Algo, 1001)
g_XY.setTitle(
    rf"$r = ${ratio_Algo.getR():.3g}, $\tau = ${ratio_Algo.getAcceptanceRatio():.3g}"
)

grid_default = ot.GridLayout(1, 2)
grid_default.setGraph(0, 0, g_UV)
grid_default.setGraph(0, 1, g_XY)
grid_default.setTitle(r"$\log f(x) = \log(\cos(x)) + x, x \in ]-\pi/2, \pi/2[$")
view = otv.View(grid_default)

# %%
# We change the :math:`r` value and we draw the same plots.
grid_full = ot.GridLayout(4, 2)
for i, r in enumerate([0.5, 0.8, 1.2, 2.0]):
    ratio_Algo.setR(r)
    g_UV = draw_RatioOfUniforms_UV(ratio_Algo)
    g_UV.setLegendPosition("")
    g_UV.setTitle(
        rf"$r = ${ratio_Algo.getR():.3g}, $\tau = ${ratio_Algo.getAcceptanceRatio():.3g}"
    )

    g_XY = draw_RatioOfUniforms_XY(ratio_Algo, 1001)
    g_XY.setLegendPosition("")
    g_XY.setTitle(
        rf"$r = ${ratio_Algo.getR():.3g}, $\tau = ${ratio_Algo.getAcceptanceRatio():.3g}"
    )

    grid_full.setGraph(i, 0, g_UV)
    grid_full.setGraph(i, 1, g_XY)

grid_full.setTitle(r"$\log f(x) = \log(\cos(x)) + x,  \quad x \in ]-\pi/2, \pi/2[$")
view = otv.View(grid_full)

# %%
# Study of the acceptance ratio.
# ------------------------------
# We study the acceptance ratio :math:`\tau(r)` as a function of the parameter :math:`r` for some distributions.
# In these cases, the value is exact and computed from the volume of :math:`A_{f,r}` and the one of :math:`\tilde{A}_{f,r}`.
# We draw the function :math:`r \rightarrow \tau(r)`. We conclude that the default value :math:`r=1` seems to be correct
# in all these cases.
#
mixt_gaus = ot.Mixture([ot.Normal(-2.0, 1.0), ot.Normal(4.0, 1.0)], [0.2, 0.8])
logPDF_list = [
    ot.Normal().getLogPDF(),
    ot.Normal(2).getLogPDF(),
    ot.Normal(3).getLogPDF(),
    ot.Exponential().getLogPDF(),
    mixt_gaus.getLogPDF(),
]
range_list = [
    ot.Normal().getRange(),
    ot.Normal(2).getRange(),
    ot.Normal(3).getRange(),
    ot.Exponential().getRange(),
    mixt_gaus.getRange(),
]
name_list = ["Normal(1)", "Normal(2)", "Normal(3)", "Exponential()", "Mixture Normal"]

g = ot.Graph(r"Acceptance ratio $\tau$ as a function of $r$", r"$r$", r"$\tau$")
g.setLegendPosition("upper right")
g.setLogScale(1)
numberR = 50
a = (1e2) ** (2.0 / numberR)
data_r_ratio = ot.Sample(numberR + 1, 2)
isScaled = True
for k in range(len(logPDF_list)):
    ratio_Algo = otexp.RatioOfUniforms(logPDF_list[k], range_list[k], isScaled)
    for i in range(numberR + 1):
        r = a ** (i - 0.5 * numberR)
        ratio_Algo.setR(r)
        accept_Ratio = ratio_Algo.getAcceptanceRatio()
        data_r_ratio[i] = [r, accept_Ratio]
    draw = ot.Curve(data_r_ratio)
    draw.setLegend(name_list[k])
    g.add(draw)

isScaled = False
ratio_Algo = otexp.RatioOfUniforms(log_UnscaledPDF, range_PDF, isScaled)
for i in range(numberR + 1):
    r = a ** (i - 0.5 * numberR)
    ratio_Algo.setR(r)
    accept_Ratio = ratio_Algo.getAcceptanceRatio()
    data_r_ratio[i] = [r, accept_Ratio]
draw = ot.Curve(data_r_ratio)
draw.setLegend(
    r"$f(x) = \cos(x)e^x\mathbf{1}_{\left[-\frac{\pi}{2}, \frac{\pi}{2}\right]}(x)$"
)
g.add(draw)

g.setLegendPosition("upper left")
g.setLegendCorner([1.0, 1.0])
view = otv.View(g)

# %%
# Exploration of the optimization problems.
# -----------------------------------------
# We illustrate how to explore the results of the optimization problems that have been solved to define
# :math:`\tilde{A}_{f,r}`.
# We consider the example defined by the function:
#
# .. math::
#
#     f(x) = \frac{1}{2} (2 + \sin(x)^2) \exp \left[- \left(2 + \cos(3x)^3 + \sin(2x)^3 \right) x
#     \right]  \mathbf{1}_{[0, 2 \pi]}(x).
#
# The method *initialize* returns all the starting points that have been used to compute respectively :math:`u_{sup}`,
# :math:`v_{inf}` and :math:`v_{sup}`.
# The maximum number of starting points is defined in :class:`~openturns.ResourceMap`, entry
# *RatioOfUniformas-MaximumMultiStart*. These
# points are selected scrolling through  the Sobol low discrepance sequence in dimension :math:`\inputDim =  1`
# and which maximum size is defined in :class:`~openturns.ResourceMap`, entry *RatioOfUniforms-CandidateNumber*.
f = ot.SymbolicFunction(
    "x", "0.5 * (2 + sin(x)^2) * exp( -( 2 + cos(3*x)^3 + sin(2*x)^3) * x )"
)
log_UnscaledPDF = ot.ComposedFunction(ot.SymbolicFunction("x", "log(x)"), f)
lower_bound = 0.0
upper_bound = 2.0 * pi
range_PDF = ot.Interval(lower_bound, upper_bound)

graph = f.draw(lower_bound, upper_bound, 512)
graph.setTitle(
    r"Christian Robert function: $f(x) =  0.5(2 + \sin^2 x) e^{ -x( 2 + \cos^3 3x + \sin^3 2x)}, x \in [0, 2\pi]$"
)
graph.setXTitle(r"$x$")
graph.setYTitle(r"$f(x)$")
view = otv.View(graph)

ratio_Algo = otexp.RatioOfUniforms(log_UnscaledPDF, range_PDF)

# %%
# In order to get all the intermediate results, we change the :class:`~openturns.ResourceMap`,
# entry *MultiStart-KeepResults*.
#
# Furthermore, as :math:`I = [0, 2 \pi]=  [a_i, b_i]`, then :math:`a_i \geq 0` implies that :math:`v_{inf} = 0`.
# There are only the bounds :math:`u_{sup}` and :math:`v_{sup}` to find: the collection of optimization problems
# only contains 2 items.
ot.ResourceMap.SetAsBool("MultiStart-KeepResults", True)
coll_Multi_Start = ratio_Algo.initialize()
start_Sample_SupU = coll_Multi_Start[0].getStartingSample()
start_Sample_SupV = coll_Multi_Start[1].getStartingSample()

# %%
# We plot on the previous graph all the starting points (red) used to find :math:`u_{sup}` and the associated maxima found by the algorithm. We
# also link each starting point to its associated optimal value (green).
Cloud_start_Usup = ot.Cloud(
    start_Sample_SupU, ot.Sample(start_Sample_SupU.getSize(), [0.0])
)
Cloud_start_Usup.setColor("red")
Cloud_start_Usup.setPointStyle("fsquare")
graph.add(Cloud_start_Usup)

coll_results = coll_Multi_Start[0].getResultCollection()
sample_optima = ot.Sample(0, 2)
for i in range(coll_results.getSize()):
    opt_x = coll_results[i].getOptimalPoint()[0]
    opt_value = f([opt_x])[0]
    sample_optima.add([opt_x, opt_value])
    line_convergence = ot.Curve([[start_Sample_SupU[i, 0], 0.0], [opt_x, opt_value]])
    line_convergence.setLineStyle("dashed")
    line_convergence.setColor("black")
    graph.add(line_convergence)
cloud_optima = ot.Cloud(sample_optima)
cloud_optima.setColor("green")
cloud_optima.setPointStyle("star")
graph.add(cloud_optima)
graph.setYTitle("")
view = otv.View(graph)


# %%
view.ShowAll()
