#include "manip.h"
#include <iostream>
#include <cmath>
#include <cassert>
#include <vector>
#include <utility>
#include <algorithm>
#include <boost/multi_array.hpp>
#include <boost/bind.hpp>
#include "gd.h"
#include "../image/image.h"

#define GOSA_KAKUSAN_YCBCR(dx, dy, rate, rate_total) do { \
        const int tx = x + d * dx; \
        const int ty = y + dy; \
        if(0 <= tx && tx < width && 0 <= ty && ty < height) { \
            gosa[tx][ty][0] += gosa_y  * rate / rate_total; \
            gosa[tx][ty][1] += gosa_cb * rate / rate_total; \
            gosa[tx][ty][2] += gosa_cr * rate / rate_total; \
        } \
    } while(0)

#define GOSA_KAKUSAN_RGB(dx, dy, rate, rate_total) do { \
        const int tx = x + d * dx; \
        const int ty = y + dy; \
        if(0 <= tx && tx < width && 0 <= ty && ty < height) { \
            gosa[tx][ty][0] += gosa_r * rate / rate_total; \
            gosa[tx][ty][1] += gosa_g * rate / rate_total; \
            gosa[tx][ty][2] += gosa_b * rate / rate_total; \
        } \
    } while(0)

namespace manip {
    namespace {
        // 大津の手法による閾値の計算
        int otsu_threshold(const double average, const int histgram[256]) {
            double max = 0.;
            int max_no = 0;
            for(int i = 0; i < 256; ++i) {
                long count1 = 0, count2 = 0;
                long long data = 0;
                double breakup1 = 0., breakup2 = 0.;
                double average1 = 0., average2 = 0.;
                for(int j = 0; j < i; ++j) {
                    count1 += histgram[j];
                    data += histgram[j] * j;
                }
                if(count1 > 0) {
                    average1 = (double)data / (double)count1;
                    for(int j = 0; j < i; ++j) {
                        breakup1 += pow(j - average1, 2) * histgram[j];
                    }
                    breakup1 /= (double)count1;
                }

                data = 0;
                for(int j = i; j < 256; ++j) {
                    count2 += histgram[j];
                    data += histgram[j] * j;
                }
                if(count2 > 0) {
                    average2 = (double)data / (double)count2;
                    for(int j = i; j < 256; ++j) {
                        breakup2 += pow(j - average2, 2) * histgram[j];
                    }
                    breakup2 /= (double)count2;
                }
                const double class1 = (double)count1 * breakup1 + (double)count2 * breakup2;
                const double class2 = (double)count1 * pow(average1 - average, 2) + (double)count2 * pow(average2 - average, 2);
                const double tmp = class2 / class1;
                if(max < tmp) {
                    max = tmp;
                    max_no = i;
                }
            }
            return max_no;
        }

        // 3x3 微分オペレータを適用する
        bool apply_operator_3x3(gd &img, const double op[3][3], double filter_div, double offset) {
            const int width = img.width();
            const int height = img.height();
            gd dst(width, height);
            dst.alpha_blending(false);

            for(int y = 0; y < height; ++y) {
                for(int x = 0; x < width; ++x) {
                    double sum_r = 0., sum_g = 0., sum_b = 0., sum_a = 0.;
                    for(int j = -1; j <= 1; ++j) {
                        const int ref_y = std::min(std::max(0, y + j), height - 1);
                        for(int i = -1; i <= 1; ++i) {
                            const int ref_x = std::min(std::max(0, x + i), width - 1);
                            const double ref_op = op[i + 1][j + 1];
                            const gd::color color = img.pixel_fast(ref_x, ref_y);
                            const int a = (color & 0x7f000000) >> 24;
                            if(a == 0x7f) {
                                sum_a += a * ref_op;
                            } else if(std::abs(ref_op) > 0) {
                                const int r = (color & 0xff0000) >> 16;
                                const int g = (color & 0x00ff00) >> 8;
                                const int b = (color & 0x0000ff);
                                sum_r += r * ref_op;
                                sum_g += g * ref_op;
                                sum_b += b * ref_op;
                                sum_a += a * ref_op;
                            }
                        }
                    }

                    const int r = std::max(0, std::min(255, static_cast<int>(sum_r / filter_div + offset + 0.5)));
                    const int g = std::max(0, std::min(255, static_cast<int>(sum_g / filter_div + offset + 0.5)));
                    const int b = std::max(0, std::min(255, static_cast<int>(sum_b / filter_div + offset + 0.5)));
                    const int a = std::max(0, std::min(127, static_cast<int>(sum_a + 0.5)));
                    dst.pixel_fast(x, y, (a << 24) | (r << 16) | (g << 8) | b);
                }
            }
            img.swap(dst);
            return true;
        }

        bool do_grayscale(gd &img, int black = 0, int white = 255, gd::color background = 0x7fffffff) {
            img.convert_to_true_color();
            img.alpha_blending(false);
            if((background & 0x7f000000) != 0x7f000000) {
                if(!fill_background(img, background)) {
                    return false;
                }
            }
            const int width = img.width();
            const int height = img.height();
            const double range = white - black;
            for(int y = 0; y < height; ++y) {
                for(int x = 0; x < width; ++x) {
                    const gd::color color = img.pixel_fast(x, y);
                    const int r = (color & 0xff0000) >> 16;
                    const int g = (color & 0x00ff00) >>  8;
                    const int b = (color & 0x0000ff);
                    const int a = (color & 0x7f000000); // そのまま使うのでビットシフトしない
                    const double gray = (r * 0.298912) + (g * 0.586611) + (b * 0.114478);
                    const int put =
                        (gray <= black)
                            ? 0
                            : (gray >= white)
                                ? 255
                                : std::min(255, std::max(0, static_cast<int>((gray - black) * 255. / range + 0.5)));
                    img.pixel_fast(x, y, a | (put * 0x010101));
                }
            }
            return true;
        }

        inline int convert_to_websafe_color(int c) {
            return c < 0 ? 0 : c > 255 ? 255 : ((int)round(c / 51.)) * 51;
        }

        inline double rgb_to_y(int r, int g, int b) {
            return (0.29900 * r) + (0.58700 * g) + (0.11400 * b);
        }

        inline double rgb_to_cb(int r, int g, int b) {
            return (-0.16874 * r) - (0.33126 * g) + (0.50000 * b) + 128.;
        }

        inline double rgb_to_cr(int r, int g, int b) {
            return (0.50000 * r) - (0.41869 * g) - (0.08131 * b) + 128.;
        }

        inline int ycbcr_to_rgb(double y, double cb, double cr) {
            cb -= 128.;
            cr -= 128.;
            const double r_ = y                  + (1.40200 * cr);
            const double g_ = y - (0.34414 * cb) - (0.71414 * cr);
            const double b_ = y + (1.77200 * cb);
            const int r = static_cast<int>(r_ < 0.5 ? 0 : r_ >= 254.5 ? 255.0 : (r_ + 0.5));
            const int g = static_cast<int>(g_ < 0.5 ? 0 : g_ >= 254.5 ? 255.0 : (g_ + 0.5));
            const int b = static_cast<int>(b_ < 0.5 ? 0 : b_ >= 254.5 ? 255.0 : (b_ + 0.5));
            return (r << 16) | (g << 8) | b;
        }

        inline size_t find_nearest_color_y_cb_cr(double y, double cb, double cr, const int palette[], const size_t palette_size) {
            y  = y  < 0.0 ? 0.0 : y  > 255.0 ? 255.0 : y;
            cb = cb < 0.0 ? 0.0 : cb > 255.0 ? 255.0 : cb;
            cr = cr < 0.0 ? 0.0 : cr > 255.0 ? 255.0 : cr;
            size_t min_index = 0;
            double min_score = INFINITY;
            for(size_t i = 0; i < palette_size; ++i) {
                const int p_color = palette[i];
                const double p_y  = (double)((p_color & 0xff0000) >> 16);
                const double p_cb = (double)((p_color & 0x00ff00) >>  8);
                const double p_cr = (double)((p_color & 0x0000ff)      );
                const double d_cbcr = sqrt(pow(cb - p_cb, 2.) + pow(cr - p_cr, 2.));
                const double score = sqrt(pow(y - p_y, 2.) + pow(d_cbcr, 2));
                if(score < min_score) {
                    min_index = i;
                    min_score = score;
                }
            }
            return min_index;
        }

        int famicom_convert(gd &img, DITHERING_METHOD dither, const int palette[], const size_t palette_size, int used_count[]) {
            img.alpha_blending(false);
            const int width = img.width();
            const int height = img.height();
            boost::multi_array<double, 3> gosa(boost::extents[width][height][3]);
            std::fill(gosa.origin(), gosa.origin() + gosa.size(), 0);

            for(int y = 0; y < height; ++y) {
                const int d = (y % 2 == 0) ? 1 : -1;
                const int x_begin = (y % 2 == 0) ? 0 : width - 1;
                const int x_end   = (y % 2 == 0) ? width : -1;
                for(int x = x_begin; x != x_end; x += d) {
                    const gd::color color = img.pixel_fast(x, y);
                    const int r = (color & 0x00ff0000) >> 16;
                    const int g = (color & 0x0000ff00) >> 8;
                    const int b = (color & 0x000000ff);
                    const double cy = rgb_to_y(r, g, b)  + gosa[x][y][0];
                    const double cb = rgb_to_cb(r, g, b) + gosa[x][y][1];
                    const double cr = rgb_to_cr(r, g, b) + gosa[x][y][2];
                    const size_t index = find_nearest_color_y_cb_cr(cy, cb, cr, palette, palette_size);
                    ++used_count[index];
                    const int p_y  = (palette[index] & 0xff0000) >> 16;
                    const int p_cb = (palette[index] & 0x00ff00) >>  8;
                    const int p_cr = (palette[index] & 0x0000ff);
                    img.pixel_fast(x, y, ycbcr_to_rgb(p_y, p_cb, p_cr));
                    const double gosa_y  = cy - p_y;
                    const double gosa_cb = cb - p_cb;
                    const double gosa_cr = cr - p_cr;

                    switch(dither) {
                    case DITHERING_NONE:
                    default:
                        break;
                    case DITHERING_FLOYD_STEINBERG:
                        // - X 7
                        // 3 5 1
                        GOSA_KAKUSAN_YCBCR( 1, 0, 7., 16.);
                        GOSA_KAKUSAN_YCBCR(-1, 0, 3., 16.);
                        GOSA_KAKUSAN_YCBCR( 0, 0, 5., 16.);
                        GOSA_KAKUSAN_YCBCR( 1, 0, 1., 16.);
                        break;
                    case DITHERING_SIERRA_3LINE:
                        // - - X 5 3
                        // 2 4 5 4 2
                        // 0 2 3 2 0
                        GOSA_KAKUSAN_YCBCR( 1, 0, 5., 32.);
                        GOSA_KAKUSAN_YCBCR( 2, 0, 3., 32.);
                        GOSA_KAKUSAN_YCBCR(-2, 1, 2., 32.);
                        GOSA_KAKUSAN_YCBCR(-1, 1, 4., 32.);
                        GOSA_KAKUSAN_YCBCR( 0, 1, 5., 32.);
                        GOSA_KAKUSAN_YCBCR( 1, 1, 4., 32.);
                        GOSA_KAKUSAN_YCBCR( 2, 1, 2., 32.);
                        GOSA_KAKUSAN_YCBCR(-1, 2, 2., 32.);
                        GOSA_KAKUSAN_YCBCR( 0, 2, 3., 32.);
                        GOSA_KAKUSAN_YCBCR( 1, 2, 2., 32.);
                        break;
                    case DITHERING_SIERRA_2LINE:
                        // - - X 4 3
                        // 1 2 3 2 1
                        GOSA_KAKUSAN_YCBCR( 1, 0, 4., 16.);
                        GOSA_KAKUSAN_YCBCR( 2, 0, 3., 16.);
                        GOSA_KAKUSAN_YCBCR(-2, 1, 1., 16.);
                        GOSA_KAKUSAN_YCBCR(-1, 1, 2., 16.);
                        GOSA_KAKUSAN_YCBCR( 0, 1, 3., 16.);
                        GOSA_KAKUSAN_YCBCR( 1, 1, 2., 16.);
                        GOSA_KAKUSAN_YCBCR( 2, 1, 1., 16.);
                        break;
                    case DITHERING_SIERRA_LITE:
                        // - X 2
                        // 1 1 0
                        GOSA_KAKUSAN_YCBCR( 1, 0, 2., 4.);
                        GOSA_KAKUSAN_YCBCR(-1, 1, 1., 4.);
                        GOSA_KAKUSAN_YCBCR( 0, 1, 1., 4.);
                        break;
                    case DITHERING_ATKINSON:
                        // - - X 1 1
                        // 0 1 1 1 0
                        // 0 0 1 0 0 ※合計6だが8で割る(75%拡散)
                        GOSA_KAKUSAN_YCBCR( 1, 0, 1., 8.);
                        GOSA_KAKUSAN_YCBCR( 2, 0, 1., 8.);
                        GOSA_KAKUSAN_YCBCR(-1, 1, 1., 8.);
                        GOSA_KAKUSAN_YCBCR( 0, 1, 1., 8.);
                        GOSA_KAKUSAN_YCBCR( 1, 1, 1., 8.);
                        GOSA_KAKUSAN_YCBCR( 0, 2, 1., 8.);
                        break;
                    }
                }
            }
            int count = 0;
            for(size_t i = 0; i < palette_size; ++i) {
                if(used_count[i] > 0) {
                    ++count;
                }
            }
            return count;
        }
    }

    bool grayscale(gd &img) {
        return do_grayscale(img, 0x00, 0xff, (0x7f << 24));
    }

    bool colorize(gd &img, int red, int green, int blue, int alpha) {
        img.convert_to_true_color();
        img.alpha_blending(false);
        const int width = img.width();
        const int height = img.height();
        for(int y = 0; y < height; ++y) {
            for(int x = 0; x < width; ++x) {
                const gd::color color = img.pixel_fast(x, y);
                const int r = ((color & 0xff0000) >> 16) + red;
                const int g = ((color & 0x00ff00) >>  8) + green;
                const int b = (color & 0x0000ff) + blue;
                const int a = ((color & 0x7f000000) >> 24) + alpha;
                const int new_color =
                    ((a < 0 ? 0 : a > 127 ? 127 : a) << 24) |
                    ((r < 0 ? 0 : r > 255 ? 255 : r) << 16) |
                    ((g < 0 ? 0 : g > 255 ? 255 : g) <<  8) |
                     (b < 0 ? 0 : b > 255 ? 255 : b);
                img.pixel_fast(x, y, new_color);
            }
        }
        return true;
    }

    bool binarize(gd &img, bool is_grayscaled, THRESHOLDING thresholding, DITHERING_METHOD dithering) {
        img.convert_to_true_color();
        img.alpha_blending(false);
        int hist_r[256] = {};
        int hist_g[256] = {};
        int hist_b[256] = {};
        long total_r = 0;
        long total_g = 0;
        long total_b = 0;
        int pixel_count = 0;
        const int width = img.width();
        const int height = img.height();

        // ヒストグラムを取得
        for(int y = 0; y < height; ++y) {
            for(int x = 0; x < width; ++x) {
                const gd::color color = img.pixel_fast(x, y);
                const int a = (color & 0x7f000000) >> 24;
                if(a == 0x7f) {
                    // 完全透明なのでスキップ
                    continue;
                }
                ++pixel_count;
                const int r = (color & 0xff0000) >> 16;
                ++hist_r[r];
                total_r += r;
                if(!is_grayscaled) {
                    const int g = (color & 0x00ff00) >> 8;
                    ++hist_g[g];
                    total_g += g;
                    const int b = (color & 0x0000ff);
                    ++hist_b[b];
                    total_b += b;
                }
            }
        }

        // 完全に透明な画像だったわ
        if(pixel_count < 1) {
            return true;
        }

        boost::multi_array<double, 3> gosa(boost::extents[width][height][3]);
        std::fill(gosa.origin(), gosa.origin() + gosa.size(), 0);

        // 閾値を計算
        const int threshold_r = (thresholding == THRESHOLD_HALF) ? 128 : otsu_threshold((double)total_r / (double)pixel_count, hist_r);
        const int threshold_g = is_grayscaled ? threshold_r : (thresholding == THRESHOLD_HALF) ? 128 : otsu_threshold((double)total_g / (double)pixel_count, hist_g);
        const int threshold_b = is_grayscaled ? threshold_r : (thresholding == THRESHOLD_HALF) ? 128 : otsu_threshold((double)total_b / (double)pixel_count, hist_b);
        for(int y = 0; y < height; ++y) {
            const int d = (y % 2 == 0) ? 1 : -1;
            const int x_begin = (y % 2 == 0) ? 0 : width - 1;
            const int x_end   = (y % 2 == 0) ? width : -1;
            for(int x = x_begin; x != x_end; x += d) {
                const gd::color color = img.pixel_fast(x, y);
                const int a = (color & 0x7f000000);
                const double r = static_cast<double>((color & 0xff0000) >> 16) + gosa[x][y][0];
                const double g = static_cast<double>((color & 0x00ff00) >>  8) + gosa[x][y][1];
                const double b = static_cast<double>((color & 0x0000ff)      ) + gosa[x][y][2];
                const int put_r = (r < threshold_r) ? 0x00 : 0xff;
                const int put_g = (g < threshold_g) ? 0x00 : 0xff;
                const int put_b = (b < threshold_b) ? 0x00 : 0xff;
                img.pixel_fast(x, y, a | (put_r << 16) | (put_g << 8) | put_b);
                if(a < 0x7f000000) {
                    const double gosa_r = r - static_cast<double>(put_r);
                    const double gosa_g = g - static_cast<double>(put_g);
                    const double gosa_b = b - static_cast<double>(put_b);
                    switch(dithering) {
                    case DITHERING_NONE:
                    default:
                        break;
                    case DITHERING_FLOYD_STEINBERG:
                        // - X 7
                        // 3 5 1
                        GOSA_KAKUSAN_RGB( 1, 0, 7., 16.);
                        GOSA_KAKUSAN_RGB(-1, 0, 3., 16.);
                        GOSA_KAKUSAN_RGB( 0, 0, 5., 16.);
                        GOSA_KAKUSAN_RGB( 1, 0, 1., 16.);
                        break;
                    case DITHERING_SIERRA_3LINE:
                        // - - X 5 3
                        // 2 4 5 4 2
                        // 0 2 3 2 0
                        GOSA_KAKUSAN_RGB( 1, 0, 5., 32.);
                        GOSA_KAKUSAN_RGB( 2, 0, 3., 32.);
                        GOSA_KAKUSAN_RGB(-2, 1, 2., 32.);
                        GOSA_KAKUSAN_RGB(-1, 1, 4., 32.);
                        GOSA_KAKUSAN_RGB( 0, 1, 5., 32.);
                        GOSA_KAKUSAN_RGB( 1, 1, 4., 32.);
                        GOSA_KAKUSAN_RGB( 2, 1, 2., 32.);
                        GOSA_KAKUSAN_RGB(-1, 2, 2., 32.);
                        GOSA_KAKUSAN_RGB( 0, 2, 3., 32.);
                        GOSA_KAKUSAN_RGB( 1, 2, 2., 32.);
                        break;
                    case DITHERING_SIERRA_2LINE:
                        // - - X 4 3
                        // 1 2 3 2 1
                        GOSA_KAKUSAN_RGB( 1, 0, 4., 16.);
                        GOSA_KAKUSAN_RGB( 2, 0, 3., 16.);
                        GOSA_KAKUSAN_RGB(-2, 1, 1., 16.);
                        GOSA_KAKUSAN_RGB(-1, 1, 2., 16.);
                        GOSA_KAKUSAN_RGB( 0, 1, 3., 16.);
                        GOSA_KAKUSAN_RGB( 1, 1, 2., 16.);
                        GOSA_KAKUSAN_RGB( 2, 1, 1., 16.);
                        break;
                    case DITHERING_SIERRA_LITE:
                        // - X 2
                        // 1 1 0
                        GOSA_KAKUSAN_RGB( 1, 0, 2., 4.);
                        GOSA_KAKUSAN_RGB(-1, 1, 1., 4.);
                        GOSA_KAKUSAN_RGB( 0, 1, 1., 4.);
                        break;
                    case DITHERING_ATKINSON:
                        // - - X 1 1
                        // 0 1 1 1 0
                        // 0 0 1 0 0 ※合計6だが8で割る(75%拡散)
                        GOSA_KAKUSAN_RGB( 1, 0, 1., 8.);
                        GOSA_KAKUSAN_RGB( 2, 0, 1., 8.);
                        GOSA_KAKUSAN_RGB(-1, 1, 1., 8.);
                        GOSA_KAKUSAN_RGB( 0, 1, 1., 8.);
                        GOSA_KAKUSAN_RGB( 1, 1, 1., 8.);
                        GOSA_KAKUSAN_RGB( 0, 2, 1., 8.);
                        break;
                    }
                }
            }
        }
        return true;
    }

    bool websafe(gd &img, DITHERING_METHOD dither) {
        img.convert_to_true_color();
        img.alpha_blending(false);
        const int width = img.width();
        const int height = img.height();
        boost::multi_array<double, 3> gosa(boost::extents[width][height][3]);
        std::fill(gosa.origin(), gosa.origin() + gosa.size(), 0); // 要らない？

        for(int y = 0; y < height; ++y) {
            const int d = (y % 2 == 0) ? 1 : -1;
            const int x_begin = (y % 2 == 0) ? 0 : width - 1;
            const int x_end   = (y % 2 == 0) ? width : -1;
            for(int x = x_begin; x != x_end; x += d) {
                const gd::color color = img.pixel_fast(x, y);
                const int a = (color & 0x7f000000);
                const double r = static_cast<double>((color & 0x00ff0000) >> 16) + gosa[x][y][0];
                const double g = static_cast<double>((color & 0x0000ff00) >>  8) + gosa[x][y][1];
                const double b = static_cast<double>((color & 0x000000ff)      ) + gosa[x][y][2];
                const int web_r = convert_to_websafe_color(static_cast<int>(round(r)));
                const int web_g = convert_to_websafe_color(static_cast<int>(round(g)));
                const int web_b = convert_to_websafe_color(static_cast<int>(round(b)));
                if(a < 0x7f000000) {
                    const double gosa_r = r - static_cast<double>(web_r);
                    const double gosa_g = g - static_cast<double>(web_g);
                    const double gosa_b = b - static_cast<double>(web_b);
                    switch(dither) {
                    case DITHERING_NONE:
                    default:
                        break;
                    case DITHERING_FLOYD_STEINBERG:
                        // - X 7
                        // 3 5 1
                        GOSA_KAKUSAN_RGB( 1, 0, 7., 16.);
                        GOSA_KAKUSAN_RGB(-1, 0, 3., 16.);
                        GOSA_KAKUSAN_RGB( 0, 0, 5., 16.);
                        GOSA_KAKUSAN_RGB( 1, 0, 1., 16.);
                        break;
                    case DITHERING_SIERRA_3LINE:
                        // - - X 5 3
                        // 2 4 5 4 2
                        // 0 2 3 2 0
                        GOSA_KAKUSAN_RGB( 1, 0, 5., 32.);
                        GOSA_KAKUSAN_RGB( 2, 0, 3., 32.);
                        GOSA_KAKUSAN_RGB(-2, 1, 2., 32.);
                        GOSA_KAKUSAN_RGB(-1, 1, 4., 32.);
                        GOSA_KAKUSAN_RGB( 0, 1, 5., 32.);
                        GOSA_KAKUSAN_RGB( 1, 1, 4., 32.);
                        GOSA_KAKUSAN_RGB( 2, 1, 2., 32.);
                        GOSA_KAKUSAN_RGB(-1, 2, 2., 32.);
                        GOSA_KAKUSAN_RGB( 0, 2, 3., 32.);
                        GOSA_KAKUSAN_RGB( 1, 2, 2., 32.);
                        break;
                    case DITHERING_SIERRA_2LINE:
                        // - - X 4 3
                        // 1 2 3 2 1
                        GOSA_KAKUSAN_RGB( 1, 0, 4., 16.);
                        GOSA_KAKUSAN_RGB( 2, 0, 3., 16.);
                        GOSA_KAKUSAN_RGB(-2, 1, 1., 16.);
                        GOSA_KAKUSAN_RGB(-1, 1, 2., 16.);
                        GOSA_KAKUSAN_RGB( 0, 1, 3., 16.);
                        GOSA_KAKUSAN_RGB( 1, 1, 2., 16.);
                        GOSA_KAKUSAN_RGB( 2, 1, 1., 16.);
                        break;
                    case DITHERING_SIERRA_LITE:
                        // - X 2
                        // 1 1 0
                        GOSA_KAKUSAN_RGB( 1, 0, 2., 4.);
                        GOSA_KAKUSAN_RGB(-1, 1, 1., 4.);
                        GOSA_KAKUSAN_RGB( 0, 1, 1., 4.);
                        break;
                    case DITHERING_ATKINSON:
                        // - - X 1 1
                        // 0 1 1 1 0
                        // 0 0 1 0 0 ※合計6だが8で割る(75%拡散)
                        GOSA_KAKUSAN_RGB( 1, 0, 1., 8.);
                        GOSA_KAKUSAN_RGB( 2, 0, 1., 8.);
                        GOSA_KAKUSAN_RGB(-1, 1, 1., 8.);
                        GOSA_KAKUSAN_RGB( 0, 1, 1., 8.);
                        GOSA_KAKUSAN_RGB( 1, 1, 1., 8.);
                        GOSA_KAKUSAN_RGB( 0, 2, 1., 8.);
                        break;
                    }
                }
                img.pixel_fast(x, y, a | (web_r << 16) | (web_g << 8) | web_b);
            }
        }
        return true;
    }

    bool famicom(gd &img, DITHERING_METHOD dither) {
        const int fixed_palette_y_cb_cr[] = {
            0x008080, 0x1bde71, 0x24fe6c, 0x32975e, 0x3462bc, 0x34635b, 0x346394, 0x3774d1,
            0x3c5ecd, 0x3d5d54, 0x3fa9bf, 0x46584e, 0x47c582, 0x5185eb, 0x57df47, 0x5ec1db,
            0x63483a, 0x63973c, 0x6ad43a, 0x6b43e5, 0x6bd683, 0x6c4333, 0x6c6a34, 0x7c3aa2,
            0x7c43ca, 0x7c8080, 0x808080, 0x93bf67, 0x9485cb, 0x98ba86, 0x9d5ac3, 0xa5534b,
            0xa5b53b, 0xa79f0d, 0xb342b5, 0xb4aab6, 0xb619af, 0xb8aa02, 0xc26b38, 0xc28080,
            0xc680a8, 0xc7a07b, 0xcc1b72, 0xd09a8b, 0xda4a98, 0xda959b, 0xdb6a93, 0xdb975e,
            0xe3457a, 0xe36b65, 0xe75f93, 0xe87b63, 0xed8a8d, 0xff8080,
        };
        const size_t fixed_palette_count = sizeof(fixed_palette_y_cb_cr) / sizeof(fixed_palette_y_cb_cr[0]);

        if(!fill_background(img, 0xffffff)) {
            return false;
        }
        // それっぽさを出すために解像度を落とす
        img.resize_fit((img.width() + 1) / 2, (img.height() + 1) / 2);

        // 画像を破壊しないためにコピーを作る
        gd img_tmp(img.width(), img.height());
        img_tmp.alpha(false, true);
        img_tmp.copy(img, 0, 0, 0, 0, img.width(), img.height());

        // 1 パス目: とりあえず変換する
        std::vector<int> palette_use_count(fixed_palette_count);
        const int color_count = famicom_convert(img_tmp, dither, fixed_palette_y_cb_cr, fixed_palette_count, &palette_use_count[0]);
        if(color_count <= 25) {
            // 1 パスで同時発色可能数に収まった
            img.swap(img_tmp);
        } else {
            // 2 パス目: 25 色以内に抑える

            // 使われた回数が多い順に並び替える
            std::vector<std::pair<size_t, int> > counts;
            for(size_t i = 0; i < fixed_palette_count; ++i) {
                counts.push_back(std::pair<size_t, int>(i, palette_use_count[i]));
            }
            std::sort(counts.begin(), counts.end(), boost::bind(&std::pair<size_t, int>::second, _1) > boost::bind(&std::pair<size_t, int>::second, _2));

            // 多い方から 25 色取得する
            int palette[25] = {};
            {
                size_t i;
                std::vector<std::pair<size_t, int> >::iterator it;
                for(i = 0, it = counts.begin(); i < 25; ++i, ++it) {
                    palette[i] = fixed_palette_y_cb_cr[it->first];
                }
            }

            palette_use_count.resize(25);
            famicom_convert(img, dither, palette, 25, &palette_use_count[0]);
        }

        // 解像度を下げたので元（とほとんど同じ）サイズに変更する
        // GD の拡大が nearest neighbor なことに依存している
        img.resize_fit(img.width() * 2, img.height() * 2);
        return true;
    }

    bool gameboy(gd &img, bool scale, DITHERING_METHOD dither) {
        const int    palette_y_cb_cr[4] = { 0x3a8080, 0x6b8080, 0xb08080, 0xde8080 };
        const size_t palette_size = 4;
        int palette_use_count[4] = {};

        // 出力色が完全な白や黒でないので調整する
        if(!do_grayscale(img, 28, 236, 0xffffff)) {
            return false;
        }

        if(scale) {
            img.resize_fit(160, 144);
        }

        // ファミコンパレットへの変換を流用して変換する
        famicom_convert(img, dither, palette_y_cb_cr, palette_size, palette_use_count);

        // パレット置換
        const int width = img.width();
        const int height = img.height();
        for(int y = 0; y < height; ++y) {
            for(int x = 0; x < width; ++x) {
                const gd::color color = img.pixel_fast(x, y);
                const int c1 = (color & 0x0000ff);
                const int c2 = (c1 < 0x52) ? 0x204631 : (c1 < 0x8d) ? 0x527f39 : (c1 < 0xc7) ? 0xaec440 : 0xd7e894;
                img.pixel_fast(x, y, c2);
            }
        }

        if(scale) {
            img.resize_fit(480, 432);
        }

        return true;
    }

    bool virtualboy(gd &img, bool scale, DITHERING_METHOD dither) {
        const int    palette_y_cb_cr[4] = { 0x008080, 0x558080, 0xaa8080, 0xff8080 };
        const size_t palette_size = 4;
        int palette_use_count[4] = {};

        if(!do_grayscale(img, 0, 255, 0xffffff)) {
            return false;
        }
        if(scale) {
            img.resize_fit(384, 288);   // 4:3 で切り出す
            img.resize_force(384, 224); // ピクセル数を整合する
        }
        // ファミコンパレットへの変換を流用して変換する
        famicom_convert(img, dither, palette_y_cb_cr, palette_size, palette_use_count);

        // パレット置換
        const int palette_red[4] = { 0x170515, 0x5c020a, 0xa20000, 0xe70000 };
        const int width = img.width();
        const int height = img.height();
        for(int y = 0; y < height; ++y) {
            for(int x = 0; x < width; ++x) {
                const gd::color color = img.pixel_fast(x, y);
                const int c1 = (color & 0x0000ff);
                const int c2 = palette_red[c1 / 85];
                img.pixel_fast(x, y, c2);
            }
        }

        if(scale) {
            img.resize_force(384, 288); // 歪ませたので戻す
        }

        gaussian_blur(img);

        return true;
    }

    bool negate(gd &img) {
        img.convert_to_true_color();
        img.alpha_blending(false);
        const int width = img.width();
        const int height = img.height();
        for(int y = 0; y < height; ++y) {
            for(int x = 0; x < width; ++x) {
                const gd::color color = img.pixel_fast(x, y);
                const int a = (color & 0x7f000000);
                const int r = 255 - ((color & 0x00ff0000) >> 16);
                const int g = 255 - ((color & 0x0000ff00) >> 8);
                const int b = 255 - ((color & 0x000000ff));
                img.pixel_fast(x, y, a | (r << 16) | (g << 8) | b);
            }
        }
        return true;
    }

    bool pixelate(gd &img, int size) {
        if(size < 1) {
            std::cerr << "モザイクのサイズは 1 以上である必要があります" << std::endl;
            return false;
        }
        img.convert_to_true_color();
        const int width = img.width();
        const int height = img.height();
        const int blocks_x = (width + size - 1) / size;
        const int blocks_y = (height + size - 1) / size;
        gd out(width, height);
        out.alpha_blending(false);
        for(int block_y = 0; block_y < blocks_y; ++block_y) {
            const int min_y = block_y * size;
            const int max_y = std::min((block_y + 1) * size, height);
            for(int block_x = 0; block_x < blocks_x; ++block_x) {
                const int min_x = block_x * size;
                const int max_x = std::min((block_x + 1) * size, width);
                int total_r = 0, total_g = 0, total_b = 0, total_a = 0;
                int pixel_count = 0;
                for(int y = min_y; y < max_y; ++y) {
                    for(int x = min_x; x < max_x; ++x) {
                        const gd::color color = img.pixel_fast(x, y);
                        const int a = (color & 0x7f000000) >> 24;
                        if(a == 0x7f) {
                            // 完全透明
                            continue;
                        }
                        const int r = (color & 0xff0000) >> 16;
                        const int g = (color & 0x00ff00) >> 8;
                        const int b = (color & 0x0000ff);
                        total_r += r;
                        total_g += g;
                        total_b += b;
                        total_a += a;
                        ++pixel_count;
                    }
                }
                const int fill_color =
                    (pixel_count < 1)
                        ? 0x7fffffff
                        : ((static_cast<int>((double)total_a / (double)pixel_count + 0.5) << 24) |
                           (static_cast<int>((double)total_r / (double)pixel_count + 0.5) << 16) |
                           (static_cast<int>((double)total_g / (double)pixel_count + 0.5) <<  8) |
                           (static_cast<int>((double)total_b / (double)pixel_count + 0.5)));
                out.fill_rect(min_x, min_y, max_x, max_y, fill_color);
            }
        }
        img.swap(out);
        return true;
    }

    bool emboss(gd &img) {
        const double filter[3][3] = {
            { 1.5, 0.0, 0.0},
            { 0.0, 0.0, 0.0},
            { 0.0, 0.0,-1.5}
        };
        return fill_background(img, 0xffffff) && grayscale(img) && apply_operator_3x3(img, filter, 1, 127);
    }

    bool gaussian_blur(gd &img) {
        const double filter[3][3] = {
            { 1./16., 2./16., 1./16. },
            { 2./16., 4./16., 2./16. },
            { 1./16., 2./16., 1./16. }
        };
        return apply_operator_3x3(img, filter, 1, 0);
    }

    bool sharpen(gd &img) {
        const double filter[3][3] = {
            { -1., -1., -1. },
            { -1.,  9., -1. },
            { -1., -1., -1. }
        };
        return apply_operator_3x3(img, filter, 1, 0);
    }

    bool edge(gd &img) {
        const double sobel_h[3][3] = {
            { 1, 0, -1 },
            { 2, 0, -2 },
            { 1, 0, -1 }
        };
        const double sobel_v[3][3] = {
            {  1,  2,  1 },
            {  0,  0,  0 },
            { -1, -2, -1 }
        };

        if(!fill_background(img, 0xffffff) || !grayscale(img)) {
            return false;
        }

        // ぼかす
        if(!gaussian_blur(img) || !gaussian_blur(img) || !gaussian_blur(img)) {
            return false;
        }
        const int width = img.width();
        const int height = img.height();

        // 勾配検出
        {
            gd horizontal(width, height);
            horizontal.alpha(false, false);
            horizontal.copy(img, 0, 0, 0, 0, width, height);
            if(!apply_operator_3x3(horizontal, sobel_h, 1, 127)) {
                return false;
            }
            gd vertical(width, height);
            vertical.alpha(false, false);
            vertical.copy(img, 0, 0, 0, 0, width, height);
            if(!apply_operator_3x3(vertical, sobel_v, 1, 127)) {
                return false;
            }
            for(int y = 0; y < height; ++y) {
                for(int x = 0; x < width; ++x) {
                    const double h = std::min(1.0, static_cast<double>((horizontal.pixel(x, y) & 0xff) - 127) / 127.0);
                    const double v = std::min(1.0, static_cast<double>((vertical.pixel(x, y) & 0xff) - 127) / 127.0);
                    const int edge = std::max(0, std::min(255, static_cast<int>(sqrt(h * h + v * v) * 255.0 + 0.5)));
                    const double theta = atan2(h, v) * 180.0 / M_PI;
                    const int direction_code =
                        (theta < 22.5) ? 1 :                // 右
                        (theta < 22.5 + 45.0 * 1) ? 2 :     // 右上
                        (theta < 22.5 + 45.0 * 2) ? 3 :     // 上
                        (theta < 22.5 + 45.0 * 3) ? 4 : 5;  // 左上・左 180°を超えることはない…はず
                    img.pixel(x, y, (direction_code << 8) | edge);
                }
            }
        }

        // 細線化（と後の処理のためにヒストグラム作成）
        int hist[256] = {};
        long long sum = 0;
        for(int y = 0; y < height; ++y) {
            for(int x = 0; x < width; ++x) {
                const int direction_code = (img.pixel(x, y) & 0xff00) >> 8;
                const int c0 = img.pixel(x, y) & 0xff;
                int c1, c2;
                switch(direction_code) {
                case 1: case 5: // 左右
                    c1 = img.pixel_safe(x + 1, y) & 0xff;
                    c2 = img.pixel_safe(x - 1, y) & 0xff;
                    break;
                case 2: // 右上
                    c1 = img.pixel_safe(x + 1, y - 1) & 0xff;
                    c2 = img.pixel_safe(x - 1, y + 1) & 0xff;
                    break;
                case 3: // 上
                    c1 = img.pixel_safe(x, y - 1) & 0xff;
                    c2 = img.pixel_safe(x, y + 1) & 0xff;
                    break;
                case 4: // 左上
                    c1 = img.pixel_safe(x - 1, y - 1) & 0xff;
                    c2 = img.pixel_safe(x + 1, y + 1) & 0xff;
                    break;
                default:
                    assert(false);
                    return false;
                }
                if(c0 < c1 || c0 < c2) {
                    img.pixel(x, y, 0x000000);
                    ++hist[0];
                } else {
                    img.pixel(x, y, 0x010101 * c0);
                    ++hist[c0];
                    sum += c0;
                }
            }
        }

        // 閾値(HI)を計算
        const int hi_threshold = otsu_threshold((double)sum / ((double)width * (double)height), hist);

        // 閾値(LO)を計算
        sum = 0;
        for(int i = 0; i < 256; ++i) {
            if(i < hi_threshold) {
                sum += hist[i] * i;
            } else if(i == hi_threshold) {
                // nothing to do.
            } else {
                // 閾値(HI)以上の画素は全部閾値(HI)と見なす
                sum += hist[i] * hi_threshold;
                hist[hi_threshold] += hist[i];
                hist[i] = 0;
            }
        }
        const int lo_threshold = otsu_threshold((double)sum / ((double)width * (double)height), hist);

        gd dst(width, height);
        dst.alpha(false, true);
        // 閾値HI以上の画素を描画
        for(int y = 0; y < height; ++y) {
            for(int x = 0; x < width; ++x) {
                dst.pixel_fast(x, y, (img.pixel_fast(x, y) & 0xff) >= hi_threshold ? 0xffffff : 0x000000);
            }
        }
        // 閾値LO以上HI未満の画素をそれなりに描画
        for(int y = 0; y < height; ++y) {
            for(int x = 0; x < width; ++x) {
                const int c = img.pixel_fast(x, y) & 0xff;
                if(lo_threshold <= c && c < hi_threshold) {
                    bool found = false;
                    for(int j = -1; j <= 1; ++j) {
                        for(int i = -1; i <= 1; ++i) {
                            if(i == 0 && j == 0) {
                                // 自分の画素を見ても仕方ない
                                continue;
                            }
                            // 近所にエッジは居る？
                            if((img.pixel_safe(x + i, y + j) & 0xff) >= hi_threshold) {
                                found = true;
                                break;
                            }
                        }
                        if(found) {
                            break;
                        }
                    }
                    if(found) {
                        dst.pixel_fast(x, y, 0xffffff);
                    }
                }
            }
        }

        img.swap(dst);

        return true;
    }

    bool rotate_fast(gd &img, int degree) {
        degree %= 360;
        assert(degree % 90 == 0);
        if(degree == 0) {
            return true;
        }
        const int src_width  = img.width();
        const int src_height = img.height();
        const int dst_width  = degree % 180 == 0 ? src_width : src_height;
        const int dst_height = degree % 180 == 0 ? src_height : src_width;
        gd dst(dst_width, dst_height);
        dst.alpha(false, true);
        dst.copy_rotated(img, dst_width / 2.0, dst_height / 2.0, 0, 0, src_width, src_height, degree);
        img.swap(dst);
        return true;
    }

    bool flip_horizontal(gd &img) {
        const int width = img.width();
        const int height = img.height();
        gd dst(width, height);
        dst.alpha(false, true);
        for(int x = 0; x < width; ++x) {
            const int x2 = width - x - 1;
            dst.copy(img, x2, 0, x, 0, 1, height);
        }
        img.swap(dst);
        return true;
    }

    bool flip_vertical(gd &img) {
        const int width = img.width();
        const int height = img.height();
        gd dst(width, height);
        dst.alpha(false, true);
        for(int y = 0; y < height; ++y) {
            const int y2 = height - y - 1;
            dst.copy(img, 0, y2, 0, y, width, 1);
        }
        img.swap(dst);
        return true;
    }

    bool fill_background(gd &img, gd::color bg) {
        const int bg_r = (bg & 0xff0000) >> 16;
        const int bg_g = (bg & 0x00ff00) >>  8;
        const int bg_b = (bg & 0x0000ff);
        img.convert_to_true_color();
        img.alpha_blending(false);
        const int width = img.width();
        const int height = img.height();
        for(int y = 0; y < height; ++y) {
            for(int x = 0; x < width; ++x) {
                const gd::color color = img.pixel_fast(x, y);
                const int a = (color & 0x7f000000) >> 24;
                if(a == 0) {
                    continue;
                }
                const double alpha = static_cast<double>(127 - a) / 127.0;
                const int org_r = (color & 0xff0000) >> 16;
                const int org_g = (color & 0x00ff00) >>  8;
                const int org_b = (color & 0x0000ff);
                const int r = static_cast<int>(org_r * alpha + bg_r * (1.0 - alpha) + 0.5);
                const int g = static_cast<int>(org_g * alpha + bg_g * (1.0 - alpha) + 0.5);
                const int b = static_cast<int>(org_b * alpha + bg_b * (1.0 - alpha) + 0.5);
                img.pixel_fast(x, y, (r << 16) | (g << 8) | b);
            }
        }
        return true;
    }
}


// Copyright (C) 2013 Massachusetts Institute of Technology, Lincoln Laboratory
// License: Boost Software License   See LICENSE.txt for the full license.
// Authors: Davis E. King (davis@dlib.net)

#include <iostream>
#include <dlib/cmd_line_parser.h>
#include <mitie/gigaword_reader.h>
#include <dlib/dir_nav.h>
#include <mitie/count_min_sketch.h>
#include <mitie/group_tokenizer.h>
#include <mitie/unigram_tokenizer.h>
#include "basic_morph.h"
#include "word_vects.h"
#include "cca_morph.h"
#include "doc_vects.h"

#include <mitie/total_word_feature_extractor.h>

#include <dlib/graph_utils_threaded.h>
#include <dlib/clustering.h>

using namespace std;
using namespace dlib;
using namespace mitie;

// ----------------------------------------------------------------------------------------

template <typename tokenizer_type>
void get_top_word_counts (
    tokenizer_type& tok,
    std::map<std::string, unsigned long>& top_words,
    const unsigned long max_top_words
)
{
    count_min_sketch counts(5000000);

    // get counts on all the words first
    std::string token;
    while (tok(token))
    {
        counts.increment(token);
    }

    std::priority_queue<std::pair<long, std::string> > best_words;

    dlib::set<std::string>::kernel_1a cur_words;

    // now pass over the data again and pick out only words with high counts
    tok.reset();
    while (tok(token))
    {
        if (cur_words.is_member(token))
            continue;

        const long hits = counts.get_count(token);
        if (best_words.size() < max_top_words || hits > -best_words.top().first)
        {
            if (best_words.size() >= max_top_words)
            {
                cur_words.destroy(best_words.top().second);
                best_words.pop();
            }

            best_words.push(std::make_pair(-hits, token));
            cur_words.add(token);
        }
    }

    // now put the best words into top_words
    top_words.clear();
    while (best_words.size() != 0)
    {
        top_words[best_words.top().second] = -best_words.top().first;
        best_words.pop();
    }
}

// ----------------------------------------------------------------------------------------

void count_words(const command_line_parser& parser);
void test(const command_line_parser& parser);
void cluster_words(const command_line_parser& parser);

int main(int argc, char** argv)
{
    try
    {
        command_line_parser parser;
        parser.add_option("h","Display this help message.");

        parser.add_option("e", "Make a total_word_feature_extractor from a folder of text files.   This option is a shortcut for executing"
                               " the following options together --count-words 200000 --word-vects --basic-morph --cca-morph.");

        parser.set_group_name("Other Options");
        parser.add_option("convert-gigaword", "Take a folder of gigaword XML documents and convert them "
            "into a regular ASCII file named <arg>.",1);
        parser.add_option("count-words", "Make a file containing the top <arg> most common words and their occurrence counts.",1);
        parser.add_option("basic-morph", "Make a word morphology extractor.");
        parser.add_option("cca-morph", "Make a CCA based word morphology extractor object as well as a total word "
                                        "feature extractor.");
        parser.add_option("word-vects", "Use CCA to create distributional word vectors.");
        parser.add_option("test", "Print out the feature vectors for the word given on the command line.");
        parser.add_option("cluster-words", "Generate word clusters based on a saved total_word_feature_extractor.");

        parser.set_group_name("Document Vector Level Features");
        parser.add_option("doc-vects", "Generate CCA based word features where we assume the important thing about a "
            "word is what other words it shows up with in the same document.  For this, we take a folder of gigaword XML files "
            "as input.");
        parser.add_option("dims", "When doing --doc-vects, make the output vectors have <arg> dimensions (default: 500).",1);

        parser.parse(argc,argv);
        parser.check_option_arg_range("count-words", 1, 1000000000);
        parser.check_option_arg_range("dims", 1, 100000);
        parser.check_sub_option("doc-vects", "dims");
        parser.check_incompatible_options("e", "word-vects");
        parser.check_incompatible_options("e", "count-words");
        parser.check_incompatible_options("e", "basic-morph");
        parser.check_incompatible_options("e", "cca-morph");

        if (parser.option("h"))
        {
            cout << "Main Usage: wordrep -e <folder of text files>\n";
            parser.print_options();
            cout << endl;
            return 0;
        }

        if (parser.option("convert-gigaword"))
        {
            ofstream fout(parser.option("convert-gigaword").argument().c_str());
            std::vector<dlib::file> files = get_files_in_directory_tree(directory(parser[0]), match_all());
            cout << "number of gigaword files found: " << files.size() << endl;
            mitie::gigaword_reader reader(files);
            std::string data;
            while (reader(data))
                fout << data << "\n\n";

        }

        if (parser.option("e"))
        {
            count_words(parser);
            word_vects(parser);
            basic_morph(parser);
            cca_morph(parser);
            return 0;
        }

        if (parser.option("count-words"))
        {
            count_words(parser);
        }

        if (parser.option("word-vects"))
        {
            word_vects(parser);
        }

        if (parser.option("basic-morph"))
        {
            basic_morph(parser);
        }

        if (parser.option("cca-morph"))
        {
            cca_morph(parser);
        }

        if (parser.option("test"))
        {
            test(parser);
        }

        if (parser.option("cluster-words"))
        {
            cluster_words(parser);
        }

        if (parser.option("doc-vects"))
        {
            make_doc_vects(parser);
        }

        return 0;
    }
    catch (std::exception& e)
    {
        cout << e.what() << endl;
        return 1;
    }
}

// ----------------------------------------------------------------------------------------

void cluster_words(const command_line_parser& parser)
{
    string classname;
    total_word_feature_extractor fe;
    deserialize("total_word_feature_extractor.dat") >> classname >> fe;

    cout << "words in dictionary: " << fe.get_num_words_in_dictionary() << endl;
    cout << "num features: " << fe.get_num_dimensions() << endl;
    const std::vector<std::string>& words = fe.get_words_in_dictionary();
    std::vector<matrix<float,0,1> > vects;
    vects.resize(words.size());
    for (unsigned long i = 0; i < vects.size(); ++i)
        fe.get_feature_vector(words[i], vects[i]);

    cout << "Making graph" << endl;
    std::vector<sample_pair> edges;
    find_k_nearest_neighbors_lsh(vects, cosine_distance(), hash_similar_angles_256(), 100, 4, edges);
    cout << "edges.size(): "<< edges.size() << endl;

    // change all the edge weights to use the formula from the unsupervised POS paper.
    for (unsigned long i = 0; i < edges.size(); ++i)
    {
        const double dist = edges[i].distance();
        double edge_weight;
        if (dist <= 0.1)
            edge_weight = 10;
        else
            edge_weight = 1.0/dist;
        edges[i] = sample_pair(edges[i].index1(), edges[i].index2(), edge_weight);
    }
    remove_short_edges(edges, 3.0);
    //remove_short_edges(edges, 5.0);
    cout << "edges.size(): "<< edges.size() << endl;

    std::vector<unsigned long> labels;
    const unsigned long num_clusters = chinese_whispers(edges, labels);
    //const unsigned long num_clusters = newman_cluster(edges, labels);
    cout << "num_clusters: "<< num_clusters << endl;
    cout << "labels.size(): "<< labels.size() << endl;

    std::vector<std::vector<std::string> > groups(num_clusters);
    for (unsigned long i = 0; i < labels.size(); ++i)
    {
        groups[labels[i]].push_back(words[i]);
    }

    ofstream fout("word_clusters.txt");
    for (unsigned long i = 0; i < groups.size(); ++i)
    {
        for (unsigned long j = 0; j < groups[i].size(); ++j)
        {
            fout << groups[i][j] << " ";
        }
        fout << "\n\n********************************************************************************\n\n";
    }

    // also save a more machine readable copy of the clusters
    ofstream fout2("word_clusters.dat", ios::binary);
    serialize(groups, fout2);
}

// ----------------------------------------------------------------------------------------

void count_words(const command_line_parser& parser)
{
    const unsigned long num_top_words = get_option(parser, "count-words", 200000);

    std::vector<dlib::file> files = get_files_in_directory_tree(directory(parser[0]), match_all());
    cout << "number of raw ASCII files found: " << files.size() << endl;

    group_tokenizer<unigram_tokenizer> tok(files);
    std::map<std::string, unsigned long> words;
    get_top_word_counts(tok, words, num_top_words);

    cout << "num words: "<< words.size() << endl;
    cout << "saving word counts to top_word_counts.dat" << endl;
    ofstream fout("top_word_counts.dat", ios::binary);
    serialize(words, fout);

    // lets also save it as a .txt file
    std::vector<std::pair<unsigned long, std::string> > temp;
    for (std::map<std::string,unsigned long>::iterator i = words.begin(); i != words.end(); ++i)
        temp.push_back(make_pair(i->second, i->first));
    std::sort(temp.begin(), temp.end());
    fout.close(); fout.clear();
    fout.open("top_words.txt");
    for (unsigned long i = 0; i < temp.size(); ++i)
    {
        fout << temp[i].first << " \t" << temp[i].second << "\n";
    }

}

// ----------------------------------------------------------------------------------------

void test(const command_line_parser& parser)
{
    string classname;
    total_word_feature_extractor fe;
    deserialize("total_word_feature_extractor.dat") >> classname >> fe;

    cout << "words in dictionary: " << fe.get_num_words_in_dictionary() << endl;
    cout << "num features: " << fe.get_num_dimensions() << endl;

    string word = parser[0];

    matrix<float,0,1> feats;
    fe.get_feature_vector(word,feats);
    cout << "feature vector: "<< trans(feats) << endl;
}

// ----------------------------------------------------------------------------------------


//  Copyright John Maddock 2007.

//  Use, modification and distribution are subject to the
//  Boost Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#include <boost/math/tools/traits.hpp>
#include <boost/static_assert.hpp>
#include <boost/math/distributions.hpp>

using namespace boost::math;

BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<double>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<int>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<bernoulli>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<beta_distribution<> >::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<binomial>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<cauchy>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<chi_squared>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<exponential>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<extreme_value>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<fisher_f>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<gamma_distribution<> >::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<lognormal>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<negative_binomial>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<normal>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<pareto>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<poisson>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<rayleigh>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<students_t>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<triangular>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<uniform>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_distribution<weibull>::value);

BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<double>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<int>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<bernoulli>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<beta_distribution<> >::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<binomial>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<cauchy>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<chi_squared>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<exponential>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<extreme_value>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<fisher_f>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<gamma_distribution<> >::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<lognormal>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<negative_binomial>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<normal>::value);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<pareto>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<poisson>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<rayleigh>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<students_t>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<triangular>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<uniform>::value == false);
BOOST_STATIC_ASSERT(::boost::math::tools::is_scaled_distribution<weibull>::value == false);


//| Copyright Inria May 2015
//| This project has received funding from the European Research Council (ERC) under
//| the European Union's Horizon 2020 research and innovation programme (grant
//| agreement No 637972) - see http://www.resibots.eu
//|
//| Contributor(s):
//|   - Jean-Baptiste Mouret (jean-baptiste.mouret@inria.fr)
//|   - Antoine Cully (antoinecully@gmail.com)
//|   - Konstantinos Chatzilygeroudis (konstantinos.chatzilygeroudis@inria.fr)
//|   - Federico Allocati (fede.allocati@gmail.com)
//|   - Vaios Papaspyros (b.papaspyros@gmail.com)
//|   - Roberto Rama (bertoski@gmail.com)
//|
//| This software is a computer library whose purpose is to optimize continuous,
//| black-box functions. It mainly implements Gaussian processes and Bayesian
//| optimization.
//| Main repository: http://github.com/resibots/limbo
//| Documentation: http://www.resibots.eu/limbo
//|
//| This software is governed by the CeCILL-C license under French law and
//| abiding by the rules of distribution of free software.  You can  use,
//| modify and/ or redistribute the software under the terms of the CeCILL-C
//| license as circulated by CEA, CNRS and INRIA at the following URL
//| "http://www.cecill.info".
//|
//| As a counterpart to the access to the source code and  rights to copy,
//| modify and redistribute granted by the license, users are provided only
//| with a limited warranty  and the software's author,  the holder of the
//| economic rights,  and the successive licensors  have only  limited
//| liability.
//|
//| In this respect, the user's attention is drawn to the risks associated
//| with loading,  using,  modifying and/or developing or reproducing the
//| software by the user in light of its specific status of free software,
//| that may mean  that it is complicated to manipulate,  and  that  also
//| therefore means  that it is reserved for developers  and  experienced
//| professionals having in-depth computer knowledge. Users are therefore
//| encouraged to load and test the software's suitability as regards their
//| requirements in conditions enabling the security of their systems and/or
//| data to be ensured and,  more generally, to use and operate it in the
//| same conditions as regards security.
//|
//| The fact that you are presently reading this means that you have had
//| knowledge of the CeCILL-C license and that you accept its terms.
//|
#ifndef LIMBO_MODEL_GP_HPP
#define LIMBO_MODEL_GP_HPP

#include <cassert>
#include <iostream>
#include <limits>
#include <vector>

#include <Eigen/Cholesky>
#include <Eigen/Core>
#include <Eigen/LU>

// Quick hack for definition of 'I' in <complex.h>
#undef I

#include <limbo/kernel/matern_five_halves.hpp>
#include <limbo/kernel/squared_exp_ard.hpp>
#include <limbo/mean/constant.hpp>
#include <limbo/mean/data.hpp>
#include <limbo/model/gp/kernel_lf_opt.hpp>
#include <limbo/model/gp/no_lf_opt.hpp>
#include <limbo/tools.hpp>

namespace limbo {
    namespace model {
        /// @ingroup model
        /// A classic Gaussian process.
        /// It is parametrized by:
        /// - a kernel function
        /// - a mean function
        /// - [optional] an optimizer for the hyper-parameters
        template <typename Params, typename KernelFunction = kernel::MaternFiveHalves<Params>, typename MeanFunction = mean::Data<Params>, typename HyperParamsOptimizer = gp::NoLFOpt<Params>>
        class GP {
        public:
            /// useful because the model might be created before knowing anything about the process
            GP() : _dim_in(-1), _dim_out(-1), _inv_kernel_updated(false) {}

            /// useful because the model might be created before having samples
            GP(int dim_in, int dim_out)
                : _dim_in(dim_in), _dim_out(dim_out), _kernel_function(dim_in), _mean_function(dim_out), _inv_kernel_updated(false) {}

            /// Compute the GP from samples and observations. This call needs to be explicit!
            void compute(const std::vector<Eigen::VectorXd>& samples,
                const std::vector<Eigen::VectorXd>& observations, bool compute_kernel = true)
            {
                assert(samples.size() != 0);
                assert(observations.size() != 0);
                assert(samples.size() == observations.size());

                if (_dim_in != samples[0].size()) {
                    _dim_in = samples[0].size();
                    _kernel_function = KernelFunction(_dim_in); // the cost of building a functor should be relatively low
                }

                if (_dim_out != observations[0].size()) {
                    _dim_out = observations[0].size();
                    _mean_function = MeanFunction(_dim_out); // the cost of building a functor should be relatively low
                }

                _samples = samples;

                _observations.resize(observations.size(), _dim_out);
                for (int i = 0; i < _observations.rows(); ++i)
                    _observations.row(i) = observations[i];

                _mean_observation = _observations.colwise().mean();

                this->_compute_obs_mean();
                if (compute_kernel)
                    this->_compute_full_kernel();
            }

            /// Do not forget to call this if you use hyper-parameters optimization!!
            void optimize_hyperparams()
            {
                _hp_optimize(*this);
            }

            /// add sample and update the GP. This code uses an incremental implementation of the Cholesky
            /// decomposition. It is therefore much faster than a call to compute()
            void add_sample(const Eigen::VectorXd& sample, const Eigen::VectorXd& observation)
            {
                if (_samples.empty()) {
                    if (_dim_in != sample.size()) {
                        _dim_in = sample.size();
                        _kernel_function = KernelFunction(_dim_in); // the cost of building a functor should be relatively low
                    }
                    if (_dim_out != observation.size()) {
                        _dim_out = observation.size();
                        _mean_function = MeanFunction(_dim_out); // the cost of building a functor should be relatively low
                    }
                }
                else {
                    assert(sample.size() == _dim_in);
                    assert(observation.size() == _dim_out);
                }

                _samples.push_back(sample);

                _observations.conservativeResize(_observations.rows() + 1, _dim_out);
                _observations.bottomRows<1>() = observation.transpose();

                _mean_observation = _observations.colwise().mean();

                this->_compute_obs_mean();
                this->_compute_incremental_kernel();
            }

            /**
             \\rst
             return :math:`\mu`, :math:`\sigma^2` (un-normalized). If there is no sample, return the value according to the mean function. Using this method instead of separate calls to mu() and sigma() is more efficient because some computations are shared between mu() and sigma().
             \\endrst
            */
            std::tuple<Eigen::VectorXd, double> query(const Eigen::VectorXd& v) const
            {
                if (_samples.size() == 0)
                    return std::make_tuple(_mean_function(v, *this),
                        _kernel_function(v, v) + _kernel_function.noise());

                Eigen::VectorXd k = _compute_k(v);
                return std::make_tuple(_mu(v, k), _sigma(v, k) + _kernel_function.noise());
            }

            /**
             \\rst
             return :math:`\mu` (un-normalized). If there is no sample, return the value according to the mean function.
             \\endrst
            */
            Eigen::VectorXd mu(const Eigen::VectorXd& v) const
            {
                if (_samples.size() == 0)
                    return _mean_function(v, *this);
                return _mu(v, _compute_k(v));
            }

            /**
             \\rst
             return :math:`\sigma^2` (un-normalized). If there is no sample, return the max :math:`\sigma^2`.
             \\endrst
            */
            double sigma(const Eigen::VectorXd& v) const
            {
                if (_samples.size() == 0)
                    return _kernel_function(v, v) + _kernel_function.noise();
                return _sigma(v, _compute_k(v)) + _kernel_function.noise();
            }

            /// return the number of dimensions of the input
            int dim_in() const
            {
                assert(_dim_in != -1); // need to compute first!
                return _dim_in;
            }

            /// return the number of dimensions of the output
            int dim_out() const
            {
                assert(_dim_out != -1); // need to compute first!
                return _dim_out;
            }

            const KernelFunction& kernel_function() const { return _kernel_function; }

            KernelFunction& kernel_function() { return _kernel_function; }

            const MeanFunction& mean_function() const { return _mean_function; }

            MeanFunction& mean_function() { return _mean_function; }

            /// return the maximum observation (only call this if the output of the GP is of dimension 1)
            Eigen::VectorXd max_observation() const
            {
                if (_observations.cols() > 1)
                    std::cout << "WARNING max_observation with multi dimensional "
                                 "observations doesn't make sense"
                              << std::endl;
                return tools::make_vector(_observations.maxCoeff());
            }

            /// return the mean observation (only call this if the output of the GP is of dimension 1)
            Eigen::VectorXd mean_observation() const
            {
                assert(_dim_out > 0);
                return _samples.size() > 0 ? _mean_observation
                                           : Eigen::VectorXd::Zero(_dim_out);
            }

            const Eigen::MatrixXd& mean_vector() const { return _mean_vector; }

            const Eigen::MatrixXd& obs_mean() const { return _obs_mean; }

            /// return the number of samples used to compute the GP
            int nb_samples() const { return _samples.size(); }

            ///  recomputes the GP
            void recompute(bool update_obs_mean = true, bool update_full_kernel = true)
            {
                assert(!_samples.empty());

                if (update_obs_mean)
                    this->_compute_obs_mean();

                if (update_full_kernel)
                    this->_compute_full_kernel();
                else
                    this->_compute_alpha();
            }

            void compute_inv_kernel()
            {
                size_t n = _obs_mean.rows();
                // K^{-1} using Cholesky decomposition
                _inv_kernel = Eigen::MatrixXd::Identity(n, n);

                _matrixL.template triangularView<Eigen::Lower>().solveInPlace(_inv_kernel);
                _matrixL.template triangularView<Eigen::Lower>().transpose().solveInPlace(_inv_kernel);

                _inv_kernel_updated = true;
            }

            /// compute and return the log likelihood
            double compute_log_lik()
            {
                size_t n = _obs_mean.rows();

                // --- cholesky ---
                // see:
                // http://xcorr.net/2008/06/11/log-determinant-of-positive-definite-matrices-in-matlab/
                long double logdet = 2 * _matrixL.diagonal().array().log().sum();

                double a = (_obs_mean.transpose() * _alpha)
                               .trace(); // generalization for multi dimensional observation

                _log_lik = -0.5 * a - 0.5 * logdet - 0.5 * n * std::log(2 * M_PI);

                return _log_lik;
            }

            /// compute and return the gradient of the log likelihood wrt to the kernel parameters
            Eigen::VectorXd compute_kernel_grad_log_lik()
            {
                size_t n = _obs_mean.rows();

                // compute K^{-1} only if needed
                if (!_inv_kernel_updated) {
                    compute_inv_kernel();
                }
                Eigen::MatrixXd w = _inv_kernel;

                // alpha * alpha.transpose() - K^{-1}
                w = _alpha * _alpha.transpose() - w;

                // only compute half of the matrix (symmetrical matrix)
                Eigen::VectorXd grad = Eigen::VectorXd::Zero(_kernel_function.h_params_size());
                for (size_t i = 0; i < n; ++i) {
                    for (size_t j = 0; j <= i; ++j) {
                        Eigen::VectorXd g = _kernel_function.grad(_samples[i], _samples[j], i, j);
                        if (i == j)
                            grad += w(i, j) * g * 0.5;
                        else
                            grad += w(i, j) * g;
                    }
                }

                return grad;
            }

            /// compute and return the gradient of the log likelihood wrt to the mean parameters
            Eigen::VectorXd compute_mean_grad_log_lik()
            {
                size_t n = _obs_mean.rows();

                // compute K^{-1} only if needed
                if (!_inv_kernel_updated) {
                    compute_inv_kernel();
                }

                Eigen::VectorXd grad = Eigen::VectorXd::Zero(_mean_function.h_params_size());
                for (int i_obs = 0; i_obs < _dim_out; ++i_obs)
                    for (size_t n_obs = 0; n_obs < n; n_obs++) {
                        grad += _obs_mean.col(i_obs).transpose() * _inv_kernel.col(n_obs) * _mean_function.grad(_samples[n_obs], *this).row(i_obs);
                    }

                return grad;
            }

            /// return the likelihood (do not compute it -- return last computed)
            double get_log_lik() const { return _log_lik; }

            /// set the log likelihood (e.g. computed from outside)
            void set_log_lik(double log_lik) { _log_lik = log_lik; }

            /// compute and return the log probability of LOO CV
            double compute_log_loo_cv()
            {
                // compute K^{-1} only if needed
                if (!_inv_kernel_updated) {
                    compute_inv_kernel();
                }

                Eigen::VectorXd inv_diag = _inv_kernel.diagonal().array().inverse();

                _log_loo_cv = (((-0.5 * (_alpha.array().square().array().colwise() * inv_diag.array())).array().colwise() - 0.5 * inv_diag.array().log().array()) - 0.5 * std::log(2 * M_PI)).colwise().sum().sum();

                return _log_loo_cv;
            }

            /// compute and return the gradient of the log probability of LOO CV wrt to the kernel parameters
            Eigen::VectorXd compute_kernel_grad_log_loo_cv()
            {
                size_t n = _obs_mean.rows();
                size_t n_params = _kernel_function.h_params_size();

                // compute K^{-1} only if needed
                if (!_inv_kernel_updated) {
                    compute_inv_kernel();
                }

                Eigen::VectorXd grad = Eigen::VectorXd::Zero(n_params);
                Eigen::MatrixXd grads = Eigen::MatrixXd::Zero(n_params, _dim_out);

                // only compute half of the matrix (symmetrical matrix)
                // TO-DO: Make it better
                std::vector<std::vector<Eigen::VectorXd>> full_dk;
                for (size_t i = 0; i < n; i++) {
                    full_dk.push_back(std::vector<Eigen::VectorXd>());
                    for (size_t j = 0; j <= i; j++)
                        full_dk[i].push_back(_kernel_function.grad(_samples[i], _samples[j], i, j));
                    for (size_t j = i + 1; j < n; j++)
                        full_dk[i].push_back(Eigen::VectorXd::Zero(n_params));
                }
                for (size_t i = 0; i < n; i++)
                    for (size_t j = 0; j < i; ++j)
                        full_dk[j][i] = full_dk[i][j];

                Eigen::VectorXd inv_diag = _inv_kernel.diagonal().array().inverse();

                for (int j = 0; j < grad.size(); j++) {
                    Eigen::MatrixXd dKdTheta_j = Eigen::MatrixXd::Zero(n, n);
                    for (size_t i = 0; i < n; i++) {
                        for (size_t k = 0; k < n; k++)
                            dKdTheta_j(i, k) = full_dk[i][k](j);
                    }
                    Eigen::MatrixXd Zeta_j = _inv_kernel * dKdTheta_j;
                    Eigen::MatrixXd Zeta_j_alpha = Zeta_j * _alpha;
                    Eigen::MatrixXd Zeta_j_K = Zeta_j * _inv_kernel;

                    grads.row(j) = ((_alpha.array() * Zeta_j_alpha.array() - 0.5 * ((1. + _alpha.array().square().array().colwise() * inv_diag.array()).array().colwise() * Zeta_j_K.diagonal().array())).array().colwise() * inv_diag.array()).colwise().sum();

                    // for (size_t i = 0; i < n; i++)
                    //     grads.row(j).array() += (_alpha.row(i).array() * Zeta_j_alpha.row(i).array() - 0.5 * (1. + _alpha.row(i).array().square() / _inv_kernel.diagonal()(i)) * Zeta_j_K.diagonal()(i)) / _inv_kernel.diagonal()(i);
                }

                grad = grads.rowwise().sum();

                return grad;
            }

            /// return the LOO-CV log probability (do not compute it -- return last computed)
            double get_log_loo_cv() const { return _log_loo_cv; }

            /// set the LOO-CV log probability (e.g. computed from outside)
            void set_log_loo_cv(double log_loo_cv) { _log_loo_cv = log_loo_cv; }

            /// LLT matrix (from Cholesky decomposition)
            const Eigen::MatrixXd& matrixL() const { return _matrixL; }

            const Eigen::MatrixXd& alpha() const { return _alpha; }

            /// return the list of samples
            const std::vector<Eigen::VectorXd>& samples() const { return _samples; }

            /// return the list of observations
            std::vector<Eigen::VectorXd> observations() const
            {
                std::vector<Eigen::VectorXd> observations;
                for (int i = 0; i < _observations.rows(); i++) {
                    observations.push_back(_observations.row(i));
                }

                return observations;
            }

            /// return the observations (in matrix form)
            /// (NxD), where N is the number of points and D is the dimension output
            const Eigen::MatrixXd& observations_matrix() const
            {
                return _observations;
            }

            bool inv_kernel_computed() { return _inv_kernel_updated; }

            /// save the parameters and the data for the GP to the archive (text or binary)
            template <typename A>
            void save(const std::string& directory) const
            {
                A archive(directory);
                save(archive);
            }

            /// save the parameters and the data for the GP to the archive (text or binary)
            template <typename A>
            void save(const A& archive) const
            {
                if (_kernel_function.h_params_size() > 0) {
                    archive.save(_kernel_function.h_params(), "kernel_params");
                }
                if (_mean_function.h_params_size() > 0) {
                    archive.save(_mean_function.h_params(), "mean_params");
                }
                archive.save(_samples, "samples");
                archive.save(_observations, "observations");
                archive.save(_matrixL, "matrixL");
                archive.save(_alpha, "alpha");
            }

            /// load the parameters and the data for the GP from the archive (text or binary)
            /// if recompute is true, we do not read the kernel matrix
            /// but we recompute it given the data and the hyperparameters
            template <typename A>
            void load(const std::string& directory, bool recompute = true)
            {
                A archive(directory);
                load(archive, recompute);
            }

            /// load the parameters and the data for the GP from the archive (text or binary)
            /// if recompute is true, we do not read the kernel matrix
            /// but we recompute it given the data and the hyperparameters
            template <typename A>
            void load(const A& archive, bool recompute = true)
            {
                _samples.clear();
                archive.load(_samples, "samples");

                archive.load(_observations, "observations");

                _dim_in = _samples[0].size();
                _kernel_function = KernelFunction(_dim_in);

                if (_kernel_function.h_params_size() > 0) {
                    Eigen::VectorXd h_params;
                    archive.load(h_params, "kernel_params");
                    assert(h_params.size() == (int)_kernel_function.h_params_size());
                    _kernel_function.set_h_params(h_params);
                }

                _dim_out = _observations.cols();
                _mean_function = MeanFunction(_dim_out);

                if (_mean_function.h_params_size() > 0) {
                    Eigen::VectorXd h_params;
                    archive.load(h_params, "mean_params");
                    assert(h_params.size() == (int)_mean_function.h_params_size());
                    _mean_function.set_h_params(h_params);
                }

                _mean_observation = _observations.colwise().mean();

                if (recompute)
                    this->recompute(true, true);
                else {
                    archive.load(_matrixL, "matrixL");
                    archive.load(_alpha, "alpha");
                }
            }

        protected:
            int _dim_in;
            int _dim_out;

            KernelFunction _kernel_function;
            MeanFunction _mean_function;

            std::vector<Eigen::VectorXd> _samples;
            Eigen::MatrixXd _observations;
            Eigen::MatrixXd _mean_vector;
            Eigen::MatrixXd _obs_mean;

            Eigen::MatrixXd _alpha;
            Eigen::VectorXd _mean_observation;

            Eigen::MatrixXd _kernel, _inv_kernel;

            Eigen::MatrixXd _matrixL;

            double _log_lik, _log_loo_cv;
            bool _inv_kernel_updated;

            HyperParamsOptimizer _hp_optimize;

            void _compute_obs_mean()
            {
                assert(!_samples.empty());
                _mean_vector.resize(_samples.size(), _dim_out);
                for (int i = 0; i < _mean_vector.rows(); i++) {
                    assert(_samples[i].cols() == 1);
                    assert(_samples[i].rows() != 0);
                    assert(_samples[i].rows() == _dim_in);
                    _mean_vector.row(i) = _mean_function(_samples[i], *this);
                }
                _obs_mean = _observations - _mean_vector;
            }

            void _compute_full_kernel()
            {
                size_t n = _samples.size();
                _kernel.resize(n, n);

                // O(n^2) [should be negligible]
                for (size_t i = 0; i < n; i++)
                    for (size_t j = 0; j <= i; ++j)
                        _kernel(i, j) = _kernel_function(_samples[i], _samples[j], i, j);

                for (size_t i = 0; i < n; i++)
                    for (size_t j = 0; j < i; ++j)
                        _kernel(j, i) = _kernel(i, j);

                // O(n^3)
                _matrixL = Eigen::LLT<Eigen::MatrixXd>(_kernel).matrixL();

                this->_compute_alpha();

                // notify change of kernel
                _inv_kernel_updated = false;
            }

            void _compute_incremental_kernel()
            {
                // Incremental LLT
                // This part of the code is inspired from the Bayesopt Library (cholesky_add_row function).
                // However, the mathematical foundations can be easily retrieved by detailing the equations of the
                // extended L matrix that produces the desired kernel.

                size_t n = _samples.size();
                _kernel.conservativeResize(n, n);

                for (size_t i = 0; i < n; ++i) {
                    _kernel(i, n - 1) = _kernel_function(_samples[i], _samples[n - 1], i, n - 1);
                    _kernel(n - 1, i) = _kernel(i, n - 1);
                }

                _matrixL.conservativeResizeLike(Eigen::MatrixXd::Zero(n, n));

                double L_j;
                for (size_t j = 0; j < n - 1; ++j) {
                    L_j = _kernel(n - 1, j) - (_matrixL.block(j, 0, 1, j) * _matrixL.block(n - 1, 0, 1, j).transpose())(0, 0);
                    _matrixL(n - 1, j) = (L_j) / _matrixL(j, j);
                }

                L_j = _kernel(n - 1, n - 1) - (_matrixL.block(n - 1, 0, 1, n - 1) * _matrixL.block(n - 1, 0, 1, n - 1).transpose())(0, 0);
                _matrixL(n - 1, n - 1) = sqrt(L_j);

                this->_compute_alpha();

                // notify change of kernel
                _inv_kernel_updated = false;
            }

            void _compute_alpha()
            {
                // alpha = K^{-1} * this->_obs_mean;
                Eigen::TriangularView<Eigen::MatrixXd, Eigen::Lower> triang = _matrixL.template triangularView<Eigen::Lower>();
                _alpha = triang.solve(_obs_mean);
                triang.adjoint().solveInPlace(_alpha);
            }

            Eigen::VectorXd _mu(const Eigen::VectorXd& v, const Eigen::VectorXd& k) const
            {
                return (k.transpose() * _alpha) + _mean_function(v, *this).transpose();
            }

            double _sigma(const Eigen::VectorXd& v, const Eigen::VectorXd& k) const
            {
                Eigen::VectorXd z = _matrixL.triangularView<Eigen::Lower>().solve(k);
                double res = _kernel_function(v, v) - z.dot(z);

                return (res <= std::numeric_limits<double>::epsilon()) ? 0 : res;
            }

            Eigen::VectorXd _compute_k(const Eigen::VectorXd& v) const
            {
                Eigen::VectorXd k(_samples.size());
                for (int i = 0; i < k.size(); i++)
                    k[i] = _kernel_function(_samples[i], v);
                return k;
            }
        };
        /// GPBasic is a GP with a "mean data" mean function, Exponential kernel,
        /// and NO hyper-parameter optimization
        template <typename Params>
        using GPBasic = GP<Params, kernel::MaternFiveHalves<Params>, mean::Data<Params>, gp::NoLFOpt<Params>>;

        /// GPOpt is a GP with a "mean data" mean function, Exponential kernel with Automatic Relevance
        /// Determination (ARD), and hyper-parameter optimization based on Rprop
        template <typename Params>
        using GPOpt = GP<Params, kernel::SquaredExpARD<Params>, mean::Data<Params>, gp::KernelLFOpt<Params>>;
    } // namespace model
} // namespace limbo

#endif


#ifndef DISTRIBUTED_TRAINING_OF_RECURRENT_NEURAL_NETWORKS_BY_FGM_PROTOCOL_GM_HH
#define DISTRIBUTED_TRAINING_OF_RECURRENT_NEURAL_NETWORKS_BY_FGM_PROTOCOL_GM_HH

#include <boost/range/adaptors.hpp>
#include <boost/shared_ptr.hpp>
#include <random>
#include "protocols.hh"
#include "cpp/models/rnn.hh"

namespace algorithms {

    namespace gm {

        using namespace rnn;
        using namespace dds;
        using namespace protocols;

        struct Coordinator;
        struct CoordinatorProxy;
        struct LearningNode;
        struct LearningNodeProxy;

        // This is the GM Network implementation for the classic Geometric Method protocol.
        struct GmNet : LearningNetwork<GmNet, Coordinator, LearningNode> {

            typedef LearningNetwork<network_t, coordinator_t, node_t> gmLearningNetwork;

            GmNet(const set<source_id> &_hids, const string &_name, Query *Q);
        };

        // This is the hub/coordinator implementation for the classic Geometric Method protocol.
        struct Coordinator : process {
            typedef Coordinator coordinator_t;
            typedef LearningNode node_t;
            typedef LearningNodeProxy node_proxy_t;
            typedef GmNet network_t;

            proxy_map<node_proxy_t, node_t> proxy;

            // Protocol Stuff
            RnnLearner *globalLearner;          // ML model
            arma::cube trainX, trainY;          // Trainset data points and labels for warmup
            arma::cube testX, testY;            // Testset data points and labels
            size_t trainPoints;
            Query *Q;                           // query
            QueryState *queryState;             // current query state
            SafeFunction *safeFunction;         // the safe zone wrapper
            size_t k;                           // number of sites
            map<node_t *, size_t> nodeIndex;    // index the nodes
            vector<node_t *> nodePtr;
            set<node_t *> B;                    // initialized by local_violation(), updated by rebalancing algorithm
            set<node_t *> Bcompl;               // Complement of B, updated by rebalancing algo
            arma::mat Mean;                     // Used to compute the mean model
            size_t numViolations;               // Number of violations in the same round (for rebalancing)
            int cnt;                            // Helping counter.

            // Statistics
            size_t nRounds;                     // Total number of rounds
            size_t nRebalances;                 // Total number of rebalances
            size_t nSzSent;                     // Total safe zones sent
            size_t nUpdates;                    // Number of stream updates received

            // Constructor and Destructor
            Coordinator(network_t *nw, Query *Q);

            ~Coordinator() override;

            network_t *Net();

            const ProtocolConfig &Cfg() const;

            // Initialize the learner model and its' variables.
            void InitializeGlobalLearner();

            // Pre-train global model with a small batch to override the initial unstable situation of the distributed algorithm.
            void WarmupGlobalLearner();

            // Method used by the hub to establish the connections with the nodes of the star network.
            void SetupConnections();

            // Start a new round.
            void StartRound();

            // Rebalance algorithm by Kamp.
            void Rebalance(node_t *n);

            // Get a model of a node
            void FetchUpdates(node_t *n);

            // Remote call on host violation
            oneway LocalViolation(sender<node_t> ctx);

            // Printing overall stats of the training process.
            void ShowProgress();

            // Stuff for finalizing a round and start a new one.
            void FinishRound();

            // Getting the accuracy of the global learner.
            double Accuracy();

            // Printing overall stats of the training process.
            void ShowOverallStats();
        };

        struct CoordinatorProxy : remote_proxy<Coordinator> {
            using coordinator_t = Coordinator;

            REMOTE_METHOD(coordinator_t, LocalViolation);

            explicit CoordinatorProxy(process *c) : remote_proxy<coordinator_t>(c) {}
        };

        // This is the site/learning node implementation for the classic Geometric Method protocol.
        struct LearningNode : local_site {

            typedef Coordinator coordinator_t;
            typedef LearningNode node_t;
            typedef LearningNodeProxy node_proxy_t;
            typedef GmNet network_t;
            typedef CoordinatorProxy coord_proxy_t;
            typedef Query query_t;

            query_t *Q;                         // The query management object
            Safezone szone;                     // The safezone object
            RnnLearner *learner{};              // The learning algorithm
            arma::mat currentEstimate;
            arma::mat drift;                    // The drift of the node
            coord_proxy_t coord;                // The proxy of the coordinator/hub
            size_t datapointsPassed;

            // Constructor
            LearningNode(network_t *net, source_id hid, query_t *Q);

            const ProtocolConfig &Cfg() const;

            // Initialize the learner model and its' variables.
            void InitializeLearner();

            // Train the local model with a batch and update protocol variables to check for violations.
            void UpdateState(arma::cube &x, arma::cube &y);

            // Call at the start of each round.
            oneway Reset(const Safezone &newsz);

            // Transfer (drift) model to the coordinator.
            ModelState SendDrift();

            // Get the parameters from the hub at the start of a new round. (for rebalancing)
            void ReceiveRebGlobalParameters(const ModelState &mdl);

            // Get the parameters from the hub at the start of a new round.
            void ReceiveGlobalParameters(const ModelState &params);
        };

        struct LearningNodeProxy : remote_proxy<LearningNode> {
            typedef LearningNode node_t;

            REMOTE_METHOD(node_t, Reset);
            REMOTE_METHOD(node_t, SendDrift);
            REMOTE_METHOD(node_t, ReceiveRebGlobalParameters);
            REMOTE_METHOD(node_t, ReceiveGlobalParameters);

            explicit LearningNodeProxy(process *p) : remote_proxy<node_t>(p) {}
        };

    } // end namespace gm
} // end namespace algorithms

namespace dds {
    template<>
    inline size_t byte_size<algorithms::gm::LearningNode *>(
            algorithms::gm::LearningNode *const &) { return 4; }
}

#endif


//
// Copyright © 2017 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//

#include <boost/test/unit_test.hpp>
#include "ParserFlatbuffersFixture.hpp"
#include "../TfLiteParser.hpp"

#include <string>
#include <iostream>

BOOST_AUTO_TEST_SUITE(TensorflowLiteParser)

struct SoftmaxFixture : public ParserFlatbuffersFixture
{
    explicit SoftmaxFixture()
    {
        m_JsonString = R"(
            {
                "version": 3,
                "operator_codes": [ { "builtin_code": "SOFTMAX" } ],
                "subgraphs": [ {
                    "tensors": [
                        {
                            "shape": [ 1, 7 ],
                            "type": "UINT8",
                            "buffer": 0,
                            "name": "inputTensor",
                            "quantization": {
                                "min": [ 0.0 ],
                                "max": [ 255.0 ],
                                "scale": [ 1.0 ],
                                "zero_point": [ 0 ],
                            }
                        },
                        {
                            "shape": [ 1, 7 ],
                            "type": "UINT8",
                            "buffer": 1,
                            "name": "outputTensor",
                            "quantization": {
                                "min": [ 0.0 ],
                                "max": [ 255.0 ],
                                "scale": [ 0.00390625 ],
                                "zero_point": [ 0 ],
                            }
                        }
                    ],
                    "inputs": [ 0 ],
                    "outputs": [ 1 ],
                    "operators": [
                        {
                          "opcode_index": 0,
                          "inputs": [ 0 ],
                          "outputs": [ 1 ],
                          "builtin_options_type": "SoftmaxOptions",
                          "builtin_options": {
                            "beta": 1.0
                          },
                          "custom_options_format": "FLEXBUFFERS"
                        }
                    ],
                } ],
                "buffers" : [ {}, {} ]
            }
        )";
        SetupSingleInputSingleOutput("inputTensor", "outputTensor");
    }
};

BOOST_FIXTURE_TEST_CASE(ParseSoftmaxLite, SoftmaxFixture)
{
    RunTest<2, armnn::DataType::QAsymmU8>(0, { 0, 0, 100, 0, 0, 0, 0 }, { 0, 0, 255, 0, 0, 0, 0 });
}

BOOST_AUTO_TEST_SUITE_END()


#include "pybind11/pybind11.h"
#include <Eigen/Dense>

#include "drake/bindings/pydrake/common/eigen_geometry_pybind.h"
#include "drake/bindings/pydrake/pydrake_pybind.h"
#include "drake/common/drake_throw.h"
#include "drake/common/eigen_types.h"

using Eigen::Isometry3d;
using Eigen::Quaterniond;
using Eigen::Translation3d;

namespace drake {
namespace pydrake {
namespace {

const double kTolerance = 1e-8;

}  // namespace

PYBIND11_MODULE(eigen_geometry_test_util, m) {
  m.doc() = "Bindings for Eigen geometric types.";

  using T = double;

  // Given an Eigen `x`, return max(abs(x[:])).
  auto max_abs = [](const auto& x) { return x.array().abs().maxCoeff(); };

  m.def("create_isometry", []() { return Isometry3<T>::Identity(); });
  m.def("check_isometry", [max_abs](const Isometry3d& X) {
    const T error = max_abs(X.matrix() - Isometry3<T>::Identity().matrix());
    DRAKE_THROW_UNLESS(error < kTolerance);
  });

  m.def("create_translation",
      []() { return Translation3<T>(Vector3<T>::Zero()); });
  m.def("check_translation", [max_abs](const Translation3<T>& p) {
    const T error = max_abs(p.vector());
    DRAKE_THROW_UNLESS(error < kTolerance);
  });

  m.def("create_quaternion", []() { return Quaternion<T>::Identity(); });
  m.def("check_quaternion", [max_abs](const Quaternion<T>& q) {
    const T error = max_abs(q.coeffs() - Quaternion<T>::Identity().coeffs());
    DRAKE_THROW_UNLESS(error < kTolerance);
  });
}

}  // namespace pydrake
}  // namespace drake