"""

Rectangular Refined Stratified Sampling - Random Refinement
============================================================

In this example, the stratified sampling method is employed to generate samples from an exponential distribution using
Voronoi stratification.
"""

# %% md
#
# In this example, Stratified sampling is used to generate samples from Uniform probability distribution and sample are
# added using adaptive approach Refined Stratified Sampling.
#
# Import the necessary libraries. Here we import standard libraries such as numpy and matplotlib, but also need to
# import the :class:`.TrueStratifiedSampling` and :class:`.RefinedStratifiedSampling` class from :py:mod:`UQpy`.

# %%

from UQpy.sampling.stratified_sampling.strata import RectangularStrata
from UQpy.sampling import TrueStratifiedSampling, RefinedStratifiedSampling, RandomRefinement
from UQpy.distributions import Uniform
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np

# %% md
#
# Create a distribution object.

# %%

marginals = [Uniform(loc=0., scale=1.), Uniform(loc=0., scale=1.)]

# %% md
#
# Create a strata object

# %%

strata = RectangularStrata(strata_number=[4, 4])

# %% md
#
# Run stratified sampling.

# %%

x = TrueStratifiedSampling(distributions=marginals, strata_object=strata, nsamples_per_stratum=1, random_state=1)

# %% md
#
# Using UQpy :class:`.RefinedStratifiedSampling` class to expand samples generated by :class:`.TrueStratifiedSampling`
# class. In this example, two new samples are generated using rectangular stratification .

# %%

refinement = RandomRefinement(strata=strata)
y = RefinedStratifiedSampling(stratified_sampling=x, refinement_algorithm=refinement,
                              nsamples=18, samples_per_iteration=2, random_state=2)

# %% md
#
# In the first figure shown below, samples generated from :class:`.TrueStratifiedSampling` class are plotted. Second
# figure, shows the new stratification corresponding to two generated samples using :class:`.RefinedStratifiedSampling`
# class.

# %%

fig1 = strata.plot_2d()
plt.plot(x.samples[:16, 0], x.samples[:16, 1], 'ro')
plt.xlim([0, 1])
plt.ylim([0, 1])
plt.show()

fig2 = strata.plot_2d()
plt.xlim([0, 1])
plt.ylim([0, 1])
plt.plot(y.samples[:16, 0], y.samples[:16, 1], 'ro')
plt.plot(y.samples[16:, 0], y.samples[16:, 1], 'yo')
plt.show()

# %% md
#
# Further, :class:`.RefinedStratifiedSampling` class is used to adaptively increase the sample size. In this example,
# samples are added in strata with maximum weights associated with it and strata's are cut randomly along the maximum
# width.

# %%

y.run(nsamples=100)

# %% md
#
# In the figure shown below, all samples generated from :class:`.TrueStratifiedSampling` and
# :class:`.RefinedStratifiedSampling` class are plotted.

# %%

fig3 = strata.plot_2d()
plt.xlim([0, 1])
plt.ylim([0, 1])
plt.plot(y.samplesU01[:16, 0], y.samplesU01[:16, 1], 'ro')
plt.plot(y.samplesU01[16:, 0], y.samplesU01[16:, 1], 'go')
plt.show()