/*
 * Copyright (c) 2022, Miroslav Stoyanov & Weiwei Kong
 *
 * This file is part of
 * Toolkit for Adaptive Stochastic Modeling And Non-Intrusive ApproximatioN: TASMANIAN
 *
 * Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following
 * conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions
 *    and the following disclaimer in the documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse
 *    or promote products derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
 * OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
 * OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 * UT-BATTELLE, LLC AND THE UNITED STATES GOVERNMENT MAKE NO REPRESENTATIONS AND DISCLAIM ALL WARRANTIES, BOTH EXPRESSED AND
 * IMPLIED. THERE ARE NO EXPRESS OR IMPLIED WARRANTIES OF MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE, OR THAT THE USE OF
 * THE SOFTWARE WILL NOT INFRINGE ANY PATENT, COPYRIGHT, TRADEMARK, OR OTHER PROPRIETARY RIGHTS, OR THAT THE SOFTWARE WILL
 * ACCOMPLISH THE INTENDED RESULTS OR THAT THE SOFTWARE OR ITS USE WILL NOT RESULT IN INJURY OR DAMAGE. THE USER ASSUMES
 * RESPONSIBILITY FOR ALL LIABILITIES, PENALTIES, FINES, CLAIMS, CAUSES OF ACTION, AND COSTS AND EXPENSES, CAUSED BY, RESULTING
 * FROM OR ARISING OUT OF, IN WHOLE OR IN PART THE USE, STORAGE OR DISPOSAL OF THE SOFTWARE.
 */

#ifndef __TASMANIAN_TASDREAM_OPTIMIZATION_TESTS_CPP
#define __TASMANIAN_TASDREAM_OPTIMIZATION_TESTS_CPP

#include "TasmanianOptimization.hpp"
#include "tasdreamExternalTests.hpp"

struct test_result{
    void tassert(bool result){
        pass_last = result;
        pass_all = pass_all and pass_last;
    }
    bool failed_last() const{ return not pass_last; }
    bool failed_any() const{ return not pass_all; }
    bool success() const{ return pass_all; }
    bool last_success() const{ return pass_last; }
private:
    bool pass_all = true;
    bool pass_last = true;
};

namespace TasOptimization {

// Unit tests for TasOptimization::ParticleSwarmState.
bool testParticleSwarmState(bool verbose) {
    test_result test;

    // Check size of the accessible vectors generated by different constructors.
    int num_dimensions = 2;
    int num_particles = 15;
    std::vector<double> dummy_positions(num_dimensions * num_particles);
    std::vector<double> dummy_velocities(num_dimensions * num_particles);
    std::vector<ParticleSwarmState> states = {
        ParticleSwarmState(num_dimensions, num_particles),
        ParticleSwarmState(num_dimensions, std::move(dummy_positions), std::move(dummy_velocities))
    };
    for (int i=0; i<2; i++) {
        test.tassert(states[i].getParticlePositions().size() == Utils::size_mult(num_dimensions, num_particles));
        test.tassert(states[i].getParticleVelocities().size() == Utils::size_mult(num_dimensions, num_particles));
        test.tassert(states[i].getBestParticlePositions().size() == Utils::size_mult(num_dimensions, num_particles+1));
        test.tassert(states[i].getBestPosition().size() == (size_t) num_dimensions);
        std::vector<bool> init_vector = states[i].getStateVector();
        test.tassert(i == 0 ? not init_vector[0] : init_vector[0]);
        test.tassert(i == 0 ? not init_vector[1] : init_vector[1]);
        test.tassert(not init_vector[2]);
        test.tassert(not init_vector[3]);
    }

    // Check TasOptimization::ParticleSwarmState::initializeParticlesInsideBox().
    std::vector<double> lower = {-1.0, 1.0};
    std::vector<double> upper = {2.0, 3.0};
    std::minstd_rand park_miller(42);
    std::uniform_real_distribution<double> unif(0.0, 1.0);
    auto get_rand = [&]()->double{ return unif(park_miller); };
    states[0].initializeParticlesInsideBox(lower, upper, get_rand);
    std::vector<double> positions = states[0].getParticlePositions();
    std::vector<double> velocities = states[0].getParticleVelocities();
    for (int i=0; i<num_particles * num_dimensions; i++) {
        test.tassert(positions[i] >= lower[i % num_dimensions] - TasGrid::Maths::num_tol);
        test.tassert(positions[i] <= upper[i % num_dimensions] + TasGrid::Maths::num_tol);
        double range = fabs(upper[i % num_dimensions] - lower[i % num_dimensions]);
        test.tassert(velocities[i] >= -range - TasGrid::Maths::num_tol);
        test.tassert(velocities[i] <=  range + TasGrid::Maths::num_tol);
    }
    std::vector<bool> init_vector = states[0].getStateVector();
    test.tassert(init_vector[0] and init_vector[1]);

    // Check the nontrivial setters.
    ParticleSwarmState state(num_dimensions, num_particles);
    std::vector<double> ones1(num_dimensions * num_particles, 1);
    std::vector<double> ones2(num_dimensions * (num_particles + 1), 1);
    state.setParticlePositions(ones1);
    for (auto p : state.getParticlePositions()) test.tassert(p == 1);
    state.setParticleVelocities(ones1);
    for (auto v : state.getParticleVelocities()) test.tassert(v == 1);
    state.setBestParticlePositions(ones2);
    for (auto bp : state.getBestParticlePositions()) test.tassert(bp == 1);
    init_vector = state.getStateVector();
    test.tassert(init_vector[0] and init_vector[1] and init_vector[2]);

    // Check TasOptimization::ParticleSwarmState::clearBestParticles().
    state.clearBestParticles();
    for (auto bp : states[0].getBestParticlePositions()) test.tassert(bp == 0);
    init_vector = state.getStateVector();
    test.tassert(not init_vector[2]);

    // Reporting.
    if (test.failed_any() or verbose) reportPassFail(test.success(), "Particle Swarm", "State Unit Tests");
    return test.success();
}

// Unit tests for TasOptimization::ParticleSwarm on a single objective function.
bool testParticleSwarmSingle(ObjectiveFunction f, ParticleSwarmState state, TasDREAM::DreamDomain inside, int iterations,
                             double optimal_val) {
    test_result test;

    // Run the particle swarm algorithm.
    std::minstd_rand park_miller(42);
    std::uniform_real_distribution<double> unif(0.0, 1.0);
    auto get_rand = [&]()->double{ return unif(park_miller); };
    ParticleSwarm(f, inside, 0.5, 2, 2, iterations, state, get_rand);

    // Check optimality and state changes of the run.
    std::vector<double> best_swarm_point = state.getBestPosition();
    std::vector<double> best_swarm_value_vec(1);
    f(best_swarm_point, best_swarm_value_vec);
    test.tassert(std::fabs(best_swarm_value_vec[0] - optimal_val) <= 10 * TasGrid::Maths::num_tol);
    std::vector<bool> init_vector = state.getStateVector();
    test.tassert(init_vector[3]);

    // Make sure subsequent runs do not make any strange modifications.
    ParticleSwarm(f, inside, 0.5, 2, 2, 1, state, get_rand);
    f(best_swarm_point, best_swarm_value_vec);
    test.tassert(std::fabs(best_swarm_value_vec[0] - optimal_val) <= 10 * TasGrid::Maths::num_tol);
    init_vector = state.getStateVector();
    test.tassert(init_vector[3]);

    // TasOptimization::ParticleSwarmState::clearCache().
    state.clearCache();
    init_vector = state.getStateVector();
    test.tassert(init_vector[0] and init_vector[1] and init_vector[2] and not init_vector[3]);

    return test.success();
}

// Unit tests for TasOptimization::ParticleSwarm on multiple objective functions.
bool testParticleSwarm(bool verbose) {
    test_result test;

    // l1 norm over the domain [-5, 2] ^ 6.
    int num_dimensions = 6;
    int num_particles = 100;
    int iterations = 300;
    std::vector<double> lower(num_dimensions, -5.0);
    std::vector<double> upper(num_dimensions, 2.0);
    TasOptimization::ObjectiveFunctionSingle l1_single =
            [](const std::vector<double> &x)->double {
                double sum = 0;
                for (auto xi : x) sum += std::fabs(xi);
                return sum;
            };
    TasOptimization::ObjectiveFunction l1 = TasOptimization::makeObjectiveFunction(num_dimensions, l1_single);
    TasOptimization::ParticleSwarmState state(num_dimensions, num_particles);
    state.initializeParticlesInsideBox(lower, upper);
    test.tassert( testParticleSwarmSingle(l1, state, TasDREAM::hypercube(lower, upper), iterations, 0) );
    if (test.failed_last()) std::cout << "ERROR: failed l1 example 1 for Particle Swarm optimization.\n";

    // Six hump-camel function over the domain [-3, 3] x [-2, 2].
    num_dimensions = 2;
    num_particles = 50;
    iterations = 100;
    lower = {-3.0, -2.0};
    upper = {3.0, 2.0};
    TasOptimization::ObjectiveFunctionSingle shc_single =
            [](const std::vector<double> &x)->double {
                return (4.0 - 2.1 * x[0]*x[0] + x[0]*x[0]*x[0]*x[0] / 3.0) * x[0]*x[0] +
                        x[0] * x[1] +
                        (-4.0 + 4.0 * x[1]*x[1]) * x[1]*x[1];};
    TasOptimization::ObjectiveFunction shc = TasOptimization::makeObjectiveFunction(num_dimensions, shc_single);
    state = ParticleSwarmState(num_dimensions, num_particles);
    state.initializeParticlesInsideBox(lower, upper);
    test.tassert( testParticleSwarmSingle(shc, state, TasDREAM::hypercube(lower, upper), iterations, -1.031628453489877) );
    if (test.failed_last()) std::cout << "ERROR: failed l1 example 2 for Particle Swarm optimization.\n";

    // Reporting.
    if (test.failed_any() or verbose) reportPassFail(test.success(), "Particle Swarm", "Algorithm Unit Tests");
    return test.success();
}

// Unit tests for TasOptimization::GradientDescentState.
bool testGradientDescentState(bool verbose) {
    test_result test;

    // Check contructor.
    size_t num_dimensions = 5;
    std::vector<double> dummy_x(num_dimensions, 1);
    GradientDescentState state = GradientDescentState(dummy_x, 0.1);
    test.tassert(num_dimensions == state.getNumDimensions());
    test.tassert(state.getAdaptiveStepsize() == 0.1);
    test.tassert(state.getX().size() == num_dimensions);

    // Check getters and coverters.
    std::vector<double> compare_x = state.getX();
    for (size_t i=0; i<compare_x.size(); i++) test.tassert(compare_x[i] == dummy_x[i]);
    std::fill(compare_x.begin(), compare_x.end(), 0);
    state.getX(compare_x.data());
    for (size_t i=0; i<compare_x.size(); i++) test.tassert(compare_x[i] == dummy_x[i]);
    std::fill(compare_x.begin(), compare_x.end(), 0);
    compare_x = state;
    for (size_t i=0; i<compare_x.size(); i++) test.tassert(compare_x[i] == dummy_x[i]);

    // Check setters.
    std::vector<double> new_x(num_dimensions, 2);
    state.setX(new_x);
    compare_x = state;
    for (size_t i=0; i<compare_x.size(); i++) test.tassert(new_x[i] == compare_x[i]);
    std::fill(compare_x.begin(), compare_x.end(), 0);
    state.setX(new_x.data());
    compare_x = state;
    for (size_t i=0; i<compare_x.size(); i++) test.tassert(new_x[i] == compare_x[i]);
    std::fill(compare_x.begin(), compare_x.end(), 0);
    state.setAdaptiveStepsize(0.2);
    test.tassert(state.getAdaptiveStepsize() == 0.2);

    // Reporting.
    if (test.failed_any() or verbose) reportPassFail(test.success(), "Gradient Descent", "State Unit Tests");
    return test.success();
}

/*! Nesterov's "worst function in the world" for convex optimization. This function is difficult for nearly all first-order
 * iterative optimization methods starting from x0=0. This generator writes the function and gradient to \b func and \b grad, and
 * the minimum is written to \b minimum. It is expected that 1 <= k <= (number of dimensions - 1) / 2 and the initial point
 * of an iterative algorithm is the zero vector.
 */
void makeNesterovTestFunction(const double L, const int k, ObjectiveFunctionSingle &func, GradientFunctionSingle &grad,
                              std::vector<double> &minimum) {
    func = [=](const std::vector<double> &x)->double {
        double result = (L / 4.0) * ((1.0 / 2.0) * (x[0] * x[0] + x[k-1] * x[k-1]) - x[0]);
        double delta;
        for (int i=0; i<k-1; i++) {
            delta = x[i] - x[i+1];
            result += (L / 8.0) * (delta * delta);
        }
        return result;
    };
    grad = [=](const std::vector<double> &x, std::vector<double> &gx)->void {
        std::fill(gx.begin(), gx.end(), 0);
        gx[0] = (L / 4.0) * (x[0] - 1.0);
        gx[k-1] = (L / 4.0) * x[k-1];
        for (int i=0; i<k-1; i++) {
            gx[i] += (L / 4.0) * (x[i] - x[i+1]);
            gx[i+1] -= (L / 4.0) * (x[i] - x[i+1]);
        }
    };
    std::fill(minimum.begin() + k, minimum.end(), 0);
    for (int i=0; i<k; i++) {
        minimum[i] = 1.0 - ((double) i + 1.0) / (k + 1.0);
    }
}

// Unit tests for TasOptimization::GradientDescent on a difficult convex problem.
bool testGradientDescent(bool verbose) {
    test_result test;

    // Nesterov's "worst function in the world".
    ObjectiveFunctionSingle func;
    GradientFunctionSingle grad;
    double L = 10;
    int num_dimensions = 11;
    std::vector<double> x_optimal(num_dimensions);
    makeNesterovTestFunction(L, (num_dimensions-1)/2, func, grad, x_optimal);

    // Constant stepsize gradient descent.
    std::vector<double> x0(num_dimensions, 0);
    GradientDescentState state(x0, 0);
    GradientDescent(grad, 1.0/L, 300, 1E-6, state);
    std::vector<double> x_gd = state.getX();
    for (int i=0; i<num_dimensions; i++) test.tassert(std::abs(x_gd[i] - x_optimal[i]) <= 1E-6);
    state.setX(x0);
    for (int t=0; t<300; t++) GradientDescent(grad, 1.0/L, 1, 1E-6, state);
    x_gd = state.getX();
    for (int i=0; i<num_dimensions; i++) test.tassert(std::abs(x_gd[i] - x_optimal[i]) <= 1E-6);

    // Variable stepsize gradient descent.
    state.setAdaptiveStepsize(10.0/L);
    state.setX(x0);
    GradientDescent(func, grad, 1.5, 1.25, 300, 1E-6, state);
    x_gd = state.getX();
    for (int i=0; i<num_dimensions; i++) test.tassert(std::abs(x_gd[i] - x_optimal[i]) <= 1E-6);
    state.setX(x0);
    for (int t=0; t<300; t++) GradientDescent(func, grad, 1.5, 1.25, 1, 1E-6, state);
    x_gd = state.getX();
    for (int i=0; i<num_dimensions; i++) test.tassert(std::abs(x_gd[i] - x_optimal[i]) <= 1E-6);

    // Proximal/Projected gradient descent (optimum now lies on the boundary).
    ProjectionFunctionSingle proj = [](const std::vector<double> &x, std::vector<double> &p) {
        for (size_t i=0; i<p.size(); i++) p[i] = std::min(std::max(x[i], -0.5), 0.5);
    };
    for (int i=0; i<(num_dimensions-1)/2; i++) x_optimal[i] = 0.5 - 0.1 * i;
    state.setAdaptiveStepsize(10.0/L);
    state.setX(x0);
    GradientDescent(func, grad, proj, 1.5, 1.25, 300, 1E-6, state);
    x_gd = state.getX();
    for (int i=0; i<num_dimensions; i++) test.tassert(std::abs(x_gd[i] - x_optimal[i]) <= 1E-6);
    for (int t=0; t<300; t++) GradientDescent(func, grad, proj, 1.5, 1.25, 1, 1E-6, state);
    x_gd = state.getX();
    for (int i=0; i<num_dimensions; i++) test.tassert(std::abs(x_gd[i] - x_optimal[i]) <= 1E-6);

    // Reporting.
    if (test.failed_any() or verbose) reportPassFail(test.success(), "Gradient Descent", "Algorithm Unit Tests");
    return test.success();
}


} // end namespace

bool DreamExternalTester::testOptimization(){
    test_result test;

    // Test Particle Swarm State and Algorithm.
    test.tassert(TasOptimization::testParticleSwarmState(verbose));
    test.tassert(TasOptimization::testParticleSwarm(verbose));
    reportPassFail(test.last_success(), "Optimization", "Particle Swarm");

    // Test Gradient Descent State and Algorithm.
    test.tassert(TasOptimization::testGradientDescentState(verbose));
    test.tassert(TasOptimization::testGradientDescent(verbose));
    reportPassFail(test.last_success(), "Optimization", "Gradient Descent");

    return test.success();
}

#endif