/// Nelder--Mead method

#include <stdio.h>
#include <math.h>

const int n = 2; /// number of dimensions

const double Delta = 1.0;
const double Theta = 0.5;
const double Gamma = 2.9;

void initSimplex(double x0[], double dx[], double x[][n]) {
    double p = (sqrt(n+1.0) - 1.0 + n) / (n * sqrt(2.0));
    double u = (sqrt(n+1.0) - 1.0) / (n * sqrt(2.0));

    for (int i = 0; i < n; i++) x[0][i] = x0[i];
    for (int i = 1; i <= n; i++) {
        for (int j = 0; j < n; j++) {
            x[i][j] = x0[j] + (i == j + 1 ? p : u) * dx[j];
        }
    }
}

/// The first argument (f) is a function that takes a point (an array of n coordinates).
/// The second argument (x0) is the initial point of the algorithm.
/// The third argument (dx) is the scaling of the simplex in each dimension.
/// The fourth argument is the epsilon that determines the termination condition.
/// The return value is the found maximum of the function f.
/// The coordinates of the point the maximum is achieved at are returned in the array x0.

double NelderMead(double (*f)(double *), double x0[], double dx[], double eps) {
    double x[n + 1][n]; // the vertices of the simplex
    double q[n + 1]; // the values f(x[i])
    double xc[n], qc; // the center of mass
    double xN[n], qN;
    double _xN[n], _qN;
    double xM[n], qM;
    double crit;
    int w = -1, w1 = -1, b = -1; // worst(min)/second-worst/best(max) index

    initSimplex(x0, dx, x);
    for (int i = 0; i <= n; i++) q[i] = f(x[i]);
    
    while (1) {
        /// Find the index of the worst, second-worst and best vertex.
        w = b = n;
        for (int i = 0; i < n; i++)
            if (q[i] < q[w]) w = i;
        for (int i = 0; i < n; i++)
            if (q[i] > q[b]) b = i;
        w1 = b;
        for (int i = 0; i <= n; i++)
            if (q[i] < q[w1] && i != w) w1 = i;

        /// Calculate the center of mass.
        for (int i = 0; i < n; i++) {
            xc[i] = -x[w][i];
            for (int j = 0; j <= n; j++) {
                xc[i] += x[j][i];
            }
            xc[i] /= n;
        }
        
        /// Termination condition
        qc = f(xc);
        crit = 0.0;
        for (int i = 0; i <= n; i++) {
            crit += (q[i] - qc) * (q[i] - qc);
        }
        crit = sqrt(crit / (n + 1.0));
        if (crit <= eps) break;

        /// reflect vertex
        for (int i = 0; i < n; i++)
            xN[i] = xc[i] + Delta * (xc[i] - x[w][i]);
        qN = f(xN);
        
        /// simplex deformations
        if (qN > q[b]) { /// expand the simplex
            for (int i = 0; i < n; i++)
                _xN[i] = xc[i] + Gamma * (xN[i] - xc[i]);
            _qN = f(_xN);
            if (_qN > q[b]) {
                for (int i = 0; i < n; i++) x[w][i] = _xN[i];
                q[w] = _qN;
            } else {
                for (int i = 0; i < n; i++) x[w][i] = xN[i];
                q[w] = qN;
            }
        } else if (qN >= q[w] && qN <= q[w1]) { /// shrink the simplex
            for (int i = 0; i < n; ++i) xM[i] = xc[i] + Theta * (x[w][i] - xc[i]);
            qM = f(xM);
            for (int i = 0; i < n; ++i) x[w][i] = xM[i];
            q[w] = qM;
        } else if (qN < q[w]) { /// reduce the simplex
            for (int j = 0; j <= n; j++) {
                if (j == b) continue;
                for (int i = 0; i < n; i++)
                    x[j][i] = x[b][i] + 0.5 * (x[j][i] - x[b][i]);
                q[j] = f(x[j]);
            }
        } else { /// no simplex deformations (just reflection)
            for (int i = 0; i < n; i++) x[w][i] = xN[i];
            q[w] = f(xN);
        }
    }
    for (int i = 0; i < n; i++) x0[i] = x[b][i];
    return q[b];
}

int it1, it2, it3;

double f1(double *_x) { /// https://en.wikipedia.org/wiki/Rosenbrock_function
    it1++;
    const double a = 1.0, b = 100.0;
    double x = _x[0], y = _x[1];
    return -(a-x)*(a-x)-b*(y-x*x)*(y-x*x);
}

double f2(double *_x) { /// https://en.wikipedia.org/wiki/Himmelblau%27s_function
    it2++;
    double x = _x[0], y = _x[1];
    return -(x*x+y-11)*(x*x+y-11)-(x+y*y-7)*(x+y*y-7);
}

/// Goldstein--Price function
double f3(double *_x) { /// https://en.wikipedia.org/wiki/Test_functions_for_optimization
    it3++;
    double x = _x[0], y = _x[1];
    return -(1+(x+y+1)*(x+y+1)*(19-14*x+3*x*x-14*y+6*x*y+3*y*y))*(30+(2*x-3*y)*(2*x-3*y)*(18-32*x+12*x*x+48*y-36*x*y+27*y*y));
}

void runTests() {
    double x0[2];
    double dx[2];
    double maxf;

    dx[0] = dx[1] = 1.0;

    x0[0] = x0[1] = 0.0;
    maxf = NelderMead(f1, x0, dx, 1e-18);
    printf("(Local) maximum: %.20lf\nAt (%.20lf, %.20lf)\nIn %d iterations.\n\n", maxf, x0[0], x0[1], it1);

    x0[0] = x0[1] = 0.0;
    maxf = NelderMead(f2, x0, dx, 1e-18);
    printf("(Local) maximum: %.20lf\nAt (%.20lf, %.20lf)\nIn %d iterations.\n\n", maxf, x0[0], x0[1], it2);

    x0[0] = x0[1] = 0.0;
    maxf = NelderMead(f3, x0, dx, 1e-18);
    printf("(Local) maximum: %.20lf\nAt (%.20lf, %.20lf)\nIn %d iterations.\n\n", maxf, x0[0], x0[1], it3);
}

int main() {
    runTests();
    return 0;
}