// Copyright (c) Dietmar Wolz.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory.

// Eigen based implementation of dual annealing 
// derived from https://github.com/scipy/scipy/blob/master/scipy/optimize/_dual_annealing.py
// Implementation only differs regarding boundary handling - this implementattion 
// uses boundary-normalized X values. Local search is fixed to LBFGS-B, see
// https://github.com/yixuan/LBFGSpp/tree/master/include 
// requires https://github.com/imneme/pcg-cpp

#include <Eigen/Core>
#include <iostream>
#include <float.h>
#include <math.h>
#include <ctime>
#include <random>
#include "pcg_random.hpp"
#include <LBFGSB.h>

using namespace LBFGSpp;
using namespace std;

typedef Eigen::Matrix<double, Eigen::Dynamic, 1> vec;
typedef Eigen::Matrix<int, Eigen::Dynamic, 1> ivec;
typedef Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> mat;

namespace dual_annealing {

typedef double (*callback_type)(int, double[]);

// wrapper around the fittness function, scales according to boundaries

class Fittness;

static uniform_real_distribution<> distr_01 = std::uniform_real_distribution<>(0, 1);
static normal_distribution<> gauss_01 = std::normal_distribution<>(0, 1);

static vec zeros(int n) {
	return  Eigen::MatrixXd::Zero(n, 1);
}

static Eigen::MatrixXd normalVec(int dim, pcg64& rs) {
	return Eigen::MatrixXd::NullaryExpr( dim, 1, [&](){return gauss_01(rs);});
}

static Eigen::MatrixXd uniformVec(int dim, pcg64& rs) {
	return Eigen::MatrixXd::NullaryExpr( dim, 1, [&](){return distr_01(rs);});
}

static vec emptyVec = {};

static vec logv(vec v) {
	return v.unaryExpr([](double x) { return log(x);});
}

static vec expv(vec v) {
	return  v.unaryExpr([](double x) { return exp(x);});
}

double minLBFGS(Fittness* fitfun, vec& X0_, int maxIterations);

class Fittness {

public:

	vec lower;
	vec upper;

	Fittness(callback_type pfunc, vec *lower_limit, vec *upper_limit,
			long maxEvals_) {
		func = pfunc;
		lower = *lower_limit;
		upper = *upper_limit;
		if (lower.size() > 0) // bounds defined
			scale = (upper - lower);
		maxEvals = maxEvals_;
	}

    vec getClosestFeasible(const vec& X) const {
        if (lower.size() > 0) {
        	return X.cwiseMin(1.0).cwiseMax(-1.0);
        }
        return X;
    }

    double eval(const vec& X) {
        int n = X.size();
        double parg[n];
        for (int i = 0; i < n; i++)
            parg[i] = X(i);
        double res = func(n, parg);
        evaluationCounter++;
        return res;
    }

    double value(const vec& X) {
    	double res = DBL_MAX;
        if (lower.size() > 0)
        	res = eval(decode(getClosestFeasible(X)));
        else
            res = eval(X);
        if (res < bestY) {
        	bestY = res;
        	bestX = vec(X);
        }
        return res;
    }

    const double LS_MAXITER_RATIO = 6;
    const double LS_MAXITER_MIN = 100;
	const double LS_MAXITER_MAX = 1000;

    double local_search(const vec& x0, double currval, vec& res) {
    	vec init = getClosestFeasible(x0);
    	bestY = DBL_MAX;
		int maxIter = LS_MAXITER_RATIO*x0.size();
		if (maxIter > LS_MAXITER_MAX)
			maxIter =  LS_MAXITER_MAX;
		if (maxIter < LS_MAXITER_MIN)
			maxIter = LS_MAXITER_MIN;
    	minLBFGS(this, init, maxIter);
    	if (bestY < DBL_MAX) {
    		for (int i = 0; i < res.size(); i++)
    			res[i] = bestX(i);
    	}
    	return bestY;
	}

    vec encode(const vec& X) const {
        if (lower.size() > 0)
            return (X - lower).array() / scale.array();
        else
            return X;
    }

    vec decode(const vec& X) const {
        if (lower.size() > 0)
            return X.cwiseProduct(scale) + lower;
        else
            return X;
    }

    int getEvaluations() {
        return evaluationCounter;
    }

    bool maxEvalReached() {
    	return evaluationCounter >= maxEvals;
    }

private:
   callback_type func;
   long evaluationCounter = 0;
   long maxEvals;
   vec scale;
   double bestY = DBL_MAX;
   vec bestX;
};


class LBFGSFunc {
private:
   Fittness* func;
   int dim;

public:

	LBFGSFunc(Fittness *fittness_, int dim_) {
		func = fittness_;
		dim = dim_;
	}

	double operator()(const vec& x, vec& grad) {
        if (!x.allFinite())
            return DBL_MAX;
		double eps = 1E-6;
		vec arg = vec(dim);
		for (int i = 0; i < dim; i++)
			arg[i] = x(i);
		for (int i = 0; i < dim; i++) {
			vec x1 = vec(arg);
			vec x2 = vec(arg);
			double e1 = eps;
			double e2 = eps;
			x1[i] += eps;
			if (x1[i] > 1) {
				x1[i] = 1;
				e1 = 1 - arg[i];
			}
			x2[i] -= eps;
			if (x2[i] < 0) {
				x2[i] = 0;
				e2 = arg[i];
			}
			double f1 = func->value(x1);
			double f2 = func->value(x2);
			grad[i] = (f1 - f2) / (e1 + e2);
		}
		double f = func->value(arg);
		return f;
	}
};

double minLBFGS(Fittness* fitfun, vec& X0, int maxIterations) {
    int dim = X0.size();
    LBFGSFunc fun = LBFGSFunc(fitfun, dim);

	LBFGSBParam<double> param;
	param.max_iterations = maxIterations;
    LBFGSBSolver<double> solver(param);
    vec lb = vec::Constant(dim, 0.0);
    vec ub = vec::Constant(dim, 1.0);
    // Initial values
    vec x = vec::Constant(dim, 0);
    for (int i = 0; i < dim; i++)
    	x[i] = X0[i];
    double fx;
    int niter;
    try {
    	niter = solver.minimize(fun, x, fx, lb, ub);
    } catch (std::exception& e) {
     	//cout << e.what() << endl;
        return DBL_MAX;
    }
    return fx;
}

class VisitingDistribution {

	    //Class used to generate new coordinates based on the distorted
	    //Cauchy-Lorentz distribution. Depending on the steps within the Markov
	    //chain, the class implements the strategy for generating new location
	    //changes.

public:

	VisitingDistribution(int dim, double visiting_param_, pcg64* rs_) {
		_visiting_param = visiting_param_;
		rs = rs_;

		// these are invariant numbers unless visiting_param changes
		double factor2 = exp((4.0 - _visiting_param) * log(
				_visiting_param - 1.0));
		double factor3 = exp((2.0 - _visiting_param) * log(2.0)
							/ (_visiting_param - 1.0));
		_factor4_p = sqrt(M_PI) * factor2 / (factor3 * (
				3.0 - _visiting_param));

		double factor5 = 1.0 / (_visiting_param - 1.0) - 0.5;
		double d1 = 2.0 - factor5;
		_factor6 = M_PI * (1.0 - factor5) / sin(
				M_PI * (1.0 - factor5)) / exp(lgamma(d1));
	}

	vec visiting(const vec& x, int step, double temperature) {
		//Based on the step in the strategy chain, new coordinated are
		//generated by changing all components is the same time or only
		//one of them, the new values are computed with visit_fn method

		int dim = x.size();
		if (step < dim) {
			// Changing all coordinates with a new visiting value
			double upper_sample = distr_01(*rs);
			double lower_sample = distr_01(*rs);
			vec visits = visit_fn(temperature, dim);
			for (int i = 0; i < dim; i++) {
				if (visits[i] > TAIL_LIMIT)
					visits[i] = TAIL_LIMIT * upper_sample;
				else if (visits[i] < -TAIL_LIMIT)
					visits[i] = -TAIL_LIMIT * lower_sample;
			}
			vec x_visit = visits + x;
			vec a = x_visit;
			vec b = vec(dim);
			for (int i = 0; i < dim; i++) {
				b[i] = fmod(a[i], 1) + 1;
				x_visit[i] = fmod(b[i], 1);
				if (abs(x_visit[i]) < MIN_VISIT_BOUND)
					x_visit[i] += 1.e-10;
			}
			//cerr << step << " " << temperature <<  endl;// << x_visit << endl;
			return x_visit;
		} else {
			// Changing only one coordinate at a time based on strategy
			// chain step
			vec x_visit = vec(x);
			double visit = visit_fn(temperature, 1)[0];
			if (visit > TAIL_LIMIT)
				visit = TAIL_LIMIT * distr_01(*rs);
			else if (visit < -TAIL_LIMIT)
				visit = -TAIL_LIMIT * distr_01(*rs);
			int index = step - dim;
			x_visit[index] = visit + x[index];
			double a = x_visit[index];
			double b = fmod(a, 1) + 1;
			x_visit[index] = fmod(b, 1);
			if (abs(x_visit[index]) < MIN_VISIT_BOUND)
				x_visit[index] += MIN_VISIT_BOUND;
			//cerr << step << " " << temperature <<  endl;// << x_visit << endl;
			return x_visit;
		}
	}

	vec visit_fn(double temperature, int dim) {

	    //Formula Visita from p. 405 of reference [2]
		vec x = normalVec(dim, *rs);
		vec y = normalVec(dim, *rs);;

		double factor1 = exp(log(temperature) / (_visiting_param - 1.0));
		double factor4 = _factor4_p * factor1;

		// sigmax
		x = x * exp(-(_visiting_param - 1.0) * log(
			_factor6 / factor4) / (3.0 - _visiting_param));

		vec den = expv(logv(y.cwiseAbs() * (_visiting_param - 1.0)) / (3.0 - _visiting_param));
		return x.cwiseQuotient(den);
	}

private:

	pcg64* rs;
	double _visiting_param;
	double _factor4_p;
	double _factor6;

	const double TAIL_LIMIT = 1.e8;
	const double MIN_VISIT_BOUND = 1.e-10;
};

class nanexception: public exception {
  virtual const char* what() const throw() {
    return "Objective function is returning nan";
  }
} naneexc;

const double BIG_VALUE = 1e16;

class EnergyState {

    //Class used to record the energy state-> At any time, it knows what is the
    //currently used coordinates and the most recent best location
public:

    double ebest;
    vec xbest;
    double current_energy;
    vec current_location;

	EnergyState(int dim_) {
		dim = dim_;
		ebest = DBL_MAX;
		xbest = { };
	    current_energy = DBL_MAX;
	    current_location = {};
	}

    void reset(Fittness* owf, pcg64* rs, const vec& x0) {
        if (x0.size() == 0)
            current_location = normalVec(dim, *rs);
        else
            current_location = vec(x0);
        bool init_error = true;
        int reinit_counter = 0;
        while (init_error) {
            current_energy = owf->value(current_location);
            if (current_energy >=  BIG_VALUE || isnan(current_energy)) {
                if (reinit_counter >= MAX_REINIT_COUNT) {
                    init_error = false;
                    throw naneexc;
                }
                current_location = uniformVec(dim, *rs);
                reinit_counter++;
            } else
                init_error = false;
            // If first time reset, initialize ebest and xbest
            if (ebest == DBL_MAX && xbest.size() == 0) {
                ebest = current_energy;
                xbest = vec(current_location);
            }
            // Otherwise, keep them in case of reannealing reset
        }
    }

    void update_best(double e, const vec& x) {
         ebest = e;
         xbest = vec(x);
	}

    void update_current(double e, const vec& x) {
         current_energy = e;
         current_location = vec(x);
	}

private:
	// Maximimum number of trials for generating a valid starting point
    int MAX_REINIT_COUNT = 1000;
    int dim;
};

class StrategyChain {
    // Class used for the Markov chain and related strategy for local search
    // decision
public:

	StrategyChain(double acceptance_param_, VisitingDistribution* vd_,
			Fittness* ofw_, pcg64* rs_, EnergyState* state_) {
        // Global optimizer state
        state = state_;
        // Local markov chain minimum energy and location
        emin = state->current_energy;
        xmin = vec(state->current_location);
        // Acceptance parameter
        acceptance_param = acceptance_param_;
        // Visiting distribution instance
        vd = vd_;
        // Wrapper to objective function and related local minimizer
        ofw = ofw_;
        not_improved_idx = 0;
        not_improved_max_idx = 1000;
        rs = rs_;
        temperature_step = 0;
        K = 100 * (state->current_location).size();
	}

	void accept_reject(int j, double e, const vec& x_visit) {
		double r = distr_01(*rs);
		double pqv_temp = (acceptance_param - 1.0) * (e - state->current_energy)
				/ (temperature_step + 1.);
		double pqv = 0;
		if (pqv_temp < 0.)
			pqv = 0.;
		else
			pqv = exp(log(pqv_temp) / (1. - acceptance_param));
		if (r <= pqv) {
			// We accept the new location and update state
			state->update_current(e, x_visit);
			xmin = vec(state->current_location);
		}
		// No improvement since long time
		if (not_improved_idx >= not_improved_max_idx) {
			if (j == 0 || state->current_energy < emin) {
				emin = state->current_energy;
				xmin = vec(state->current_location);
			}
		}
	}

    void run(int step, double temperature) {
        temperature_step = temperature / (double)(step + 1);
        not_improved_idx += 1;
        for (unsigned int j = 0; j < (state->current_location).size() * 2; j++) {
            if (j == 0)
                state_improved = false;
            if (step == 0 && j == 0)
                state_improved = true;
            vec x_visit = vd->visiting(state->current_location, j, temperature);
            // Calling the objective function
            double e = ofw->value(x_visit);
            if (e < state->current_energy) {
                // We have got a better energy value
            	state->update_current(e, x_visit);
                if (e < state->ebest) {
                	state->update_best(e, x_visit);
                    state_improved = true;
                    not_improved_idx = 0;
                }
            } else {
                // We have not improved but do we accept the new location?
            	accept_reject(j, e, x_visit);
            }
        	if (ofw->maxEvalReached())
        		return;
        }	// End of StrategyChain loop
    }

    void local_search() {
        // Decision making for performing a local search
        // based on Markov chain results
        // If energy has been improved or no improvement since too long,
        // performing a local search with the best Markov chain location
    	int dim = state->xbest.size();
        if (state_improved) {
            // Global energy has improved, let's see if LS improved further
        	vec x = vec(dim);
            double e = ofw->local_search(state->xbest, state->ebest, x);
            if (e < state->ebest) {
                not_improved_idx = 0;
                state->update_best(e,x);
                state->update_current(e,x);
            	if (ofw->maxEvalReached())
            		return;
            }
        }
        // Check probability of a need to perform a LS even if no improvment
        bool do_ls = false;
        if (K < 90 * state->current_location.size()) {
            double pls = exp(K * (state->ebest - state->current_energy) / temperature_step);
            if (pls >= distr_01(*rs))
                do_ls = true;
        }
        // Global energy not improved, let's see what LS gives
        // on the best strategy chain location
        if (not_improved_idx >= not_improved_max_idx)
            do_ls = true;
        if (do_ls) {
        	vec x = vec(dim);
            double e = ofw->local_search(xmin, state->ebest, x);
            xmin = vec(x);
            emin = e;
            not_improved_idx = 0;
            not_improved_max_idx = state->current_location.size();
            if (e < state->ebest) {
            	state->update_best(emin, xmin);
            	state->update_current(e, x);
            }
        }
    }

private:

    double emin;
    vec xmin;
    EnergyState* state;
    double acceptance_param;
    VisitingDistribution* vd;
    int not_improved_idx;
    int not_improved_max_idx;
    pcg64* rs;
	Fittness* ofw;
	double temperature_step;
	double K;
    bool state_improved = false;
};

class sizeexception: public exception {
  virtual const char* what() const throw() {
    return "Bounds size does not match x0";
  }
} sizeeexc;


class DARunner {

public:

	DARunner(Fittness* fun_, vec& x0_, long seed_, bool use_local_search_) {
		owf = fun_;
        if (x0_.size() > 0 && x0_.size() != owf->lower.size())
            throw sizeeexc;
        //Initialization of RandomState for reproducible runs if seed provided
        rs = new pcg64(seed_);
        use_local_search = use_local_search_;
        // Initialization of the energy state
        es = new EnergyState(owf->lower.size());
        es->reset(owf, rs, x0_);
        // VisitingDistribution instance
        vd = new VisitingDistribution(owf->lower.size(), qv, rs);
        // Markov chain instance
        sc = new StrategyChain(qa, vd, owf, rs, es);
	}

	~DARunner() {
		delete rs;
		delete vd;
		delete sc;
		delete es;
	}

    void search() {
        iter = 0;
        double t1 = exp((qv - 1) * log(2.0)) - 1.0;
        for (;;) {
            for (int i = 0; i < maxsteps; i++) {
                // Compute temperature for this step
                double s = i + 2.0;
                double t2 = exp((qv - 1) * log(s)) - 1.0;
                double temperature = temperature_start * t1 / t2;
                if (iter++ >= maxsteps)
                    return;
                // Need a re-annealing process?
                if (temperature < temperature_restart) {
                    es->reset(owf, rs, emptyVec);
                    break;
                }
                // starting strategy chain
                sc->run(i, temperature);
                if (owf->maxEvalReached())
                    return;
                if (use_local_search) {
                    sc->local_search();
                    if (owf->maxEvalReached())
                    	return;
                }
            }
        }
    }

	vec bestX() {
		return es->xbest;
    }

	double bestY() {
		return es->ebest;
	}

private:

	int MAX_REINIT_COUNT = 1000;
	double temperature_start = 5230;
	double qv = 2.62;
	double qa = -5.0;
	bool use_local_search;
    // maximum number of step (main iteration) 
    double maxsteps = 1000;
    // minimum value of annealing temperature reached to perform
    // re-annealing temperature_start
    double temperature_restart = 0.1;
    Fittness* owf;
    pcg64* rs;
    EnergyState* es;
    StrategyChain* sc;
    VisitingDistribution* vd;
    int iter = 0;
};

double minimize(Fittness* fun, vec& x0, long seed, 
                bool use_local_search, vec& X) {
	DARunner gr = DARunner(fun, x0, seed, use_local_search);
	gr.search();
	int dim = x0.size();
	vec bx = gr.bestX();
	for (int i = 0; i < dim; i++)
		X[i] = bx[i];
	return gr.bestY();
}
}

using namespace dual_annealing;

extern "C" {
    double* optimizeDA_C(long runid, callback_type func, int dim, int seed,
    		double *init, double *lower, double *upper, int maxEvals, 
    		bool use_local_search) {
        int n = dim;
        double *res = new double[n + 4];
        vec guess(n), lower_limit(n), upper_limit(n);
        bool useLimit = false;
        for (int i = 0; i < n; i++) {
            guess[i] = init[i];
            lower_limit[i] = lower[i];
            upper_limit[i] = upper[i];
            useLimit |= (lower[i] != 0);
            useLimit |= (upper[i] != 0);
        }
        if (useLimit == false) {
            lower_limit.resize(0);
            upper_limit.resize(0);
        }
        if (maxEvals <= 0)
        	maxEvals = 1E7;
        Fittness fitfun(func, &lower_limit, &upper_limit, maxEvals);

        try {
           	vec X = zeros(dim);
           	vec enc = fitfun.encode(guess);
        	double bestY = minimize(&fitfun, enc, seed, 
        	                        use_local_search, X);
        	vec bestX = fitfun.decode(X);
            for (int i = 0; i < n; i++)
                res[i] = bestX[i];
            res[n] = bestY;
            res[n+1] = fitfun.getEvaluations();
            res[n+2] = 0;
            res[n+3] = 0;
            return res;
        } catch (std::exception& e) {
            cerr << e.what() << endl;
            return res;
        }
    }
}


