//                                               -*- C++ -*-
/**
 *  @brief The test file of FunctionalChaosAlgoritm class
 *
 *  Copyright 2005-2025 Airbus-EDF-IMACS-ONERA-Phimeca
 *
 *  This library is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This library is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with this library.  If not, see <http://www.gnu.org/licenses/>.
 *
 */
#include "openturns/OT.hxx"
#include "openturns/IshigamiUseCase.hxx"
#include "openturns/OTtestcode.hxx"

using namespace OT;
using namespace OT::Test;

// Simultaneously sort the nodes and weights
void printSobolResult(Scalar S_computed, Scalar S_exact)
{
  OStream fullprint(std::cout);
  fullprint << "   S (PCE) " << std::fixed << std::setprecision(4) << S_computed << std::endl;
  fullprint << "   S (exact) " << std::fixed << std::setprecision(4) << S_exact << std::endl;
  Scalar absoluteError = abs(S_computed - S_exact);
  fullprint << "   Abs. Error " << std::scientific << std::setprecision(4) << absoluteError << std::endl;
}

int main(int, char *[])
{
  TESTPREAMBLE;
  OStream fullprint(std::cout);

  try
  {

    // Problem parameters
    IshigamiUseCase ishigami;
    UnsignedInteger dimension = 3;
    // Reference analytical values
    Scalar mean = ishigami.getMean();
    Scalar variance = ishigami.getVariance();
    Scalar S1 = ishigami.getFirstOrderSobolIndices()[0];
    Scalar S2 = ishigami.getFirstOrderSobolIndices()[1];
    Scalar S3 = ishigami.getFirstOrderSobolIndices()[2];
    Scalar S12 = ishigami.getFirstOrderInteractionSobolIndex({0, 1});
    Scalar S13 = ishigami.getFirstOrderInteractionSobolIndex({0, 2});
    Scalar S23 = ishigami.getFirstOrderInteractionSobolIndex({1, 2});
    Scalar ST1 = S1 + S13;
    Scalar ST2 = S2;
    Scalar ST3 = S3 + S13;
    Scalar S123 = ishigami.getFirstOrderInteractionSobolIndex({0, 1, 2});
    fullprint << "mean = " << mean << std::endl;
    fullprint << "variance = " << std::fixed << std::setprecision(4) << variance << std::endl;
    fullprint << "S1 = " << std::fixed << std::setprecision(4) << S1 << std::endl;
    fullprint << "S2 = " << std::fixed << std::setprecision(4) << S2 << std::endl;
    fullprint << "S3 = " << std::fixed << std::setprecision(4) << S3 << std::endl;
    fullprint << "S13 = " << std::fixed << std::setprecision(4) << S13 << std::endl;
    fullprint << "ST1 = " << std::fixed << std::setprecision(4) << ST1 << std::endl;
    fullprint << "ST2 = " << std::fixed << std::setprecision(4) << ST2 << std::endl;
    fullprint << "ST3 = " << std::fixed << std::setprecision(4) << ST3 << std::endl;

    Function model(ishigami.getModel());  // Create the Ishigami function
    JointDistribution distribution(ishigami.getInputDistribution());  // Create the input distribution

    // Create the orthogonal basis
    Collection<OrthogonalUniVariatePolynomialFamily> polynomialCollection(dimension);
    polynomialCollection[0] = LegendreFactory();
    polynomialCollection[1] = LegendreFactory();
    polynomialCollection[2] = LegendreFactory();

    HyperbolicAnisotropicEnumerateFunction enumerateFunction( dimension,  0.6);
    OrthogonalProductPolynomialFactory productBasis(polynomialCollection, enumerateFunction);

    UnsignedInteger size = pow(2, 10);
    fullprint << "size = " << size << std::endl;
    WeightedExperiment experiment(LowDiscrepancyExperiment(SobolSequence(), distribution, size));
    // Create the adaptive strategy
    // We can choose amongst several strategies
    // First, the most efficient (but more complex!) strategy
    UnsignedInteger degree = 14;
    UnsignedInteger basisSize = enumerateFunction.getBasisSizeFromTotalDegree(degree);
    fullprint << "basisSize = " << basisSize << std::endl;
    FixedStrategy adaptiveStrategy( productBasis, basisSize );
    FittingAlgorithm fittingAlgorithm = CorrectedLeaveOneOut();
    const LeastSquaresStrategy projectionStrategy(LeastSquaresMetaModelSelectionFactory(LARS(), fittingAlgorithm));
    const Sample X(experiment.generate());
    const Sample Y(model(X));
    fullprint << "Create object" << std::endl;
    FunctionalChaosAlgorithm algo(X, Y, distribution, adaptiveStrategy, projectionStrategy);
    fullprint << "Run()" << std::endl;
    algo.run();
    fullprint << "GetResult()" << std::endl;
    FunctionalChaosResult result = algo.getResult();
    FunctionalChaosSobolIndices sensitivity(result);
    fullprint << sensitivity.__str__() << std::endl;
    fullprint << sensitivity.__repr_markdown__() << std::endl;
    //
    const Scalar rtol = 0.0;
    const Scalar atol = 0.001;
    Scalar S_computed;
    Scalar S_exact;
    //
    fullprint << "Test first order Sobol' indices" << std::endl;
    fullprint << "First order, X1" << std::endl;
    S_computed = sensitivity.getSobolIndex(0);
    printSobolResult(S_computed, S1);
    assert_almost_equal(S_computed, S1, rtol, atol);
    //
    fullprint << "First order, X2" << std::endl;
    S_computed = sensitivity.getSobolIndex(1);
    printSobolResult(S_computed, S2);
    assert_almost_equal(S_computed, S2, rtol, atol);
    //
    fullprint << "First order, X3" << std::endl;
    S_computed = sensitivity.getSobolIndex(2);
    printSobolResult(S_computed, S3);
    assert_almost_equal(S_computed, S3, rtol, atol);
    //
    fullprint << "Test total order Sobol' indices" << std::endl;
    fullprint << "Total, X1" << std::endl;
    S_computed = sensitivity.getSobolTotalIndex(0);
    printSobolResult(S_computed, ST1);
    assert_almost_equal(S_computed, ST1, rtol, atol);
    //
    fullprint << "Total, X2" << std::endl;
    S_computed = sensitivity.getSobolTotalIndex(1);
    printSobolResult(S_computed, ST2);
    assert_almost_equal(S_computed, ST2, rtol, atol);
    //
    fullprint << "Total, X3" << std::endl;
    S_computed = sensitivity.getSobolTotalIndex(2);
    printSobolResult(S_computed, ST3);
    assert_almost_equal(S_computed, ST3, rtol, atol);
    //
    fullprint << "Test first order (closed) group Sobol' indices" << std::endl;
    fullprint << "X1" << std::endl;
    S_computed = sensitivity.getSobolGroupedIndex({0});
    printSobolResult(S_computed, S1);
    assert_almost_equal(S_computed, S1, rtol, atol);
    //
    fullprint << "X2" << std::endl;
    S_computed = sensitivity.getSobolGroupedIndex({1});
    printSobolResult(S_computed, S2);
    assert_almost_equal(S_computed, S2, rtol, atol);
    //
    fullprint << "X3" << std::endl;
    S_computed = sensitivity.getSobolGroupedIndex({2});
    printSobolResult(S_computed, S3);
    assert_almost_equal(S_computed, S3, rtol, atol);
    //
    fullprint << "(X1, X2)" << std::endl;
    S_computed = sensitivity.getSobolGroupedIndex({0, 1});
    S_exact = S1 + S2 + S12;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "(X1, X3)" << std::endl;
    S_computed = sensitivity.getSobolGroupedIndex({0, 2});
    S_exact = S1 + S3 + S13;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "(X2, X3)" << std::endl;
    S_computed = sensitivity.getSobolGroupedIndex({1, 2});
    S_exact = S2 + S3 + S23;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "(X1, X2, X3)" << std::endl;
    S_computed = sensitivity.getSobolGroupedIndex({0, 1, 2});
    S_exact = 1.0;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "Test total group Sobol' indices" << std::endl;
    fullprint << "X1" << std::endl;
    S_computed = sensitivity.getSobolGroupedTotalIndex({0});
    printSobolResult(S_computed, ST1);
    assert_almost_equal(S_computed, ST1, rtol, atol);
    //
    fullprint << "X2" << std::endl;
    S_computed = sensitivity.getSobolGroupedTotalIndex({1});
    printSobolResult(S_computed, ST2);
    assert_almost_equal(S_computed, ST2, rtol, atol);
    //
    fullprint << "X3" << std::endl;
    S_computed = sensitivity.getSobolGroupedTotalIndex({2});
    printSobolResult(S_computed, ST3);
    assert_almost_equal(S_computed, ST3, rtol, atol);
    //
    fullprint << "(X1, X2)" << std::endl;
    S_computed = sensitivity.getSobolGroupedTotalIndex({0, 1});
    S_exact = S1 + S2 + S12 + S13 + S123;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "(X1, X3)" << std::endl;
    S_computed = sensitivity.getSobolGroupedTotalIndex({0, 2});
    S_exact = S1 + S3 + S13 + S23 + S123;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "(X2, X3)" << std::endl;
    S_computed = sensitivity.getSobolGroupedTotalIndex({1, 2});
    S_exact = S2 + S3 + S13 + S23 + S123;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "(X1, X2, X3)" << std::endl;
    S_computed = sensitivity.getSobolGroupedTotalIndex({0, 1, 2});
    S_exact = 1.0;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "Test interaction group Sobol' indices" << std::endl;
    fullprint << "X1" << std::endl;
    Indices indices0 = {0};
    S_computed = sensitivity.getSobolIndex(indices0);
    printSobolResult(S_computed, S1);
    assert_almost_equal(S_computed, S1, rtol, atol);
    //
    fullprint << "X2" << std::endl;
    Indices indices1 = {1};
    S_computed = sensitivity.getSobolIndex(indices1);
    printSobolResult(S_computed, S2);
    assert_almost_equal(S_computed, S2, rtol, atol);
    //
    fullprint << "X3" << std::endl;
    Indices indices2 = {2};
    S_computed = sensitivity.getSobolIndex(indices2);
    printSobolResult(S_computed, S3);
    assert_almost_equal(S_computed, S3, rtol, atol);
    //
    fullprint << "X1, X2" << std::endl;
    S_computed = sensitivity.getSobolIndex({0, 1});
    printSobolResult(S_computed, S12);
    assert_almost_equal(S_computed, S12, rtol, atol);
    //
    fullprint << "X1, X3" << std::endl;
    S_computed = sensitivity.getSobolIndex({0, 2});
    printSobolResult(S_computed, S13);
    assert_almost_equal(S_computed, S13, rtol, atol);
    //
    fullprint << "X2, X3" << std::endl;
    S_computed = sensitivity.getSobolIndex({1, 2});
    printSobolResult(S_computed, S23);
    assert_almost_equal(S_computed, S23, rtol, atol);
    //
    fullprint << "X1, X2, X3" << std::endl;
    S_computed = sensitivity.getSobolIndex({0, 1, 2});
    printSobolResult(S_computed, S123);
    assert_almost_equal(S_computed, S123, rtol, atol);
    //
    fullprint << "Test total interaction group Sobol' indices" << std::endl;
    fullprint << "X1" << std::endl;
    S_computed = sensitivity.getSobolTotalIndex(indices0);
    S_exact = S1 + S12 + S13 + S123;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "X2" << std::endl;
    S_computed = sensitivity.getSobolTotalIndex(indices1);
    S_exact = S2 + S12 + S23 + S123;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "X3" << std::endl;
    S_computed = sensitivity.getSobolTotalIndex(indices2);
    S_exact = S3 + S13 + S23 + S123;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "X1, X2" << std::endl;
    S_computed = sensitivity.getSobolTotalIndex({0, 1});
    S_exact = S12 + S123;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "X1, X3" << std::endl;
    S_computed = sensitivity.getSobolTotalIndex({0, 2});
    S_exact = S13 + S123;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "X2, X3" << std::endl;
    S_computed = sensitivity.getSobolTotalIndex({1, 2});
    S_exact = S23 + S123;
    printSobolResult(S_computed, S_exact);
    assert_almost_equal(S_computed, S_exact, rtol, atol);
    //
    fullprint << "X1, X2, X3" << std::endl;
    S_computed = sensitivity.getSobolTotalIndex({0, 1, 2});
    printSobolResult(S_computed, S123);
    assert_almost_equal(S_computed, S123, rtol, atol);
  }
  catch (TestFailed & ex)
  {
    std::cerr << ex << std::endl;
    return ExitCode::Error;
  }


  return ExitCode::Success;
}
