#include <algorithm>
#include <cmath>
#include <functional>
#include <iostream>
#include <vector>

// ガウス過程回帰
// 注意: あんまりテストしてない
class GP {
    static constexpr double EPS = 1e-10;
    using vec = std::vector<double>;
    using mat = std::vector<vec>;

    mat data_x;
    vec data_y;
    double data_mu;
    double data_sgm2;
    mat invK;
    std::function<double(const vec&, const vec&)> kernel;

    // x^T A yを計算
    double product(const vec& x, const mat& A, const vec& y) {
        double ret = 0;
        int N = A.size();
        for (int i = 0; i < N; i++) {
            double sum = 0;
            for (int j = 0; j < N; j++) sum += A[i][j] * y[j];
            ret += x[i] * sum;
        }
        return ret;
    }

    // Aの逆行列を計算
    mat inverse(mat A) {
        int N = A.size();
        // LU分解
        std::vector<int> p(N), ip(N);
        for (int i = 0; i < N; i++) p[i] = i;
        for (int i = 0; i < N; i++) {
            int pivot = i;
            for (int j = i + 1; j < N; j++) {
                if (fabs(A[j][i]) > fabs(A[pivot][i])) {
                    pivot = j;
                }
            }
            std::swap(A[pivot], A[i]);
            std::swap(p[pivot], p[i]);
            if (fabs(A[i][i]) < EPS)  // detA=0
                return mat();

            for (int j = i + 1; j < N; j++) {
                A[j][i] /= A[i][i];
                for (int k = i + 1; k < N; k++) {
                    A[j][k] -= A[i][k] * A[j][i];
                }
            }
        }
        for (int i = 0; i < N; i++) ip[p[i]] = i;
        // 逆行列
        mat B(N, vec(N, 0)), C(N, vec(N, 0)), ret(N, vec(N));
        for (int i = 0; i < N; i++) B[i][i] = 1.0 / A[i][i];
        for (int i = 1; i < N; i++) {
            for (int j = 0, k = i; k < N; j++, k++) {
                for (int l = j + 1; l <= k; l++) B[j][k] -= A[j][l] * B[l][k];
                B[j][k] /= A[j][j];
                for (int l = j; l < k; l++)
                    B[k][j] -=
                        (k <= l ? 1.0 : A[k][l]) * (l <= j ? 1.0 : B[l][j]);
            }
        }
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                for (int k = 0; k < N; k++) {
                    double u = (i > k ? 0 : B[i][k]);
                    double l = (j > k ? 0 : (j == k ? 1.0 : B[k][j]));
                    C[i][j] += u * l;
                }
            }
        }
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                ret[j][i] = C[j][ip[i]];
            }
        }
        return ret;
    }

   public:
    // カーネルの設定
    void setKernel(const std::function<double(const vec&, const vec&)>& f) {
        kernel = f;
    }
    // データ(y, x)と、yの正規ノイズの分散sgm2の設定
    void setData(const vec& y, const mat& x, double sgm2) {
        int N = y.size();
        // 平均
        double mu = 0;
        for (int i = 0; i < N; i++) mu += y[i];
        mu /= N;
        // 平均を引いたもの
        vec yy(N);
        for (int i = 0; i < N; i++) yy[i] = y[i] - mu;
        // 行列Kと逆行列
        mat K(N, vec(N, 0));
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                K[i][j] = kernel(x[i], x[j]) + (i == j ? sgm2 : 0);
            }
        }

        data_x = x;
        data_y = yy;
        data_mu = mu;
        data_sgm2 = sgm2;
        invK = inverse(K);  // O(N^3)
    }
    // xのときの予測値 Normal(mu, sgm2)
    std::pair<double, double> predict(const vec& x) {
        int N = data_y.size();
        vec kk(N);
        for (int i = 0; i < N; i++) kk[i] = kernel(data_x[i], x);
        double kkk = kernel(x, x) + data_sgm2;
        return std::make_pair(data_mu + product(kk, invK, data_y),
                              std::max(0.0, kkk - product(kk, invK, kk)));
    }
};

int main() {
    GP gp;

    // ガウスカーネルをセット
    gp.setKernel(
        [](const std::vector<double>& x, const std::vector<double>& y) {
            // パラメータは適当
            double theta1 = 1.0;
            double theta2 = 1.0;
            double sum = 0;
            for (int i = 0; i < x.size(); i++)
                sum += (x[i] - y[i]) * (x[i] - y[i]);
            return theta1 * exp(-sum / theta2);
        });

    // データ
    std::vector<double> y;
    std::vector<std::vector<double>> x;
    x.push_back({1.6});
    y.push_back(1.9);

    x.push_back({4.0});
    y.push_back(1.1);

    x.push_back({2.1});
    y.push_back(1.85);

    x.push_back({0.5});
    y.push_back(0.8);

    x.push_back({2.2});
    y.push_back(2.2);

    // データのセット
    gp.setData(y, x, 0.05);
    // 予測
    for (double i = 0; i <= 5; i += 0.05) {
        std::vector<double> tmp{i};
        std::pair<double, double> res = gp.predict(tmp);
        double mu = res.first;
        double sigma = sqrt(res.second);
        // 期待値
        std::cout << i << "\t" << mu << "\t";
        // 2σ区間(におさまる確率=約95.45%)
        std::cout << mu - 2 * sigma << "\t" << mu + 2 * sigma << std::endl;
    }

    return 0;
}