"""Example showing how to use the ProblemSpec approach.

Showcases method chaining, and parallel model runs using 2 processors.

The following caveats apply:

1. Functions passed into `sample`, `analyze` and `evaluate` must
   accept a numpy array of `X` values as the first parameter, and return a
   numpy array of results.
2. Parallel evaluation/analysis is only beneficial for long-running models
   or large datasets
3. Currently, results must fit in memory - no on-disk caching is provided.
"""

from SALib.analyze import sobol
from SALib.sample import saltelli
from SALib.test_functions import Ishigami
from SALib import ProblemSpec
import numpy as np

import time


if __name__ == "__main__":
    # Create the SALib Problem specification
    sp = ProblemSpec(
        {
            "names": ["x1", "x2", "x3"],
            "groups": None,
            "bounds": [[-np.pi, np.pi]] * 3,
            "outputs": ["Y"],
        }
    )

    # Single core example
    start = time.perf_counter()
    (
        sp.sample_saltelli(2**15)
        .evaluate(Ishigami.evaluate)
        .analyze_sobol(calc_second_order=True, conf_level=0.95)
    )
    print("Time taken with 1 core:", time.perf_counter() - start, "\n")

    # Same above example, but passing in specific functions
    # (sp.sample(saltelli.sample, 25000, calc_second_order=True)
    #     .evaluate(Ishigami.evaluate)
    #     .analyze(sobol.analyze, calc_second_order=True, conf_level=0.95))

    # Samples, model results and analyses can be extracted:
    # print(sp.samples)
    # print(sp.results)
    # print(sp.analysis)
    # print(sp.to_df())

    # Can set pre-existing samples/results as needed
    # sp.samples = some_numpy_array
    # sp.set_samples(some_numpy_array)
    #
    # Using method chaining...
    # (sp.set_samples(some_numpy_array)
    #    .set_results(some_result_array)
    #    .analyze_sobol(calc_second_order=True, conf_level=0.95))

    # Parallel example
    start = time.perf_counter()
    (
        sp.sample(saltelli.sample, 2**15)
        # can specify number of processors to use with `nprocs`
        # this will be capped to the number of detected processors
        # or, in the case of analysis, the number of outputs.
        .evaluate(Ishigami.evaluate, nprocs=2).analyze_sobol(
            calc_second_order=True, conf_level=0.95, nprocs=2
        )
    )
    print("Time taken with 2 cores:", time.perf_counter() - start, "\n")

    print(sp)

    # Distributed example
    # This example requires an additional dependency (Pathos) to be installed
    # `pip install pathos` or `conda install pathos -c conda-forge`

    # Specify itself as servers as an example
    servers = ("localhost:55774", "localhost:55775", "localhost:55776")

    start = time.perf_counter()
    (
        sp.sample(saltelli.sample, 2**15)
        .evaluate_distributed(
            Ishigami.evaluate, nprocs=2, servers=servers, verbose=True
        )
        .analyze(sobol.analyze, calc_second_order=True, conf_level=0.95)
    )
    print("Time taken with distributed cores:", time.perf_counter() - start, "\n")

    print(sp)

    # To display plots:
    # import matplotlib.pyplot as plt
    # sp.plot()
    # plt.show()

# First-order indices expected with Saltelli sampling:
# x1: 0.3139
# x2: 0.4424
# x3: 0.0
