import numpy as np
import bet.calculateP.simpleFunP as simpleFunP
import bet.calculateP.calculateP as calculateP
import bet.sample as samp
import bet.sampling.basicSampling as bsam
from myModel import my_model
from IPython.display import Image

"""
Explanation: <center> Validation Example: Linear Map
Copyright (C) 2014-2019 The BET Development Team
This 2D linear example verifies that geometrically distinct QoI can
recreate a probability measure on the input parameter space
used to define the output probability measure.
End of explanation
"""


sampler = bsam.sampler(my_model)

# Initialize 3-dimensional input parameter sample set object
input_samples = samp.sample_set(2)

# Set parameter domain
input_samples.set_domain(np.repeat([[0.0, 1.0]], 2, axis=0))

"""
Explanation: Characterize Parameter Space
Define the sampler that will be used to create the discretization
object, which is the fundamental object used by BET to compute
solutions to the stochastic inverse problem.
The sampler and my_model is the interface of BET to the model,
and it allows BET to create input/output samples of the model.
End of explanation
"""


# Generate samples on the parameter space
randomSampling = True
if randomSampling is True:
    input_samples = sampler.random_sample_set('random', input_samples, num_samples=1E3)
else:
    input_samples = sampler.regular_sample_set(input_samples, num_samples_per_dim=[30, 30])

    

"""
Explanation: Suggested Changes
Try with and without random sampling.
If using random sampling, try num_samples = 1E3 and 1E4.
What happens when num_samples = 1E2?
Try using 'lhs' instead of 'random' in the random_sample_set.
If using regular sampling, try different numbers of samples
per dimension.
End of explanation
"""


MC_assumption = True

# Estimate volumes of Voronoi cells associated with the parameter samples
if MC_assumption is False:
    input_samples.estimate_volume(n_mc_points=1E5)
else:
    input_samples.estimate_volume_mc()

# Create the discretization object using the input samples
my_discretization = sampler.compute_QoI_and_create_discretization(input_samples,
                                               savefile = 'Validation_discretization.txt.gz')

"""
Explanation: Characterize Data Space
Compute the output distribution simple function approximation by
propagating a different set of samples to implicitly define a Voronoi
discretization of the data space, corresponding to an implicitly defined
set of contour events defining a discretization of the input parameter
space. 
The probabilities of the Voronoi cells in the data space (and
thus the probabilities of the corresponding contour events in the
input parameter space) are determined by Monte Carlo sampling using
a set of i.i.d. uniform samples to bin into these cells.
Suggested Changes
A standard Monte Carlo (MC) assumption is that every Voronoi cell
has the same volume. If a regular grid of samples was used, then
the standard MC assumption is true.
See what happens if the MC assumption is not assumed to be true, and
if different numbers of points are used to estimate the volumes of
the Voronoi cells.
End of explanation
"""


num_samples_discretize_D = 1
num_iid_samples = 1E5

Partition_set = samp.sample_set(2)
Monte_Carlo_set = samp.sample_set(2)

Partition_set.set_domain(np.repeat([[0.0, 1.0]], 2, axis=0))
Monte_Carlo_set.set_domain(np.repeat([[0.0, 1.0]], 2, axis=0))

Partition_discretization = sampler.create_random_discretization('random',
                                                            Partition_set,
                                                            num_samples=num_samples_discretize_D)

Monte_Carlo_discretization = sampler.create_random_discretization('random',
                                                            Monte_Carlo_set,
                                                            num_samples=num_iid_samples)

# Compute the simple function approximation to the distribution on the data space
simpleFunP.user_partition_user_distribution(my_discretization,
                                            Partition_discretization,
                                            Monte_Carlo_discretization)

# Calculate probabilities
calculateP.prob(my_discretization)

"""
Explanation: Solve Problem
Suggested Changes
See the effect of using different values for 
num_samples_discretize_D.
Choosing num_samples_discretize_D = 1 produces exactly the right answer and is equivalent to assigning a
uniform probability to each data sample above (why?).
Try setting this to 2, 5, 10, 50, and 100. Can you explain what you
are seeing? To see an exaggerated effect, try using random sampling
above with n_samples set to 1E2.
End of explanation
"""


%store my_discretization

"""
Explanation: Store Data for Retrieval in other Jupyter Notebooks
End of explanation
"""


