#!/usr/bin/env python
# coding: utf-8

import numpy as np
import scipy.stats as stats

from wellmet.samplebox import SampleBox
from wellmet import f_models


# 1. Set up probability distribution
# Standard Gaussian variables, 2D
#f = f_models.SNorm(2)
# Just normal variables
f = f_models.Norm(mean=[-1, 0], std=[2, 1])
# Independent non-Gaussian variables
#f = f_models.UnCorD((stats.gumbel_r, stats.uniform))
# Correlated:
# Nataf model with correlations of the respective _Gaussian_ marginals
#f = f_models.Nataf((stats.gumbel_r, stats.weibull_min(c=1.5)), [[1,0.8], [0.8,1]])

# 2. Define LSF function
def my_problem(input_sample):
    # get real (physical) space coordinates
    # X is a numpy array with shape (nsim, ndim)
    # the algorithm normally sends (1, ndim) sample
    X = input_sample.R
    # LSF
    g = X[:, 0] - X[:, 1] + 3
    # we should return an instance of SampleBox class
    # it stores coordinates along with LSF calculation result
    # with kind of signature
    return SampleBox(input_sample, g, "my_problem")






# 3. Prepare storage
# no need to store anything
#sample_box = SampleBox(f)

# keep samples and estimations continiously stored 
from wellmet import reader
sample_box = reader.Reader("meow_problem", f)

# 4. Setup the algorithm
from wellmet.dicebox.circumtri import CirQTri
import quadpy

scheme = quadpy.tn.stroud_tn_3_6b(sample_box.nvar)
convex_hull_degree = 5 # degreee of Grundmann-Möller cubature scheme
q = 1 # should be > 0. Greater values slightly enforces exploration
screening_rate = 0 # 10 means to sacrifice every tenth sample for screening
box = CirQTri(sample_box, scheme, convex_hull_degree, q, screening_rate)


# 5. Here we go!
for i in range(20):
    # ask where to sample the next point
    # next_node is an f_model instance
    next_node = box()
    # call LSF
    new_sample = my_problem(next_node)
    # put calculation result to the box
    box.add_sample(new_sample)
    
print(box.get_pf_estimation())
sensitivities_results = box.Tri.perform_sensitivity_analysis()
print(sensitivities_results.sensitivities)