#include "config.h"
#include <algorithm>
#include <cassert>
#include <list>
#include <set>
#include <iostream>
#include <fstream>
#include <memory>
#include <cmath>
#include <array>
#include <thread>
#include <boost/utility.hpp>
#include <boost/format.hpp>

#ifdef USE_CAFFE
#include <caffe/proto/caffe.pb.h>
#include <caffe/util/db.hpp>
#include <caffe/util/io.hpp>
#include <caffe/blob.hpp>

using namespace caffe;
#endif
#include "Im2Col.h"
#ifdef __APPLE__
#include <Accelerate/Accelerate.h>
#endif
#ifdef USE_MKL
#include <mkl.h>
#endif
#ifdef USE_OPENBLAS
#include <cblas.h>
#ifdef USE_ONEDNN
#include <dnnl.h>
#endif
#endif
#ifdef USE_OPENCL
#include "OpenCL.h"
#include "UCTNode.h"
#endif

#include "SGFTree.h"
#include "SGFParser.h"
#include "Utils.h"
#include "FastBoard.h"
#include "Random.h"
#include "Network.h"
#include "GTP.h"
#include "Utils.h"

using namespace Utils;

Network* Network::s_Net = nullptr;
#ifdef USE_CAFFE
std::unique_ptr<caffe::Net> Network::s_net;
#endif

#ifdef USE_OPENCL
extern const std::array<float, 102400> conv1_w;
extern const std::array<float, 128> conv1_b;
extern const std::array<float, 221184> conv2_w;
extern const std::array<float, 192> conv2_b;
extern const std::array<float, 331776> conv3_w;
extern const std::array<float, 192> conv3_b;
extern const std::array<float, 331776> conv4_w;
extern const std::array<float, 192> conv4_b;
extern const std::array<float, 331776> conv5_w;
extern const std::array<float, 192> conv5_b;
extern const std::array<float, 331776> conv6_w;
extern const std::array<float, 192> conv6_b;
extern const std::array<float, 331776> conv7_w;
extern const std::array<float, 192> conv7_b;
extern const std::array<float, 331776> conv8_w;
extern const std::array<float, 192> conv8_b;
extern const std::array<float, 331776> conv9_w;
extern const std::array<float, 192> conv9_b;
extern const std::array<float, 331776> conv10_w;
extern const std::array<float, 192> conv10_b;
extern const std::array<float, 331776> conv11_w;
extern const std::array<float, 192> conv11_b;
extern const std::array<float, 331776> conv12_w;
extern const std::array<float, 192> conv12_b;
extern const std::array<float, 1728> conv13_w;
extern const std::array<float, 1> conv13_b;
#else
extern const std::array<float, 76800> conv1_w;
extern const std::array<float, 96> conv1_b;
extern const std::array<float, 110592> conv2_w;
extern const std::array<float, 128> conv2_b;
extern const std::array<float, 147456> conv3_w;
extern const std::array<float, 128> conv3_b;
extern const std::array<float, 147456> conv4_w;
extern const std::array<float, 128> conv4_b;
extern const std::array<float, 147456> conv5_w;
extern const std::array<float, 128> conv5_b;
extern const std::array<float, 147456> conv6_w;
extern const std::array<float, 128> conv6_b;
extern const std::array<float, 147456> conv7_w;
extern const std::array<float, 128> conv7_b;
extern const std::array<float, 147456> conv8_w;
extern const std::array<float, 128> conv8_b;
extern const std::array<float, 147456> conv9_w;
extern const std::array<float, 128> conv9_b;
extern const std::array<float, 147456> conv10_w;
extern const std::array<float, 128> conv10_b;
extern const std::array<float, 147456> conv11_w;
extern const std::array<float, 128> conv11_b;
extern const std::array<float, 147456> conv12_w;
extern const std::array<float, 128> conv12_b;
extern const std::array<float, 1152> conv13_w;
extern const std::array<float, 1> conv13_b;
#endif

extern const std::array<float, 51200> val_conv1_w;
extern const std::array<float, 64> val_conv1_b;
extern const std::array<float, 36864> val_conv2_w;
extern const std::array<float, 64> val_conv2_b;
extern const std::array<float, 36864> val_conv3_w;
extern const std::array<float, 64> val_conv3_b;
extern const std::array<float, 36864> val_conv4_w;
extern const std::array<float, 64> val_conv4_b;
extern const std::array<float, 36864> val_conv5_w;
extern const std::array<float, 64> val_conv5_b;
extern const std::array<float, 36864> val_conv6_w;
extern const std::array<float, 64> val_conv6_b;
extern const std::array<float, 36864> val_conv7_w;
extern const std::array<float, 64> val_conv7_b;
extern const std::array<float, 36864> val_conv8_w;
extern const std::array<float, 64> val_conv8_b;
extern const std::array<float, 36864> val_conv9_w;
extern const std::array<float, 64> val_conv9_b;
extern const std::array<float, 36864> val_conv10_w;
extern const std::array<float, 64> val_conv10_b;
extern const std::array<float, 36864> val_conv11_w;
extern const std::array<float, 64> val_conv11_b;
extern const std::array<float, 576> val_conv12_w;
extern const std::array<float, 1> val_conv12_b;
extern const std::array<float, 92416> val_ip13_w;
extern const std::array<float, 256> val_ip13_b;
extern const std::array<float, 256> val_ip14_w;
extern const std::array<float, 1> val_ip14_b;

Network * Network::get_Network(void) {
    if (!s_Net) {
        s_Net = new Network();
        s_Net->initialize();
    }
    return s_Net;
}

void Network::benchmark(FastState * state) {
    {
        // Policy
        int BENCH_AMOUNT = 2000;
        int cpus = cfg_num_threads;
        int iters_per_thread = (BENCH_AMOUNT + (cpus - 1)) / cpus;

        Time start;

        ThreadGroup tg(thread_pool);
        for (int i = 0; i < cpus; i++) {
            tg.add_task([iters_per_thread, state]() {
                FastState mystate = *state;
                for (int loop = 0; loop < iters_per_thread; loop++) {
                    auto vec = get_scored_moves(&mystate, Ensemble::RANDOM_ROTATION);
                }
            });
        };
        tg.wait_all();

        Time end;

        myprintf("%5d predictions in %5.2f seconds -> %d p/s\n",
                 BENCH_AMOUNT,
                 (float)Time::timediff(start,end)/100.0,
                 (int)((float)BENCH_AMOUNT/((float)Time::timediff(start,end)/100.0)));
    }
    {
        // Value
        int BENCH_AMOUNT = 10000;
        int cpus = cfg_num_threads;
        int iters_per_thread = (BENCH_AMOUNT + (cpus - 1)) / cpus;

        Time start;
        ThreadGroup tg(thread_pool);
        for (int i = 0; i < cpus; i++) {
            tg.add_task([iters_per_thread, state]() {
                FastState mystate = *state;
                for (int loop = 0; loop < iters_per_thread; loop++) {
                    auto vec = get_value(&mystate, Ensemble::RANDOM_ROTATION);
                }
            });
        };
        tg.wait_all();

        Time end;

        myprintf("%5d evaluations in %5.2f seconds -> %d p/s\n",
                 BENCH_AMOUNT,
                 (float)Time::timediff(start,end)/100.0,
                 (int)((float)BENCH_AMOUNT/((float)Time::timediff(start,end)/100.0)));
    }
}

void Network::initialize(void) {
#ifdef USE_OPENCL
    myprintf("Initializing OpenCL\n");
    opencl.initialize();
    myprintf("Transferring weights to GPU...");
    opencl_policy_net.push_convolve(5, conv1_w, conv1_b);
    opencl_policy_net.push_convolve(3, conv2_w, conv2_b);
    opencl_policy_net.push_convolve(3, conv3_w, conv3_b);
    opencl_policy_net.push_convolve(3, conv4_w, conv4_b);
    opencl_policy_net.push_convolve(3, conv5_w, conv5_b);
    opencl_policy_net.push_convolve(3, conv6_w, conv6_b);
    opencl_policy_net.push_convolve(3, conv7_w, conv7_b);
    opencl_policy_net.push_convolve(3, conv8_w, conv8_b);
    opencl_policy_net.push_convolve(3, conv9_w, conv9_b);
    opencl_policy_net.push_convolve(3, conv10_w, conv10_b);
    opencl_policy_net.push_convolve(3, conv11_w, conv11_b);
    opencl_policy_net.push_convolve(3, conv12_w, conv12_b);
    opencl_policy_net.push_convolve(3, conv13_w, conv13_b);

    opencl_value_net.push_convolve(5, val_conv1_w, val_conv1_b);
    opencl_value_net.push_convolve(3, val_conv2_w, val_conv2_b);
    opencl_value_net.push_convolve(3, val_conv3_w, val_conv3_b);
    opencl_value_net.push_convolve(3, val_conv4_w, val_conv4_b);
    opencl_value_net.push_convolve(3, val_conv5_w, val_conv5_b);
    opencl_value_net.push_convolve(3, val_conv6_w, val_conv6_b);
    opencl_value_net.push_convolve(3, val_conv7_w, val_conv7_b);
    opencl_value_net.push_convolve(3, val_conv8_w, val_conv8_b);
    opencl_value_net.push_convolve(3, val_conv9_w, val_conv9_b);
    opencl_value_net.push_convolve(3, val_conv10_w, val_conv10_b);
    opencl_value_net.push_convolve(3, val_conv11_w, val_conv11_b);
    opencl_value_net.push_convolve(3, val_conv12_w, val_conv12_b);
    opencl_value_net.push_innerproduct(val_ip13_w, val_ip13_b);
    opencl_value_net.push_innerproduct(val_ip14_w, val_ip14_b);
    myprintf("done\n");
#endif
#ifdef USE_BLAS
#ifndef __APPLE__
#ifdef USE_OPENBLAS
    openblas_set_num_threads(1);
    myprintf("BLAS Core: %s\n", openblas_get_corename());
#endif
#ifdef USE_MKL
    //mkl_set_threading_layer(MKL_THREADING_SEQUENTIAL);
    mkl_set_num_threads(1);
    MKLVersion Version;
    mkl_get_version(&Version);
    myprintf("BLAS core: MKL %s\n", Version.Processor);
#endif
#endif
#endif
#ifdef USE_CAFFE
    myprintf("Initializing DCNN...");
    Caffe::set_mode(Caffe::GPU);

    s_net.reset(new Net("model_5722.txt", TEST));
    s_net->CopyTrainedLayersFrom("model_5779.caffemodel");

    myprintf("Inputs: %d Outputs: %d\n",
        s_net->num_inputs(), s_net->num_outputs());

    Blob* input_layer = s_net->input_blobs()[0];
    int num_channels = input_layer->channels();
    int width = input_layer->width();
    int height = input_layer->height();
    myprintf("Input: channels=%d, width=%d, height=%d\n", num_channels, width, height);

    for (int i = 0; i < s_net->num_outputs(); i++) {
        Blob* output_layer = s_net->output_blobs()[i];
        int num_out_channels = output_layer->channels();
        width = output_layer->width();
        height = output_layer->height();
        myprintf("Output: channels=%d, width=%d, height=%d\n", num_out_channels, width, height);
    }

//#define WRITE_WEIGHTS
#ifdef WRITE_WEIGHTS
    std::ofstream out("weights.txt");
    out << "#include <array>" << std::endl << std::endl;
#endif

    int total_weights = 0;
    auto & layers = s_net->layers();
    myprintf("%d layers:\n", layers.size());
    int layer_num = 1;
#ifdef WRITE_WEIGHTS
    int conv_count = 0;
#endif
    for (auto it = layers.begin(); it != layers.end(); ++it, ++layer_num) {
        myprintf("layer %d (%s)", layer_num, (*it)->type());
        auto & blobs = (*it)->blobs();
        if (blobs.size() > 0) myprintf(" = ");
        for (auto pars = blobs.begin(); pars != blobs.end(); ++pars) {
            const Blob & blob = *(*pars);
            total_weights += blob.count();
            myprintf("%s ", blob.shape_string().c_str());
            if (boost::next(pars) != blobs.end()) myprintf("+ ");

#ifdef WRITE_WEIGHTS
            out << "// " << blob.shape_string() << std::endl;
            if (strcmp((*it)->type(), "Convolution") == 0) {
                if (pars == blobs.begin()) {
                    conv_count++;
                    out << "std::array<float, " << blob.count()
                        << "> val_conv" << conv_count << "_w = {{" << std::endl;
                } else {
                    out << "std::array<float, " << blob.count()
                        << "> val_conv" << conv_count << "_b = {{" << std::endl;
                }
            } else if (strcmp((*it)->type(), "BatchNorm") == 0) {
                out << "std::array<float, " << blob.count()
                    << "> val_bn" << conv_count << "_w" << (pars - blobs.begin()) + 1
                    << " = {{" << std::endl;
            } else if (strcmp((*it)->type(), "InnerProduct") == 0) {
                if (pars == blobs.begin()) {
                    conv_count++;
                    out << "std::array<float, " << blob.count()
                        << "> val_ip" << conv_count << "_w = {{" << std::endl;
                } else {
                    out << "std::array<float, " << blob.count()
                        << "> val_ip" << conv_count << "_b = {{" << std::endl;
                }
            } else {
                out << "std::array<float, " << blob.count()
                    << "> val_sc" << conv_count << "_w" << (pars - blobs.begin()) + 1
                    << " = {{" << std::endl;
            }
            for (int idx = 0; idx < blob.count(); idx++) {
                out << blob.cpu_data<float>()[idx];
                if (idx != blob.count() - 1) out << ", ";
                else out << " }};" << std::endl;
            }
            out << std::endl;
#endif
        }
        myprintf("\n");
    }
#ifdef WRITE_WEIGHTS
    out.close();
#endif
    myprintf("%d total DCNN weights\n", total_weights);
#endif
}

#ifdef USE_BLAS
template<unsigned int filter_size,
         unsigned int channels, unsigned int outputs,
         size_t W, size_t B>
void convolve(std::vector<float>& input,
              const std::array<float, W>& weights,
              const std::array<float, B>& biases,
              std::vector<float>& output) {
    // fixed for 19x19
    constexpr unsigned int width = 19;
    constexpr unsigned int height = 19;
    constexpr unsigned int spatial_out = width * height;

    constexpr unsigned int filter_len = filter_size * filter_size;
    constexpr unsigned int filter_dim = filter_len * channels;

    std::vector<float> col(filter_dim * width * height);
    im2col<channels, filter_size>(input, col);

    // Weight shape (output, input, filter_size, filter_size)
    // 96 22 5 5
    // outputs[96,19x19] = weights[96,22x9] x col[22x9,19x19]
    // C←αAB + βC
    // M Number of rows in matrices A and C.
    // N Number of columns in matrices B and C.
    // K Number of columns in matrix A; number of rows in matrix B.
    // lda The size of the first dimention of matrix A; if you are
    // passing a matrix A[m][n], the value should be m.
    //    cblas_sgemm(CblasRowMajor, TransA, TransB, M, N, K, alpha, A, lda, B,
    //                ldb, beta, C, N);

#ifdef USE_ONEDNN
    dnnl_sgemm('N', 'N',
                outputs, spatial_out, filter_dim,
                1.0f, &weights[0], filter_dim,
                &col[0], spatial_out,
                0.0f, &output[0], spatial_out);
#else
    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                // M        N            K
                outputs, spatial_out, filter_dim,
                1.0f, &weights[0], filter_dim,
                &col[0], spatial_out,
                0.0f, &output[0], spatial_out);
#endif

    auto lambda_ELU = [](float val) { return (val > 0.0f) ?
                                      val : 1.0f * (std::exp(val) - 1.0f); };
    //auto lambda_ReLU = [](float val) { return (val > 0.0f) ?
    //                                   val : 0.0f; };

    for (unsigned int o = 0; o < outputs; o++) {
        for (unsigned int b = 0; b < spatial_out; b++) {
            output[(o * spatial_out) + b] =
                lambda_ELU(biases[o] + output[(o * spatial_out) + b]);
        }
    }
}

template<unsigned int inputs,
         unsigned int outputs,
         size_t W, size_t B>
void innerproduct(std::vector<float>& input,
                  const std::array<float, W>& weights,
                  const std::array<float, B>& biases,
                  std::vector<float>& output) {
    assert(B == outputs);

    cblas_sgemv(CblasRowMajor, CblasNoTrans,
                // M     K
                outputs, inputs,
                1.0f, &weights[0], inputs,
                &input[0], 1,
                0.0f, &output[0], 1);

    auto lambda_ELU = [](float val) { return (val > 0.0f) ?
                                      val : 1.0f * (std::exp(val) - 1.0f); };
    //auto lambda_ReLU = [](float val) { return (val > 0.0f) ?
    //                                   val : 0.0f; };

    for (unsigned int o = 0; o < outputs; o++) {
        float val = biases[o] + output[o];
        if (outputs > 1) {
            val = lambda_ELU(val);
        }
        output[o] = val;
    }
}

template<unsigned int channels,
         unsigned int spatial_size>
void batchnorm(std::vector<float>& input,
               std::array<float, channels>& means,
               std::array<float, channels>& variances,
               std::array<float, 1> scale,
               std::vector<float>& output)
{
    constexpr float epsilon = 1e-5;

    for (unsigned int c = 0; c < channels; ++c) {
        float mean = means[c] / scale[0];
        float variance = variances[c] / scale[0];
        variance += epsilon;
        float scale_stddiv = 1.0f / std::sqrt(variance);

        float * out = &output[c * spatial_size];
        float const * in  = &input[c * spatial_size];
        for (unsigned int b = 0; b < spatial_size; b++) {
            out[b] = scale_stddiv * (in[b] - mean);
        }
    }
}
#endif

void Network::softmax(std::vector<float>& input,
                      std::vector<float>& output,
                      float temperature) {
    assert(&input != &output);

    float alpha = *std::max_element(input.begin(),
                                    input.begin() + output.size());
    alpha /= temperature;

    float denom = 0.0f;
    std::vector<float> helper(output.size());
    for (size_t i = 0; i < output.size(); i++) {
        float val  = std::exp((input[i]/temperature) - alpha);
        helper[i]  = val;
        denom     += val;
    }
    for (size_t i = 0; i < output.size(); i++) {
        output[i] = helper[i] / denom;
    }
}

#ifdef USE_OPENCL
class CallbackData {
public:
    std::atomic<int> * m_nodecount;
    FastState m_state;
    UCTNode * m_node;
    int m_rotation;
    std::atomic<int> * m_thread_results_outstanding;
    std::vector<float> m_output_data;
    std::vector<float> m_input_data;
    Network::BoardPlane m_ladder;
};

extern "C" void CL_CALLBACK forward_cb(cl_event event, cl_int status,
                                       void* data) {
    CallbackData * cb_data = static_cast<CallbackData*>(data);

    // Mark the kernels as available
    cb_data->m_thread_results_outstanding->fetch_sub(1, std::memory_order_release);

    constexpr int width = 19;
    constexpr int height = 19;
    std::vector<float> softmax_data(width * height);
    Network::softmax(cb_data->m_output_data, softmax_data, cfg_softmax_temp);
    std::vector<float>& outputs = softmax_data;

    Network::Netresult result;

    for (size_t idx = 0; idx < outputs.size(); idx++) {
        int rot_idx = Network::rev_rotate_nn_idx(idx, cb_data->m_rotation);
        float val = outputs[rot_idx];
        int x = idx % 19;
        int y = idx / 19;
        int vtx = cb_data->m_state.board.get_vertex(x, y);
        if (cb_data->m_state.board.get_square(vtx) == FastBoard::EMPTY) {
            result.push_back(std::make_pair(val, vtx));
        }
    }

    /* prune losing ladders completely */
    for (auto & sm : result) {
        std::pair<int, int> xy = cb_data->m_state.board.get_xy(sm.second);
        int bitmappos = (xy.second * 19) + xy.first;
        if (cb_data->m_ladder[bitmappos]) {
            //myprintf("Ladder at %s (%d) score %f\n",
            //         state->board.move_to_text(sm.second).c_str(),
            //         sm.second,
            //         sm.first);
            sm.first = 0.0f;
        }
    }

    // Network::show_heatmap(&cb_data->m_state, result, false);

    cb_data->m_node->scoring_cb(cb_data->m_nodecount, cb_data->m_state,
                                result, false);

    delete cb_data;

    // Reduce the count of things having pointers to UCTNodes
    // or UCTSearch. We cannot destroy the search till these
    // have finished.
    opencl.callback_finished();
}

void Network::async_scored_moves(std::atomic<int> * nodecount,
                                 FastState * state,
                                 UCTNode * node,
                                 Ensemble ensemble,
                                 int rotation) {
    if (state->board.get_boardsize() != 19) {
        return;
    }

    assert(ensemble == DIRECT || ensemble == RANDOM_ROTATION);
    if (ensemble == RANDOM_ROTATION) {
        assert(rotation == -1);
        rotation = Random::get_Rng()->randfix<8>();
    } else {
        assert(ensemble == DIRECT);
    }

    CallbackData * cb_data = new CallbackData();

    NNPlanes planes;
    BoardPlane *ladder;
    gather_features_policy(state, planes, &ladder);

    constexpr int width = 19;
    constexpr int height = 19;

    cb_data->m_nodecount = nodecount;
    cb_data->m_state = *state;
    cb_data->m_node = node;
    cb_data->m_input_data.resize(Network::MAX_CHANNELS * 19 * 19);
    cb_data->m_output_data.resize(Network::MAX_CHANNELS * 19 * 19);
    cb_data->m_thread_results_outstanding = opencl.get_thread_results_outstanding();
    //assert(cb_data->m_thread_result_outstanding.load(boost::memory_order_acquire) == 0);
    cb_data->m_rotation = rotation;
    cb_data->m_ladder = *ladder;

    for (int c = 0; c < Network::POLICY_CHANNELS; ++c) {
        for (int h = 0; h < height; ++h) {
            for (int w = 0; w < width; ++w) {
                int vtx = rotate_nn_idx(h * 19 + w, rotation);
                cb_data->m_input_data[(c * height + h) * width + w] =
                    (float)planes[c][vtx];
            }
        }
    }

    void * data = static_cast<void*>(cb_data);

    opencl_policy_net.forward(cb_data->m_input_data,
                                    cb_data->m_output_data,
                                    forward_cb, data);
}
#endif

float Network::get_value(FastState * state, Ensemble ensemble) {
    if (state->board.get_boardsize() != 19) {
        assert(false);
        return 0.5f;
    }

    NNPlanes planes;
    gather_features_value(state, planes);
    float result;

    if (ensemble == DIRECT) {
        result = get_value_internal(state, planes, 0);
    } else if (ensemble == RANDOM_ROTATION) {
        int rotation = Random::get_Rng()->randfix<8>();
        result = get_value_internal(state, planes, rotation);
    } else {
        assert(ensemble == AVERAGE_ALL);
        result = get_value_internal(state, planes, 0);
        //myprintf("%5.4f ", result);
        for (int r = 1; r < 8; r++) {
            float sum_res = get_value_internal(state, planes, r);
            //myprintf("%5.4f ", sum_res);
            result += sum_res;
        }
        result /= 8.0f;
    }

    //if (ensemble == AVERAGE_ALL || ensemble == DIRECT) {
    //    myprintf("==> %5.4f\n", result);
    //}

    return result;
}

Network::Netresult Network::get_scored_moves(
    FastState * state, Ensemble ensemble, int rotation) {
    Netresult result;
    if (state->board.get_boardsize() != 19) {
        return result;
    }

    NNPlanes planes;
    BoardPlane* ladder;
    gather_features_policy(state, planes, &ladder);

    if (ensemble == DIRECT) {
        assert(rotation >= 0 && rotation <= 7);
        result = get_scored_moves_internal(state, planes, rotation);
    } else if (ensemble == RANDOM_ROTATION) {
        assert(rotation == -1);
        int rand_rot = Random::get_Rng()->randfix<8>();
        result = get_scored_moves_internal(state, planes, rand_rot);
    } else {
        assert(ensemble == AVERAGE_ALL);
        result = get_scored_moves_internal(state, planes, 0);
        for (int r = 1; r < 8; r++) {
            auto sum_res = get_scored_moves_internal(state, planes, r);
            for (size_t i = 0; i < sum_res.size(); i++) {
                assert(result[i].second == sum_res[i].second);
                result[i].first += sum_res[i].first;
            }
        }
        std::for_each(result.begin(), result.end(),
                      [](scored_node & sn){ sn.first /= 8.0f; });
    }

    /* prune losing ladders completely */
    for (auto & sm : result) {
        std::pair<int, int> xy = state->board.get_xy(sm.second);
        int bitmappos = (xy.second * 19) + xy.first;
        if ((*ladder)[bitmappos]) {
            //myprintf("Ladder at %s (%d) score %f\n",
            //         state->board.move_to_text(sm.second).c_str(),
            //         sm.second,
            //         sm.first);
            sm.first = 0.0f;
        }
    }

    // if (ensemble == AVERAGE_ALL || ensemble == DIRECT) {
    //     show_heatmap(state, result, true);
    // }

    return result;
}

float Network::get_value_internal(
    FastState * state, NNPlanes & planes, int rotation) {
    assert(rotation >= 0 && rotation <= 7);
    float result;

    constexpr int channels = VALUE_CHANNELS;
    constexpr int width = 19;
    constexpr int height = 19;
    constexpr int max_channels = MAX_VALUE_CHANNELS;
    std::vector<float> orig_input_data(planes.size() * width * height);
    std::vector<float> input_data(max_channels * width * height);
    std::vector<float> output_data(max_channels * width * height);
    std::vector<float> winrate_data(256);
    std::vector<float> winrate_out(1);

    for (int c = 0; c < channels; ++c) {
        for (int h = 0; h < height; ++h) {
            for (int w = 0; w < width; ++w) {
                int vtx = rotate_nn_idx(h * 19 + w, rotation);
                orig_input_data[(c * height + h) * width + w] =
                    (float)planes[c][vtx];
            }
        }
    }
#ifdef USE_OPENCL
    std::copy(orig_input_data.begin(), orig_input_data.end(), input_data.begin());
    opencl_value_net.forward(input_data, output_data, nullptr, nullptr);
    // Sigmoid
    float winrate_sig = (1.0f + std::tanh(output_data[0])) / 2.0f;
    result = winrate_sig;
#elif defined(USE_BLAS)
    std::copy(orig_input_data.begin(), orig_input_data.end(), input_data.begin());

    convolve<5, 32, 64>(input_data, val_conv1_w, val_conv1_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 64, 64>(input_data, val_conv2_w, val_conv2_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 64, 64>(input_data, val_conv3_w, val_conv3_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 64, 64>(input_data, val_conv4_w, val_conv4_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 64, 64>(input_data, val_conv5_w, val_conv5_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 64, 64>(input_data, val_conv6_w, val_conv6_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 64, 64>(input_data, val_conv7_w, val_conv7_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 64, 64>(input_data, val_conv8_w, val_conv8_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 64, 64>(input_data, val_conv9_w, val_conv9_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 64, 64>(input_data, val_conv10_w, val_conv10_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 64, 64>(input_data, val_conv11_w, val_conv11_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 64,  1>(input_data, val_conv12_w, val_conv12_b, output_data);
    // Now get the score
    innerproduct<361, 256>(output_data, val_ip13_w, val_ip13_b, winrate_data);
    innerproduct<256, 1>(winrate_data, val_ip14_w, val_ip14_b, winrate_out);
    // Sigmoid
    float winrate_sig = (1.0f + std::tanh(winrate_out[0])) / 2.0f;
    result = winrate_sig;
 #endif
    return result;
}

Network::Netresult Network::get_scored_moves_internal(
    FastState * state, NNPlanes & planes, int rotation) {
    Netresult result;
    assert(rotation >= 0 && rotation <= 7);
#ifdef USE_CAFFE
    Blob* input_layer = s_net->input_blobs()[0];
    int channels = input_layer->channels();
    int width = input_layer->width();
    int height = input_layer->height();
    assert(channels == (int)planes.size());
    assert(width == state->board.get_boardsize());
    assert(height == state->board.get_boardsize());
    float* orig_input_data = input_layer->mutable_cpu_data<float>();
#else
    constexpr int channels = POLICY_CHANNELS;
    constexpr int width = 19;
    constexpr int height = 19;
    constexpr int max_channels = MAX_CHANNELS;
    std::vector<float> orig_input_data(planes.size() * width * height);
    std::vector<float> input_data(max_channels * width * height);
    std::vector<float> output_data(max_channels * width * height);
    std::vector<float> softmax_data(width * height);
#endif
    for (int c = 0; c < channels; ++c) {
        for (int h = 0; h < height; ++h) {
            for (int w = 0; w < width; ++w) {
                int vtx = rotate_nn_idx(h * 19 + w, rotation);
                orig_input_data[(c * height + h) * width + w] =
                    (float)planes[c][vtx];
            }
        }
    }
#ifdef USE_OPENCL
    std::copy(orig_input_data.begin(), orig_input_data.end(), input_data.begin());
    opencl_policy_net.forward(input_data, output_data, nullptr, nullptr);
    softmax(output_data, softmax_data, cfg_softmax_temp);
    std::vector<float>& outputs = softmax_data;
#elif defined(USE_BLAS)
    // XXX really only need the first 24
    std::copy(orig_input_data.begin(), orig_input_data.end(), input_data.begin());

    convolve<5,  32,  96>(input_data, conv1_w, conv1_b, output_data);
    std::swap(input_data, output_data);
    convolve<3,  96, 128>(input_data, conv2_w, conv2_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 128, 128>(input_data, conv3_w, conv3_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 128, 128>(input_data, conv4_w, conv4_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 128, 128>(input_data, conv5_w, conv5_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 128, 128>(input_data, conv6_w, conv6_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 128, 128>(input_data, conv7_w, conv7_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 128, 128>(input_data, conv8_w, conv8_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 128, 128>(input_data, conv9_w, conv9_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 128, 128>(input_data, conv10_w, conv10_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 128, 128>(input_data, conv11_w, conv11_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 128, 128>(input_data, conv12_w, conv12_b, output_data);
    std::swap(input_data, output_data);
    convolve<3, 128,   1>(input_data, conv13_w, conv13_b, output_data);
    softmax(output_data, softmax_data, cfg_softmax_temp);

    // Move scores
    std::vector<float>& outputs = softmax_data;
#endif
#ifdef USE_CAFFE
    s_net->Forward();
    Blob* output_layer = s_net->output_blobs()[0];
    const float* begin = output_layer->cpu_data<float>();
    const float* end = begin + output_layer->channels();
    auto outputs = std::vector<float>(begin, end);
#endif
    for (size_t idx = 0; idx < outputs.size(); idx++) {
        int rot_idx = rev_rotate_nn_idx(idx, rotation);
        float val = outputs[rot_idx];
        int x = idx % 19;
        int y = idx / 19;
        int vtx = state->board.get_vertex(x, y);
        if (state->board.get_square(vtx) == FastBoard::EMPTY) {
            result.push_back(std::make_pair(val, vtx));
        }
    }

    return result;
}

void Network::show_heatmap(FastState * state, Netresult& result, bool topmoves) {
    auto moves = result;
    std::vector<std::string> display_map;
    std::string line;

    for (unsigned int y = 0; y < 19; y++) {
        for (unsigned int x = 0; x < 19; x++) {
            int vtx = state->board.get_vertex(x, y);

            auto item = std::find_if(moves.cbegin(), moves.cend(),
                [&vtx](scored_node const & item) {
                return item.second == vtx;
            });

            float score = 0.0f;
            // Non-empty squares won't be scored
            if (item != moves.end()) {
                score = item->first;
                assert(vtx == item->second);
            }

            line += boost::str(boost::format("%3d ") % int(score * 1000));
            if (x == 18) {
                display_map.push_back(line);
                line.clear();
            }
        }
    }

    for (int i = display_map.size() - 1; i >= 0; --i) {
        myprintf("%s\n", display_map[i].c_str());
    }

    if (topmoves) {
    std::stable_sort(moves.rbegin(), moves.rend());

    float cum = 0.0f;
    size_t tried = 0;
    while (cum < 0.85f && tried < moves.size()) {
        if (moves[tried].first < 0.01f) break;
        myprintf("%1.3f (%s)\n",
                 moves[tried].first,
                 state->board.move_to_text(moves[tried].second).c_str());
        cum += moves[tried].first;
        tried++;
    }
    }
}

void Network::gather_features_policy(FastState * state, NNPlanes & planes,
                                     BoardPlane** ladder_out) {
    planes.resize(32);
    BoardPlane& empt_color   = planes[0];
    BoardPlane& move_color   = planes[1];
    BoardPlane& othr_color   = planes[2];
    BoardPlane& libs_1       = planes[3];
    BoardPlane& libs_2       = planes[4];
    BoardPlane& libs_3       = planes[5];
    BoardPlane& libs_4       = planes[6];
    BoardPlane& libs_5p      = planes[7];
    BoardPlane& libs_1_e     = planes[8];
    BoardPlane& libs_2_e     = planes[9];
    BoardPlane& libs_3_e     = planes[10];
    BoardPlane& libs_4_e     = planes[11];
    BoardPlane& libs_5p_e    = planes[12];
    BoardPlane& after_1      = planes[13];
    BoardPlane& after_2      = planes[14];
    BoardPlane& after_3      = planes[15];
    BoardPlane& after_4      = planes[16];
    BoardPlane& after_5      = planes[17];
    BoardPlane& after_6p     = planes[18];
    BoardPlane& after_1_e    = planes[19];
    BoardPlane& after_2_e    = planes[20];
    BoardPlane& after_3_e    = planes[21];
    BoardPlane& after_4_e    = planes[22];
    BoardPlane& after_5_e    = planes[23];
    BoardPlane& after_6p_e   = planes[24];
    BoardPlane& ladder       = planes[25];
    BoardPlane& ladder_win   = planes[26];
    BoardPlane& komove       = planes[27];
    BoardPlane& movehist1    = planes[28];
    BoardPlane& movehist2    = planes[29];
    BoardPlane& has_komi     = planes[30];
    BoardPlane& line_3       = planes[31];

    if (ladder_out) {
        *ladder_out = &ladder;
    }

    bool white_has_komi = true;
    if (std::fabs(state->get_komi()) <= 0.75f) {
        white_has_komi = false;
    }

    int tomove = state->get_to_move();
    // collect white, black occupation planes
    for (int j = 0; j < 19; j++) {
        for(int i = 0; i < 19; i++) {
            int vtx = state->board.get_vertex(i, j);
            FastBoard::square_t color =
                state->board.get_square(vtx);
            int idx = j * 19 + i;
            if (i == 2 || i == 16 || j == 2 || j == 16) {
                line_3[idx] = true;
            }
            if (color != FastBoard::EMPTY) {
                // White gets extra points in scoring
                if (color == FastBoard::WHITE && white_has_komi) {
                    has_komi[idx] = true;
                }
                if (color == tomove) {
                    move_color[idx] = true;
                } else {
                    othr_color[idx] = true;
                }
                int rlibs = state->board.count_rliberties(vtx);
                if (rlibs == 1) {
                    if (color == tomove) {
                        libs_1[idx] = true;
                    } else {
                        libs_1_e[idx] = true;
                    }
                } else if (rlibs == 2) {
                    if (color == tomove) {
                        libs_2[idx] = true;
                    } else {
                        libs_2_e[idx] = true;
                    }
                } else if (rlibs == 3) {
                    if (color == tomove) {
                        libs_3[idx] = true;
                    } else {
                        libs_3_e[idx] = true;
                    }
                } else if (rlibs == 4) {
                    if (color == tomove) {
                        libs_4[idx] = true;
                    } else {
                        libs_4_e[idx] = true;
                    }
                } else if (rlibs >= 5) {
                    if (color == tomove) {
                        libs_5p[idx] = true;
                    } else {
                        libs_5p_e[idx] = true;
                    }
                }
            } else {
                empt_color[idx] = true;

                std::pair<int, int> p =
                    state->board.after_liberties(tomove, vtx);
                int al = p.first;
                if (al == 1) {
                    after_1[idx] = true;
                } else if (al == 2) {
                    after_2[idx] = true;
                } else if (al == 3) {
                    after_3[idx] = true;
                } else if (al == 4) {
                    after_4[idx] = true;
                } else if (al == 5) {
                    after_5[idx] = true;
                } else if (al >= 6) {
                    after_6p[idx] = true;
                }
                int at = p.second;
                if (at == 1) {
                    after_1_e[idx] = true;
                } else if (at == 2) {
                    after_2_e[idx] = true;
                } else if (at == 3) {
                    after_3_e[idx] = true;
                } else if (at == 4) {
                    after_4_e[idx] = true;
                } else if (at == 5) {
                    after_5_e[idx] = true;
                } else if (at >= 6) {
                    after_6p_e[idx] = true;
                }

                int ae = state->board.count_pliberties(vtx);
                if (ae == 2) {
                    int ss = state->board.saving_size(tomove, vtx);
                    if (ss > 0) {
                        bool ll = state->board.check_losing_ladder(tomove, vtx);
                        ladder[idx] = ll;
                    }
                }
                if (state->board.check_winning_ladder(tomove, vtx)) {
                    ladder_win[idx] = true;
                }
            }
        }
    }

    if (state->get_last_move() > 0) {
        std::pair<int, int> lastmove = state->board.get_xy(state->get_last_move());
        int idx = lastmove.second * 19 + lastmove.first;
        movehist1[idx] = true;
        if (state->get_prevlast_move() > 0) {
            std::pair<int, int> prevlast = state->board.get_xy(state->get_prevlast_move());
            int idxp = prevlast.second * 19 + prevlast.first;
            movehist2[idxp] = true;
        }
    }

    if (state->get_komove() > 0) {
        std::pair<int, int> kosq = state->board.get_xy(state->get_komove());
        int idx = kosq.second * 19 + kosq.first;
        komove[idx] = true;
    }
}

void Network::gather_features_value(FastState * state, NNPlanes & planes) {
    planes.resize(32);
    BoardPlane& empt_color   = planes[0];
    BoardPlane& move_color   = planes[1];
    BoardPlane& othr_color   = planes[2];
    BoardPlane& libs_1       = planes[3];
    BoardPlane& libs_2       = planes[4];
    BoardPlane& libs_3       = planes[5];
    BoardPlane& libs_4       = planes[6];
    BoardPlane& libs_5       = planes[7];
    BoardPlane& libs_6p      = planes[8];
    BoardPlane& libs_1_e     = planes[9];
    BoardPlane& libs_2_e     = planes[10];
    BoardPlane& libs_3_e     = planes[11];
    BoardPlane& libs_4_e     = planes[12];
    BoardPlane& libs_5_e     = planes[13];
    BoardPlane& libs_6p_e    = planes[14];
    BoardPlane& after_1      = planes[15];
    BoardPlane& after_2      = planes[16];
    BoardPlane& after_3      = planes[17];
    BoardPlane& after_4      = planes[18];
    BoardPlane& after_5      = planes[19];
    BoardPlane& after_6p     = planes[20];
    BoardPlane& after_1_e    = planes[21];
    BoardPlane& after_2_e    = planes[22];
    BoardPlane& after_3_e    = planes[23];
    BoardPlane& after_4_e    = planes[24];
    BoardPlane& after_5_e    = planes[25];
    BoardPlane& after_6p_e   = planes[26];
    BoardPlane& ladder       = planes[27];
    BoardPlane& ladder_win   = planes[28];
    BoardPlane& has_komi     = planes[29];
    BoardPlane& line_3       = planes[30];
    BoardPlane& komove       = planes[31];

    bool white_has_komi = true;
    if (std::fabs(state->get_komi()) <= 0.75f) {
        white_has_komi = false;
    }

    int tomove = state->get_to_move();
    // collect white, black occupation planes
    for (int j = 0; j < 19; j++) {
        for(int i = 0; i < 19; i++) {
            int vtx = state->board.get_vertex(i, j);
            FastBoard::square_t color =
                state->board.get_square(vtx);
            int idx = j * 19 + i;
            if (i == 2 || i == 16 || j == 2 || j == 16) {
                line_3[idx] = true;
            }
            if (color != FastBoard::EMPTY) {
                // White gets extra points in scoring
                if (color == FastBoard::WHITE && white_has_komi) {
                    has_komi[idx] = true;
                }
                if (color == tomove) {
                    move_color[idx] = true;
                } else {
                    othr_color[idx] = true;
                }
                int rlibs = state->board.count_rliberties(vtx);
                if (rlibs == 1) {
                    if (color == tomove) {
                        libs_1[idx] = true;
                    } else {
                        libs_1_e[idx] = true;
                    }
                } else if (rlibs == 2) {
                    if (color == tomove) {
                       libs_2[idx] = true;
                    } else {
                        libs_2_e[idx] = true;
                    }
                } else if (rlibs == 3) {
                    if (color == tomove) {
                        libs_3[idx] = true;
                    } else {
                        libs_3_e[idx] = true;
                    }
                } else if (rlibs == 4) {
                    if (color == tomove) {
                        libs_4[idx] = true;
                    } else {
                        libs_4_e[idx] = true;
                    }
                } else if (rlibs == 5) {
                    if (color == tomove) {
                        libs_5[idx] = true;
                    } else {
                        libs_5_e[idx] = true;
                    }
                }  else if (rlibs >= 6) {
                    if (color == tomove) {
                        libs_6p[idx] = true;
                    } else {
                        libs_6p_e[idx] = true;
                    }
                }
            } else {
                empt_color[idx] = true;

                std::pair<int, int> p =
                    state->board.after_liberties(tomove, vtx);
                int al = p.first;
                if (al == 1) {
                    after_1[idx] = true;
                } else if (al == 2) {
                    after_2[idx] = true;
                } else if (al == 3) {
                    after_3[idx] = true;
                } else if (al == 4) {
                    after_4[idx] = true;
                } else if (al == 5) {
                    after_5[idx] = true;
                } else if (al >= 6) {
                    after_6p[idx] = true;
                }
                int at = p.second;
                if (at == 1) {
                    after_1_e[idx] = true;
                } else if (at == 2) {
                    after_2_e[idx] = true;
                } else if (at == 3) {
                    after_3_e[idx] = true;
                } else if (at == 4) {
                    after_4_e[idx] = true;
                } else if (at == 5) {
                    after_5_e[idx] = true;
                } else if (at >= 6) {
                    after_6p_e[idx] = true;
                }

                int ss = state->board.saving_size(tomove, vtx);
                int ae = state->board.count_pliberties(vtx);
                if (ae == 2) {
                    if (ss > 0) {
                        bool ll = state->board.check_losing_ladder(tomove, vtx);
                        ladder[idx] = ll;
                    }
                }
                if (state->board.check_winning_ladder(tomove, vtx)) {
                    ladder_win[idx] = true;
                }
            }
        }
    }

    if (state->get_komove() > 0) {
        std::pair<int, int> kosq = state->board.get_xy(state->get_komove());
        int idx = kosq.second * 19 + kosq.first;
        komove[idx] = true;
    }
}

void Network::gather_traindata(std::string filename, TrainVector& data) {
    std::vector<std::string> games = SGFParser::chop_all(filename);
    int gametotal = games.size();
    int gamecount = 0;

    size_t train_pos = 0;
    size_t test_pos = 0;

    myprintf("Total games in file: %d\n", gametotal);
    myprintf("Shuffling...\n");
    std::random_shuffle(games.begin(), games.end());

    while (gamecount < gametotal) {
        std::unique_ptr<SGFTree> sgftree(new SGFTree);

        try {
            sgftree->load_from_string(games[gamecount]);
        } catch (...) {
        };

        SGFTree * treewalk = &(*sgftree);
        size_t counter = 0;

        size_t movecount = sgftree->count_mainline_moves();
        std::vector<int> tree_moves = sgftree->get_mainline();
        int who_won = sgftree->get_winner();
        int handicap = sgftree->get_state()->get_handicap();
        //float komi = sgftree->get_state()->get_komi();
        if (handicap) {
            goto skipnext;
        }
        // 5.5, 6.5, 7.5
        //if (std::abs(komi) > 0.75f && (std::abs(komi) < 5.25f || std::abs(komi) > 7.75f)) {
        //    goto skipnext;
        //}
        if (who_won != FastBoard::BLACK && who_won != FastBoard::WHITE) {
            goto skipnext;
        }

        while (counter < movecount) {
            assert(treewalk != NULL);
            assert(treewalk->get_state() != NULL);
            if (treewalk->get_state()->board.get_boardsize() != 19)
                break;

            KoState * state = treewalk->get_state();
            int tomove = state->get_to_move();
            int move;

            if (treewalk->get_child(0) != NULL) {
                move = treewalk->get_child(0)->get_move(tomove);
                if (move == SGFTree::EOT) {
                    break;
                }
            } else {
                break;
            }

            assert(move == tree_moves[counter]);
            int this_move = -1;

            std::vector<int> moves = state->generate_moves(tomove);
            bool moveseen = false;
            for(auto it = moves.begin(); it != moves.end(); ++it) {
                if (*it == move) {
                    if (move != FastBoard::PASS) {
                        // get x y coords for actual move
                        std::pair<int, int> xy = state->board.get_xy(move);
                        this_move = (xy.second * 19) + xy.first;
                    }
                    moveseen = true;
                }
            }

            bool has_next_moves = counter + 2 < tree_moves.size();
            if (!has_next_moves) {
                goto skipnext;
            }

            has_next_moves  = tree_moves[counter + 1] != FastBoard::PASS;
            has_next_moves &= tree_moves[counter + 2] != FastBoard::PASS;

            if (!has_next_moves) {
                goto skipnext;
            }

            //int skip = Random::get_Rng()->randfix<8>();
            if (/*skip == 0*/ 1) {
                if (moveseen && move != FastBoard::PASS && has_next_moves) {
                    TrainPosition position;
                    //position.stm_won = (tomove == who_won ? 1.0f : 0.0f);
                    //position.stm_won_tanh = (tomove == who_won ? 1.0f : -1.0f);
                    //float frac = (float)counter / (float)movecount;
                    //position.stm_score = (frac * position.stm_won)
                        //+ ((1.0f - frac) * 0.5f);
                    //position.stm_score_tanh = (frac * position.stm_won_tanh)
                        //+ ((1.0f - frac) * 0.0f);
                    //gather_features_value(state, position.planes);
                    position.moves[0] = this_move;
                    // add next 2 moves to position
                    // we do not check them for legality
                    int next_move = tree_moves[counter + 1];
                    int next_next_move = tree_moves[counter + 2];
                    std::pair<int, int> xy = state->board.get_xy(next_move);
                    position.moves[1] = (xy.second * 19) + xy.first;
                    xy = state->board.get_xy(next_next_move);
                    position.moves[2] = (xy.second * 19) + xy.first;
                    gather_features_policy(state, position.planes);
                    data.push_back(position);
                } else if (move != FastBoard::PASS) {
                    myprintf("Mainline move not found: %d\n", move);
                    goto skipnext;
                }
            }

            counter++;
            treewalk = treewalk->get_child(0);
        }

skipnext:
        gamecount++;
        if (gamecount % 100 == 0) {
            myprintf("Game %d, %d new positions, %d total\n",
                     gamecount, data.size(), train_pos + data.size());
        }
        if (gamecount % (50000) == 0) {
            train_network(data, train_pos, test_pos);
        }
    }

    train_network(data, train_pos, test_pos);

    std::cout << train_pos << " training positions." << std::endl;
    std::cout << test_pos << " testing positions." << std::endl;

    myprintf("Gathering pass done.\n");
}

int Network::rev_rotate_nn_idx(const int vertex, int symmetry) {
    static const int invert[] = {0, 1, 2, 3, 4, 6, 5, 7};
    assert(rotate_nn_idx(rotate_nn_idx(vertex, symmetry), invert[symmetry])
           == vertex);
    return rotate_nn_idx(vertex, invert[symmetry]);
}

int Network::rotate_nn_idx(const int vertex, int symmetry) {
    assert(vertex >= 0 && vertex < 19*19);
    assert(symmetry >= 0 && symmetry < 8);
    int x = vertex % 19;
    int y = vertex / 19;
    int newx;
    int newy;

    if (symmetry >= 4) {
        std::swap(x, y);
        symmetry -= 4;
    }

    if (symmetry == 0) {
        newx = x;
        newy = y;
    } else if (symmetry == 1) {
        newx = x;
        newy = 19 - y - 1;
    } else if (symmetry == 2) {
        newx = 19 - x - 1;
        newy = y;
    } else {
        assert(symmetry == 3);
        newx = 19 - x - 1;
        newy = 19 - y - 1;
    }

    int newvtx = (newy * 19) + newx;
    assert(newvtx >= 0 && newvtx < 19*19);
    return newvtx;
}

void Network::train_network(TrainVector& data,
                            size_t& total_train_pos,
                            size_t& total_test_pos) {
#ifdef USE_CAFFE
    size_t data_size = data.size();
    size_t traincut = (data_size * 98) / 100;

    size_t train_pos = 0;
    size_t test_pos = 0;

    std::unique_ptr<caffe::db::DB> train_db(caffe::db::GetDB("leveldb"));
    std::string dbTrainName("leela_train");
    train_db->Open(dbTrainName.c_str(), caffe::db::WRITE);
    std::unique_ptr<caffe::db::DB> train_label_db(caffe::db::GetDB("leveldb"));
    std::string dbTrainLabelName("leela_train_label");
    train_label_db->Open(dbTrainLabelName.c_str(), caffe::db::WRITE);

    std::unique_ptr<caffe::db::DB> test_db(caffe::db::GetDB("leveldb"));
    std::string dbTestName("leela_test");
    test_db->Open(dbTestName.c_str(), caffe::db::WRITE);
    std::unique_ptr<caffe::db::DB> test_label_db(caffe::db::GetDB("leveldb"));
    std::string dbTestLabelName("leela_test_label");
    test_label_db->Open(dbTestLabelName.c_str(), caffe::db::WRITE);

    std::unique_ptr<caffe::db::Transaction> train_txn(train_db->NewTransaction());
    std::unique_ptr<caffe::db::Transaction> test_txn(test_db->NewTransaction());
    std::unique_ptr<caffe::db::Transaction> train_label_txn(train_label_db->NewTransaction());
    std::unique_ptr<caffe::db::Transaction> test_label_txn(test_label_db->NewTransaction());

    std::cout << "Shuffling training data...";
    std::random_shuffle(data.begin(), data.end());
    std::cout << "writing: ";

    size_t data_pos = 0;
    for (auto it = data.begin(); it != data.end(); ++it) {
        TrainPosition& position = *it;
        NNPlanes& nnplanes = position.planes;

        // train data
        caffe::Datum datum;
        size_t datum_channels = nnplanes.size();
        datum.set_channels(datum_channels);
        datum.set_height(19);
        datum.set_width(19);
        std::string buffer(datum_channels * 19 * 19, '\0');
        // check whether to rotate the position
        int symmetry = Random::get_Rng()->randfix<8>();
        for (size_t p = 0; p < nnplanes.size(); p++) {
            BoardPlane tmp;
            for (size_t b = 0; b < nnplanes[p].size(); b++) {
                float val = nnplanes[p][b];
                int rot_idx = rotate_nn_idx((int)b, symmetry);
                tmp[rot_idx] = val;
            }
            if (p == 0) {
                assert(tmp[rot_move] == true);
            } else if (p == 1 || p == 2) {
                assert(tmp[rot_move] == false);
            }
            for (size_t b = 0; b < tmp.size(); b++) {
                buffer[(p * (19 * 19)) + b] = (int)tmp[b];
            }
        }
        datum.set_data(buffer);
        std::string out;
        datum.SerializeToString(&out);

        // labels
        caffe::Datum datum_label;
        datum_label.set_channels(3);
        datum_label.set_height(1);
        datum_label.set_width(1);

        int this_move = rotate_nn_idx(position.moves[0], symmetry);
        int next_move = rotate_nn_idx(position.moves[1], symmetry);
        int next_next_move = rotate_nn_idx(position.moves[2], symmetry);

        datum_label.add_float_data(this_move);
        datum_label.add_float_data(next_move);
        datum_label.add_float_data(next_next_move);
        //datum_label.add_float_data((float)position.stm_score);
        //datum_label.add_float_data((float)position.stm_won);
        //datum_label.add_float_data((float)position.stm_score_tanh);
        //datum_label.add_float_data((float)position.stm_won_tanh);
        std::string label_out;
        datum_label.SerializeToString(&label_out);

        data_pos++;
        if (data_pos > traincut) {
            std::stringstream ss;
            ss << (total_test_pos + test_pos);
            test_pos++;
            test_txn->Put(ss.str(), out);
            test_label_txn->Put(ss.str(), label_out);
            if (test_pos % 10000 == 0) {
                std::cout << "t";
                test_txn->Commit();
                test_label_txn->Commit();
                test_txn.reset(test_db->NewTransaction());
                test_label_txn.reset(test_label_db->NewTransaction());
            }
        } else {
            std::stringstream ss;
            ss << (total_train_pos + train_pos);
            train_pos++;
            train_txn->Put(ss.str(), out);
            train_label_txn->Put(ss.str(), label_out);
            if (train_pos % 10000 == 0) {
                std::cout << symmetry;
                train_txn->Commit();
                train_label_txn->Commit();
                train_txn.reset(train_db->NewTransaction());
                train_label_txn.reset(train_label_db->NewTransaction());
            }
        }
    }
    data.clear();

    train_txn->Commit();
    test_txn->Commit();
    train_label_txn->Commit();
    test_label_txn->Commit();

    total_train_pos += train_pos;
    total_test_pos += test_pos;

    std::cout << std::endl;
#endif
}

void Network::autotune_from_file(std::string filename) {
#ifdef USE_CAFFE
#if 0
    {
        std::unique_ptr<caffe::db::DB> train_db(caffe::db::GetDB("leveldb"));
        std::string dbTrainName("leela_train");
        train_db->Open(dbTrainName.c_str(), caffe::db::NEW);
        std::unique_ptr<caffe::db::DB> train_label_db(caffe::db::GetDB("leveldb"));
        std::string dbTrainLabelName("leela_train_label");
        train_label_db->Open(dbTrainLabelName.c_str(), caffe::db::NEW);
        std::unique_ptr<caffe::db::DB> test_db(caffe::db::GetDB("leveldb"));
        std::string dbTestName("leela_test");
        test_db->Open(dbTestName.c_str(), caffe::db::NEW);
        std::unique_ptr<caffe::db::DB> test_label_db(caffe::db::GetDB("leveldb"));
        std::string dbTestLabelName("leela_test_label");
        test_label_db->Open(dbTestLabelName.c_str(), caffe::db::NEW);
    }
#endif
#endif
    TrainVector data;
    gather_traindata(filename, data);
}

std::string Network::get_backend() {
#if defined(USE_OPENCL)
    return opencl.get_device_name();
#elif defined(USE_CAFFE)
    return std::string("Caffe");
#else
#ifdef USE_BLAS
#ifndef __APPLE__
#ifdef USE_OPENBLAS
    return std::string("BLAS core: " + std::string(openblas_get_corename()));
#endif
#ifdef USE_MKL
    MKLVersion Version;
    mkl_get_version(&Version);
    return std::string("BLAS core: " + std::string(Version.Processor));
#endif
#else
    return std::string("BLAS core: Apple Accelerate");
#endif
#endif
#endif
    return std::string("No BLAS backend active");
}


// Copyright (c) 2015-2018 Daniel Cooke
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

#include "read_assigner.hpp"

#include <utility>
#include <iterator>
#include <algorithm>
#include <limits>
#include <random>
#include <stdexcept>
#include <cassert>

#include <boost/optional.hpp>

#include "utils/maths.hpp"
#include "utils/kmer_mapper.hpp"
#include "core/models/haplotype_likelihood_model.hpp"
#include "core/models/error/error_model_factory.hpp"

namespace octopus {

namespace {

using HaplotypeLikelihoods = std::vector<std::vector<double>>;

auto vectorise(const std::vector<Haplotype>& haplotypes, const HaplotypeProbabilityMap& priors)
{
    std::vector<double> result(haplotypes.size());
    std::transform(std::cbegin(haplotypes), std::cend(haplotypes), std::begin(result),
                   [&] (const auto& haplotype) { return priors.at(haplotype); });
    return result;
}

auto get_priors(const std::vector<Haplotype>& haplotypes, const HaplotypeProbabilityMap& log_priors)
{
    if (log_priors.empty()) {
        return std::vector<double>(haplotypes.size());
    } else {
        return vectorise(haplotypes, log_priors);
    }
}

void find_map_haplotypes(const std::vector<Haplotype>& haplotypes, const unsigned read,
                         const HaplotypeLikelihoods& likelihoods, const std::vector<double>& log_priors,
                         std::vector<unsigned>& result)
{
    assert(result.empty());
    auto max_likelihood = std::numeric_limits<double>::lowest();
    for (unsigned k {0}; k < haplotypes.size(); ++k) {
        const auto curr = likelihoods[k][read] + log_priors[k];
        if (maths::almost_equal(curr, max_likelihood)) {
            result.push_back(k);
        } else if (curr > max_likelihood) {
            result.assign({k});
            max_likelihood = curr;
        }
    }
    if (result.empty()) {
        result.resize(haplotypes.size());
        std::iota(std::begin(result), std::end(result), 0);
    }
}

template <typename ForwardIt, typename RandomGenerator>
ForwardIt random_select(ForwardIt first, ForwardIt last, RandomGenerator& g)
{
    if (first == last) return first;
    const auto max = static_cast<std::size_t>(std::distance(first, last));
    if (max == 1) return first;
    std::uniform_int_distribution<std::size_t> dist {0, max - 1};
    std::advance(first, dist(g));
    return first;
}

template <typename ForwardIt>
ForwardIt random_select(ForwardIt first, ForwardIt last)
{
    static thread_local std::mt19937 generator {42};
    return random_select(first, last, generator);
}

template <typename Range>
decltype(auto) random_select(const Range& values)
{
    assert(!values.empty());
    return *random_select(std::cbegin(values), std::cend(values));
}

auto calculate_support(const std::vector<Haplotype>& haplotypes,
                       const std::vector<AlignedRead>& reads,
                       const std::vector<double>& log_priors,
                       const HaplotypeLikelihoods& likelihoods,
                       boost::optional<AmbiguousReadList&> ambiguous,
                       const AssignmentConfig& config)
{
    HaplotypeSupportMap result {};
    std::vector<unsigned> top {};
    top.reserve(haplotypes.size());
    for (unsigned i {0}; i < reads.size(); ++i) {
        const auto& read = reads[i];
        find_map_haplotypes(haplotypes, i, likelihoods, log_priors, top);
        if (top.size() == 1) {
            result[haplotypes[top.front()]].push_back(read);
        } else {
            using UA = AssignmentConfig::AmbiguousAction;
            switch (config.ambiguous_action) {
                case UA::first:
                    result[haplotypes[top.front()]].push_back(read);
                    break;
                case UA::all: {
                    for (auto idx : top) result[haplotypes[idx]].push_back(read);
                    break;
                }
                case UA::random: {
                    result[haplotypes[random_select(top)]].push_back(read);
                    break;
                }
                case UA::drop:
                default:
                    break;
            }
            if (ambiguous) {
                ambiguous->emplace_back(read);
                if (config.ambiguous_record != AssignmentConfig::AmbiguousRecord::read_only) {
                    ambiguous->back().haplotypes = std::vector<Haplotype> {};
                    ambiguous->back().haplotypes->reserve(top.size());
                    for (auto idx : top) ambiguous->back().haplotypes->push_back(haplotypes[idx]);
                }
            }
        }
        top.clear();
    }
    return result;
}

template <typename MappableTp>
GenomicRegion::Size estimate_max_indel_size(const MappableTp& mappable)
{
    const auto p = std::minmax({region_size(mappable), static_cast<GenomicRegion::Size>(sequence_size(mappable))});
    return p.second - p.first;
}

template <typename MappableTp>
auto estimate_max_indel_size(const std::vector<MappableTp>& mappables)
{
    GenomicRegion::Size result {0};
    for (const auto& mappable : mappables) {
        result = std::max(result, estimate_max_indel_size(mappable));
    }
    return result;
}

template <typename Container>
auto compute_read_hashes(const Container& reads)
{
    static constexpr unsigned char mapperKmerSize {6};
    std::vector<KmerPerfectHashes> result {};
    result.reserve(reads.size());
    std::transform(std::cbegin(reads), std::cend(reads), std::back_inserter(result),
                   [=] (const AlignedRead& read) { return compute_kmer_hashes<mapperKmerSize>(read.sequence()); });
    return result;
}

auto expand_for_alignment(const Haplotype& haplotype, const GenomicRegion& reads_region,
                          const GenomicRegion::Size indel_factor)
{
    const auto min_flank_pad = 2 * HaplotypeLikelihoodModel::pad_requirement();
    const auto& haplotype_region = mapped_region(haplotype);
    unsigned min_lhs_expansion {min_flank_pad}, min_rhs_expansion {min_flank_pad};
    if (begins_before(reads_region, haplotype_region)) {
        min_lhs_expansion += begin_distance(reads_region, haplotype_region);
    }
    if (ends_before(haplotype_region, reads_region)) {
        min_rhs_expansion += end_distance(haplotype_region, reads_region);
    }
    const auto min_expansion = std::max(min_lhs_expansion, min_rhs_expansion) + indel_factor;
    return expand(haplotype, min_expansion);
}

template <typename Container>
auto calculate_likelihoods(const std::vector<Haplotype>& haplotypes, const Container& reads,
                           HaplotypeLikelihoodModel& model)
{
    assert(!haplotypes.empty());
    const auto reads_region = encompassing_region(reads);
    const auto read_hashes = compute_read_hashes(reads);
    static constexpr unsigned char mapperKmerSize {6};
    auto haplotype_hashes = init_kmer_hash_table<mapperKmerSize>();
    HaplotypeLikelihoods result {};
    result.reserve(haplotypes.size());
    const auto indel_factor = estimate_max_indel_size(haplotypes) + estimate_max_indel_size(reads);
    for (const auto& haplotype : haplotypes) {
        const auto expanded_haplotype = expand_for_alignment(haplotype, reads_region, indel_factor);
        populate_kmer_hash_table<mapperKmerSize>(expanded_haplotype.sequence(), haplotype_hashes);
        auto haplotype_mapping_counts = init_mapping_counts(haplotype_hashes);
        model.reset(expanded_haplotype);
        std::vector<double> likelihoods(reads.size());
        std::transform(std::cbegin(reads), std::cend(reads), std::cbegin(read_hashes), std::begin(likelihoods),
                       [&] (const auto& read, const auto& read_hash) {
                           auto mapping_positions = map_query_to_target(read_hash, haplotype_hashes, haplotype_mapping_counts);
                           reset_mapping_counts(haplotype_mapping_counts);
                           return model.evaluate(read, mapping_positions);
                       });
        clear_kmer_hash_table(haplotype_hashes);
        result.push_back(std::move(likelihoods));
    }
    return result;
}

} // namespace

HaplotypeSupportMap
compute_haplotype_support(const Genotype<Haplotype>& genotype,
                          const std::vector<AlignedRead>& reads,
                          const HaplotypeProbabilityMap& log_priors,
                          HaplotypeLikelihoodModel model,
                          boost::optional<AmbiguousReadList&> ambiguous,
                          AssignmentConfig config)
{
    if (!reads.empty()) {
        if (!genotype.is_homozygous()) {
            const auto unique_haplotypes = genotype.copy_unique();
            assert(unique_haplotypes.size() > 1);
            const auto priors = get_priors(unique_haplotypes, log_priors);
            const auto likelihoods = calculate_likelihoods(unique_haplotypes, reads, model);
            return calculate_support(unique_haplotypes, reads, priors, likelihoods, ambiguous, config);
        } else if (config.ambiguous_action != AssignmentConfig::AmbiguousAction::drop) {
            HaplotypeSupportMap result {};
            result.emplace(genotype[0], reads);
            return result;
        }
    }
    return {};
}

HaplotypeSupportMap
compute_haplotype_support(const Genotype<Haplotype>& genotype,
                          const std::vector<AlignedRead>& reads,
                          AmbiguousReadList& ambiguous,
                          const HaplotypeProbabilityMap& log_priors,
                          HaplotypeLikelihoodModel model,
                          AssignmentConfig config)
{
    return compute_haplotype_support(genotype, reads, log_priors, model, ambiguous, config);
}

HaplotypeSupportMap
compute_haplotype_support(const Genotype<Haplotype>& genotype,
                          const std::vector<AlignedRead>& reads,
                          HaplotypeLikelihoodModel model,
                          AmbiguousReadList& ambiguous,
                          AssignmentConfig config)
{
    return compute_haplotype_support(genotype, reads, ambiguous, {}, model, config);
}

static HaplotypeLikelihoodModel make_default_haplotype_likelihood_model()
{
    HaplotypeLikelihoodModel::Config config {};
    config.use_mapping_quality = false;
    return {nullptr, make_indel_error_model(), config};
}

HaplotypeSupportMap
compute_haplotype_support(const Genotype<Haplotype>& genotype,
                          const std::vector<AlignedRead>& reads,
                          AmbiguousReadList& ambiguous,
                          const HaplotypeProbabilityMap& log_priors,
                          AssignmentConfig config)
{
    auto model = make_default_haplotype_likelihood_model();
    return compute_haplotype_support(genotype, reads, log_priors, model, ambiguous, config);
}

HaplotypeSupportMap
compute_haplotype_support(const Genotype<Haplotype>& genotype,
                          const std::vector<AlignedRead>& reads,
                          AmbiguousReadList& ambiguous,
                          AssignmentConfig config)
{
    auto model = make_default_haplotype_likelihood_model();
    return compute_haplotype_support(genotype, reads, model, ambiguous, config);
}

HaplotypeSupportMap
compute_haplotype_support(const Genotype<Haplotype>& genotype,
                          const std::vector<AlignedRead>& reads,
                          const HaplotypeProbabilityMap& log_priors,
                          AssignmentConfig config)
{
    auto model = make_default_haplotype_likelihood_model();
    return compute_haplotype_support(genotype, reads, log_priors, model, boost::none, config);
}

HaplotypeSupportMap
compute_haplotype_support(const Genotype<Haplotype>& genotype,
                          const std::vector<AlignedRead>& reads,
                          AssignmentConfig config)
{
    auto model = make_default_haplotype_likelihood_model();
    return compute_haplotype_support(genotype, reads, model, config);
}

HaplotypeSupportMap
compute_haplotype_support(const Genotype<Haplotype>& genotype,
                          const std::vector<AlignedRead>& reads,
                          HaplotypeLikelihoodModel model,
                          AssignmentConfig config)
{
    return compute_haplotype_support(genotype, reads, {}, std::move(model), boost::none, config);
}

HaplotypeSupportMap
compute_haplotype_support(const Genotype<Haplotype>& genotype,
                          const std::vector<AlignedRead>& reads,
                          AmbiguousReadList& ambiguous,
                          HaplotypeLikelihoodModel model,
                          AssignmentConfig config)
{
    return compute_haplotype_support(genotype, reads, std::move(model), ambiguous, config);
}

AlleleSupportMap
compute_allele_support(const std::vector<Allele>& alleles, const HaplotypeSupportMap& haplotype_support)
{
    return compute_allele_support(alleles, haplotype_support,
                                  [] (const Haplotype& haplotype, const Allele& allele) {
                                      return haplotype.includes(allele);
                                  });
}

auto copy_included(const std::vector<Allele>& alleles, const Haplotype& haplotype)
{
    std::vector<Allele> result {};
    result.reserve(alleles.size());
    std::copy_if(std::cbegin(alleles), std::cend(alleles), std::back_inserter(result),
                [&] (const auto& allele) { return haplotype.includes(allele); });
    return result;
}

struct HaveDifferentAlleles
{
    bool operator()(const Haplotype& lhs, const Haplotype& rhs) const
    {
        const auto lhs_includes = copy_included(alleles, lhs);
        const auto rhs_includes = copy_included(alleles, rhs);
        return lhs_includes != rhs_includes;
    }
    HaveDifferentAlleles(const std::vector<Allele>& alleles) : alleles {alleles} {}
    const std::vector<Allele>& alleles;
};

bool have_common_alleles(const std::vector<Haplotype>& haplotypes, const std::vector<Allele>& alleles)
{
    return std::adjacent_find(std::cbegin(haplotypes), std::cend(haplotypes), HaveDifferentAlleles {alleles}) == std::cend(haplotypes);
}

void sort_and_merge(std::deque<AlignedReadConstReference>& src, ReadRefSupportSet& dst)
{
    std::sort(std::begin(src), std::end(src));
    auto itr = dst.insert(std::end(dst), std::begin(src), std::end(src));
    std::inplace_merge(std::begin(dst), itr, std::end(dst));
}

std::size_t
try_assign_ambiguous_reads_to_alleles(const std::vector<Allele>& alleles,
                                      const AmbiguousReadList& ambiguous_reads,
                                      AlleleSupportMap& allele_support)
{
    std::size_t num_assigned {0};
    std::unordered_map<Allele, std::deque<AlignedReadConstReference>> assigned {};
    assigned.reserve(alleles.size());
    for (const auto& ambiguous_read : ambiguous_reads) {
        if (ambiguous_read.haplotypes && have_common_alleles(*ambiguous_read.haplotypes, alleles)) {
            const auto supported_alleles = copy_included(alleles, ambiguous_read.haplotypes->front());
            for (const auto& allele : supported_alleles) {
                assigned[allele].emplace_back(ambiguous_read.read);
            }
        }
    }
    for (auto& p : assigned) sort_and_merge(p.second, allele_support[p.first]);
    return num_assigned;
}

AlleleSupportMap
compute_allele_support(const std::vector<Allele>& alleles,
                       const HaplotypeSupportMap& haplotype_support,
                       const AmbiguousReadList& ambiguous_reads)
{
    auto result = compute_allele_support(alleles, haplotype_support);
    try_assign_ambiguous_reads_to_alleles(alleles, ambiguous_reads, result);
    return result;
}

} // namespace octopus


#pragma once

// std
#include <iostream>
#include <string>
#include <vector>
#include <memory>
#include <cassert>
#include <tuple>
#include <unordered_map>
#include <math.h>
// ros
#include <ros/ros.h>
#include <ros/console.h>
#include <ros/time.h>
#include <message_filters/subscriber.h>
#include <message_filters/time_synchronizer.h>
#include <message_filters/sync_policies/exact_time.h>
#include <message_filters/sync_policies/approximate_time.h>
#include <sensor_msgs/Image.h>
#include <sensor_msgs/CameraInfo.h>
#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/PointCloud.h>
#include <sensor_msgs/point_cloud2_iterator.h>
#include <cv_bridge/cv_bridge.h>
#include <image_geometry/pinhole_camera_model.h>
#include <std_msgs/Float32MultiArray.h>
#include <std_msgs/MultiArrayDimension.h>

#include <Eigen/Dense>
#include <Eigen/Core>

// opencv
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/core/matx.hpp>
#include <opencv2/core/eigen.hpp>

// boost
#include <boost/shared_ptr.hpp>

// tensorflow
//#include <tensorflow/c/c_api.h>

// for the PointSegmentedDistribution class
#include "PointSegmentedDistribution.hpp"
#include "SegmentationMapping/ImageLabelDistribution.h"

//namespace tf = tensorflow;
//namespace tf_ops = tensorflow::ops;

namespace SegmentationMapping {

  typedef message_filters::sync_policies::ExactTime
  <sensor_msgs::Image, sensor_msgs::Image, sensor_msgs::CameraInfo, sensor_msgs::Image, ImageLabelDistribution> sync_pol;

  template <unsigned int NUM_CLASS>
  class StereoSegmentation {
  public:
    StereoSegmentation()
      : nh_()
      , color_topic_("/camera/color/image_raw")
      , depth_topic_("/camera/aligned_depth_to_color/image_raw")
      , depth_cam_topic_("/camera/aligned_depth_to_color/camera_info")
      , label_topic_("/labeled_image")
      , distribution_topic_("/distribution_image")
    {
      ros::NodeHandle pnh("~");
      pnh.getParam("color_topic", color_topic_);
      pnh.getParam("depth_topic", depth_topic_);
      pnh.getParam("depth_cam_topic", depth_cam_topic_);
      pnh.getParam("label_topic", label_topic_);
      pnh.getParam("distribution_topic", distribution_topic_);

      // init the message filter for image, depth, camera info
      color_sub_ = new message_filters::Subscriber<sensor_msgs::Image> (pnh, color_topic_, 50);
      depth_sub_ = new message_filters::Subscriber<sensor_msgs::Image> (pnh, depth_topic_, 50);
      depth_cam_sub_ = new message_filters::Subscriber<sensor_msgs::CameraInfo> (pnh, depth_cam_topic_, 50);
      label_sub_ = new message_filters::Subscriber<sensor_msgs::Image> (pnh, label_topic_, 5);
      distribution_sub_ = new message_filters::Subscriber<ImageLabelDistribution> (pnh, distribution_topic_, 5);
      sync_ = new message_filters::Synchronizer<sync_pol> (sync_pol(300), *depth_sub_, *color_sub_, *depth_cam_sub_, *label_sub_, *distribution_sub_);

      sync_->registerCallback(boost::bind(&StereoSegmentation::DepthColorCallback, this,_1, _2, _3, _4, _5));

      pc1_publisher_ = pnh.advertise<sensor_msgs::PointCloud>("/labeled_pointcloud", 1);
      pc2_publisher_ = pnh.advertise<sensor_msgs::PointCloud2>("/labeled_pointcloud_color_pc2", 1);

      num_skip_frames = pnh.getParam("skip_every_k_frame", num_skip_frames);

      /*
      // init tensorflow
      std::string input_tensor_name, output_distribution_tensor_name, output_label_tensor_name;
      pnh.getParam("tf_input_tensor", input_tensor_name);
      //pnh.getParam("tf_label_output_tensor", output_label_tensor_name);
      pnh.getParam("tf_distribution_output_tensor", output_distribution_tensor_name);
      std::string frozen_graph_path;
      pnh.getParam("tf_frozen_graph_path", frozen_graph_path );
      ROS_DEBUG("Init tf_Inference....");
      tf_infer.reset(new tfInference(frozen_graph_path, input_tensor_name,
                                   output_label_tensor_name, output_distribution_tensor_name));
      ROS_DEBUG_STREAM("Init tensorflow and revoke frozen graph from "<<frozen_graph_path);
      */
      label2color[2]  =std::make_tuple(250, 250, 250 ); // road
      label2color[3]  =std::make_tuple(250, 250, 250 ); // sidewalk
      label2color[5]  =std::make_tuple(250, 128, 0   ); // building
      label2color[10] =std::make_tuple(192, 192, 192 ); // pole
      label2color[12] =std::make_tuple(250, 250, 0   ); // sign
      label2color[6]  =std::make_tuple(0  , 100, 0   ); // vegetation
      label2color[4]  =std::make_tuple(128, 128, 0   ); // terrain
      label2color[13] =std::make_tuple(135, 206, 235 ); // sky
      label2color[1]  =std::make_tuple( 30, 144, 250 ); // water
      label2color[8]  =std::make_tuple(220, 20,  60  ); // person
      label2color[7]  =std::make_tuple( 0, 0,142     ); // car
      label2color[9]  =std::make_tuple(119, 11, 32   ); // bike
      label2color[11] =std::make_tuple(123, 104, 238 ); // stair
      label2color[0]  =std::make_tuple(255, 255, 255 ); // background

    }

    ~StereoSegmentation() {
      delete sync_;
      delete color_sub_;
      delete depth_sub_;
      delete depth_cam_sub_;
    }

    void DepthColorCallback(const sensor_msgs::ImageConstPtr& depth_msg,
                            const sensor_msgs::ImageConstPtr& color_msg,
                            const sensor_msgs::CameraInfoConstPtr& camera_info_msg,
                            const sensor_msgs::ImageConstPtr& labeled_msg,
                            const ImageLabelDistributionConstPtr & distribution_msg);

    void Depth2PointCloud1(const sensor_msgs::ImageConstPtr& depth_msg,
                           const sensor_msgs::ImageConstPtr& color_msg,
                           bool publish_semantic,
                           const cv::Mat & label_img,
                           const cv::Mat & distribution_img);
    void Depth2PointCloud2(const sensor_msgs::ImageConstPtr& depth_msg,
                           const sensor_msgs::ImageConstPtr& color_msg,
                           bool publish_semantic,
                           const cv::Mat & label_img);

  private:
    ros::NodeHandle nh_;
    std::string color_topic_;
    std::string depth_topic_;
    std::string depth_cam_topic_;
    std::string label_topic_;
    std::string distribution_topic_;

    message_filters::Subscriber<sensor_msgs::Image>* color_sub_;
    message_filters::Subscriber<sensor_msgs::Image> *depth_sub_;
    message_filters::Subscriber<sensor_msgs::CameraInfo>* depth_cam_sub_;
    message_filters::Subscriber<sensor_msgs::Image> *label_sub_;
    message_filters::Subscriber<ImageLabelDistribution> *distribution_sub_;
    message_filters::Synchronizer<sync_pol>* sync_;
    //message_filters::TimeSynchronizer<sync_pol>* sync_;
    //sync_pol* sync_;

    ros::Publisher pc1_publisher_;
    ros::Publisher pc2_publisher_;
    // For camera depth to point cloud
    image_geometry::PinholeCameraModel model_;
    ros::Subscriber label_sub;

    int num_skip_frames;
    //std::unique_ptr<tfInference> tf_infer;
    std::unordered_map<int, std::tuple<uint8_t, uint8_t, uint8_t>> label2color;
  };

  template <unsigned int NUM_CLASS>
  inline void
  StereoSegmentation<NUM_CLASS>::DepthColorCallback(const sensor_msgs::ImageConstPtr& depth_msg,
                                                    const sensor_msgs::ImageConstPtr& color_msg,
                                                    const sensor_msgs::CameraInfoConstPtr& camera_info_msg,
                                                    const sensor_msgs::ImageConstPtr& labeled_msg,
                                                    const ImageLabelDistributionConstPtr & distribution_msg) {
    ros::Time curr_t = ros::Time::now();
    ROS_DEBUG_STREAM("Callback starts at time "<<uint32_t(curr_t.toSec())<<". "<<(uint32_t)curr_t.toNSec() );

    // std::cout << "depth: " << depth_msg->header.stamp << std::endl;
    // std::cout << "color: " << color_msg->header.stamp << std::endl;
    // std::cout << "camera_info: " << camera_info_msg->header.stamp << std::endl;
    // std::cout << "labeled_msg: " << labeled_msg->header.stamp << std::endl;
    // std::cout << "distribution_msg: " << distribution_msg->header.stamp << std::endl;

    // std::cout<<"DepthColorCallback: New callback"<< depth_msg->header.frame_id <<"\n";
    // Check for bad inputs
    if (depth_msg->header.frame_id != color_msg->header.frame_id) {
      ROS_ERROR("Depth iamge frame id [%s] doesn't match color image frame id [%s]",
          depth_msg->header.frame_id.c_str(), color_msg->header.frame_id.c_str());
      return;
    }

    // Get rgb and depth images
    cv_bridge::CvImagePtr color_ptr;
    cv_bridge::CvImagePtr depth_ptr;
    cv_bridge::CvImagePtr label_ptr;
    try{
      color_ptr = cv_bridge::toCvCopy(color_msg, sensor_msgs::image_encodings::RGB8);
      depth_ptr = cv_bridge::toCvCopy(depth_msg, sensor_msgs::image_encodings::TYPE_16UC1);
      label_ptr = cv_bridge::toCvCopy(labeled_msg, sensor_msgs::image_encodings::TYPE_8UC1);
    } catch (cv_bridge::Exception& e) {
      ROS_ERROR("cv_bridge exception: %s", e.what());
      return;
    }
    //;cv::Mat color = color_ptr->image;
    //cv::Mat depth = depth_ptr->image;

    int rows = color_ptr->image.rows;
    int cols = color_ptr->image.cols;

    //float * buffer = new float [distribution_msg->distribution.data.layout.dims[0].stride];
    std::vector<float> data = distribution_msg->distribution.data;

    //  TensorMap<Tensor<int, 4>> t_4d(storage, 2, 4, 2, 8);
    //Eigen::Map<Eigen::MatrixXf> mat(data.data(), h, w);
    //cv::Mat distribution_output = cv::Mat(rows, cols, CV_32FC(uint32_t(distribution_msg->distribution.layout.dim[2].size)), distribution_msg->distribution.data).clone();
    cv::Mat distribution_output = cv::Mat(rows, cols, CV_32FC(NUM_CLASS), const_cast<float*>(distribution_msg->distribution.data.data())).clone();
    //std::cout<<"distribution.data size "<<distribution_output.total() * distribution_output.elemSize()<<std::endl;

    //tf_infer->segmentation(color, 20, label_output, distribution_output);

    // Update camera model
    this->model_.fromCameraInfo(camera_info_msg);

    Depth2PointCloud1(depth_msg, color_msg, true, label_ptr->image, distribution_output);

    // Debugging
#ifndef NDEBUG
    Depth2PointCloud2(depth_msg, color_msg, false, label_ptr->image);
#endif
    //curr_t = ros::Time::now();
    //ROS_DEBUG_STREAM("Callback ends at time "<<uint32_t(curr_t.toSec())<<". "<<(uint32_t)curr_t.toNSec() );

  }

  template <unsigned int NUM_CLASS>
  inline void
  StereoSegmentation<NUM_CLASS>::Depth2PointCloud1(const sensor_msgs::ImageConstPtr& depth_msg,
                                                   const sensor_msgs::ImageConstPtr& color_msg,
                                                   bool publish_semantic,
                                                   const cv::Mat & label_img,
                                                   const cv::Mat & distribution_img) {
    std::cout<<"New image with label\n";

    // Set up to-publish cloud msg
    sensor_msgs::PointCloud::Ptr cloud_msg(new sensor_msgs::PointCloud);
    cloud_msg->header = depth_msg->header;  // use depth image time stamp
    sensor_msgs::ChannelFloat32 label_channel;
    label_channel.name = "labels";
    sensor_msgs::ChannelFloat32 r_channel;
    r_channel.name = "r";
    sensor_msgs::ChannelFloat32 g_channel;
    g_channel.name = "g";
    sensor_msgs::ChannelFloat32 b_channel;
    b_channel.name = "b";
    sensor_msgs::ChannelFloat32 r_original_channel;
    r_channel.name = "r_original";
    sensor_msgs::ChannelFloat32 g_original_channel;
    g_channel.name = "g_original";
    sensor_msgs::ChannelFloat32 b_original_channel;
    b_channel.name = "b_original";

    sensor_msgs::ChannelFloat32 distribution_channel[NUM_CLASS];

    // Use correct principal point from camera info
    float center_x = model_.cx();
    float center_y = model_.cy();

    // Combine unit conversion with scaling by focal length for computing X, Y
    double unit_scaling = 0.001f;
    float constant_x = unit_scaling / model_.fx();
    float constant_y = unit_scaling / model_.fy();
    float bad_point = std::numeric_limits<float>::quiet_NaN();

    const uint16_t* depth_row = reinterpret_cast<const uint16_t*>(&depth_msg->data[0]);
    int row_step = depth_msg->step / sizeof(uint16_t);
    const uint8_t* color = &color_msg->data[0];
    int color_step = 3;  // 1 if mono
    int color_skip = color_msg->step - color_msg->width * color_step;

    cv::Mat distribution_exp;
    if (publish_semantic){
      //cv::cv2eigen(distribution_img, distribution_exp);
      distribution_exp  = distribution_img;
      //cv::exp(distribution_img, distribution_exp);
      std::cout<<"elemSize: "<<distribution_exp.elemSize()<<", total "<<distribution_exp.total()<<", channels "<<distribution_exp.channels()<<std::endl;
    }

    // Iterate through depth image
    int i = 0;  // num of point
    int points_count = 0;
    for (int v = 0; v < int(depth_msg->height); ++v, depth_row += row_step, color += color_skip){
      for (int u = 0; u < int(depth_msg->width); ++u, color += color_step, ++i) {

      // Skip the upper half of the image
      //if (v < int(depth_msg->height / 2))
      //  continue;

      uint16_t depth = depth_row[u];

      // Check for invalid measurements
      // Kaiduo Fang: max_depth = 9;
      geometry_msgs::Point32 p;
      if (depth <= 0 ) {
        p.x = p.y = p.z = bad_point;
      }
      else {
        // Fill in XYZ
        p.x = (u - center_x) * depth * constant_x;
        p.y = (v - center_y) * depth * constant_y;
        p.z = (float) depth * unit_scaling;
        points_count += 1;
      }

      if (depth > 0)
      	if (sqrt(p.x * p.x + p.y * p.y + p.z * p.z) > 9)
          p.x = p.y = p.z = bad_point;

      cloud_msg->points.push_back(p);

      if (publish_semantic) {

        // Fill in semantics
        //cv::Mat dist_cv =distribution_exp(cv::Rect(u, v, 1, 1));
        cv::Vec<float, NUM_CLASS> dist_class = distribution_exp.at<cv::Vec<float, NUM_CLASS>>(v, u);
        /*        float sums = cv::sum(dist_class)[0];
        if (sums < 0.99 || sums > 1.01) {

          std::cerr<<"Wrong sum out of range: "<<sums<<std::endl;
          }*/
        //Eigen::VectorXf dist;
        //Eigen::Map<Eigen::Matrix<float, 20, 1> > eigenT( dist_cv.data );
        //dist = eigenT;
        //cv::cv2eigen(dist_cv, dist);
        //Eigen::VectorXf dist = distribution_exp(u, v);
        //Eigen::VectorXf dist_class = dist.segment(0, NUM_CLASS);
        /*
        Eigen::VectorXf dist_background = dist.segment(NUM_CLASS, dist.size()-NUM_CLASS);
        //assume background is 0
        float sum_background = dist_background.sum() + dist_class(0);
        dist_class(0) = sum_background;
        float sum_all_exp = dist_class.sum();
        dist_class = (dist_class / sum_all_exp).eval();
        */
        int max_ind =0;
        float max_elem = 0;
        for (int i = 0; i < NUM_CLASS; i++) {
          if (dist_class(i) > max_elem) {
            max_ind = i;
            max_elem = dist_class(i);
          }
          distribution_channel[i].values.push_back( dist_class(i) );
        }

        int label = max_ind;
        //std::cout<<"label is "<<label<<std::endl;
        label_channel.values.push_back(label);
        r_channel.values.push_back(std::get<0>(label2color[label]));
        g_channel.values.push_back(std::get<1>(label2color[label]));
        b_channel.values.push_back(std::get<2>(label2color[label]));
        r_original_channel.values.push_back(color[0]);
        g_original_channel.values.push_back(color[1]);
        b_original_channel.values.push_back(color[2]);

      } else {
        // Fill in r g b channels: bgr
        r_channel.values.push_back(color[0]);
        g_channel.values.push_back(color[1]);
        b_channel.values.push_back(color[2]);
        // Fill in semantics
        label_channel.values.push_back(1);
        distribution_channel[0].values.push_back(1);
        for (int i = 1; i < NUM_CLASS; i++) {
          distribution_channel[i].values.push_back(0);
        }
      }
    }
   }
    std::cout << "count_points: " << points_count << std::endl << std::endl;

    cloud_msg->channels.push_back(label_channel);
    cloud_msg->channels.push_back(r_channel);
    cloud_msg->channels.push_back(g_channel);
    cloud_msg->channels.push_back(b_channel);
    cloud_msg->channels.push_back(r_original_channel );
    cloud_msg->channels.push_back(g_original_channel);
    cloud_msg->channels.push_back(b_original_channel);
    for (int i = 0; i < NUM_CLASS; i++)
      cloud_msg->channels.push_back(distribution_channel[i]);

    pc1_publisher_.publish(cloud_msg);
  }

  template <unsigned int NUM_CLASS>
  inline void
  StereoSegmentation<NUM_CLASS>::Depth2PointCloud2(const sensor_msgs::ImageConstPtr& depth_msg,
                                                   const sensor_msgs::ImageConstPtr& color_msg,
                                                   bool publish_semantic,
                                                   const cv::Mat & label_img) {
   // Set up to-publish cloud msg
   sensor_msgs::PointCloud2::Ptr cloud_msg (new sensor_msgs::PointCloud2);
   cloud_msg->header = depth_msg->header;  // use depth image time stamp

   cloud_msg->height = depth_msg->height;
   cloud_msg->width = depth_msg->width;
   cloud_msg->is_dense = false;
   cloud_msg->is_bigendian = false;

   sensor_msgs::PointCloud2Modifier pcd_modifier(*cloud_msg);
   pcd_modifier.setPointCloud2FieldsByString(2, "xyz", "rgb");

   // Use correct principal point from camera info
   float center_x = model_.cx();
   float center_y = model_.cy();

   // Combine unit conversion with scaling by focal length for computing X,Y
   double unit_scaling = 0.001f;
   float constant_x = unit_scaling / model_.fx();
   float constant_y = unit_scaling / model_.fy();
   float bad_point = std::numeric_limits<float>::quiet_NaN();

   const uint16_t* depth_row = reinterpret_cast<const uint16_t*>(&depth_msg->data[0]);
   int row_step = depth_msg->step / sizeof(uint16_t);
   const uint8_t* color = &color_msg->data[0];
   int color_step = 3;  // 1 if mono
   int color_skip = color_msg->step - color_msg->width * color_step;

   sensor_msgs::PointCloud2Iterator<float> iter_x(*cloud_msg, "x");
   sensor_msgs::PointCloud2Iterator<float> iter_y(*cloud_msg, "y");
   sensor_msgs::PointCloud2Iterator<float> iter_z(*cloud_msg, "z");
   sensor_msgs::PointCloud2Iterator<uint8_t> iter_r(*cloud_msg, "r");
   sensor_msgs::PointCloud2Iterator<uint8_t> iter_g(*cloud_msg, "g");
   sensor_msgs::PointCloud2Iterator<uint8_t> iter_b(*cloud_msg, "b");
   sensor_msgs::PointCloud2Iterator<uint8_t> iter_a(*cloud_msg, "a");

   // Iterate through depth image
   for (int v = 0 ; v < int(cloud_msg->height); ++v, depth_row += row_step, color += color_skip){
    for (int u = 0; u < int(cloud_msg->width); ++u, color += color_step,
        ++iter_x, ++iter_y, ++iter_z, ++iter_a, ++iter_r, ++iter_g, ++iter_b) {

      // Skip the upper half of the image
      //if (v < int(cloud_msg->height / 2))
      // continue;

      uint16_t depth = depth_row[u];

      // Check for invalid measurements
      // update by Kaiduo Fang: depth > 9, max_depth;
      if (depth <= 0)
        *iter_x = *iter_y = *iter_z = bad_point;
      else {
        // Fill in XYZ
        *iter_x = (u - center_x) * depth * constant_x;
        *iter_y = (v - center_y) * depth * constant_y;
        *iter_z = (float) depth * unit_scaling;
        if (*iter_z > 10)
        *iter_z = bad_point;
      }

      // Fill in color
      //label_channel.values.push_back(label);
      if (publish_semantic) {
      uint8_t label = label_img.at<uint8_t>(u, v);

      *iter_a = 255;
      *iter_r = std::get<0>(label2color[label]);
      *iter_g = std::get<1>(label2color[label]);
      *iter_b = std::get<2>(label2color[label]);
      } else {
      *iter_a = 255;
      *iter_r = color[0];
      *iter_g = color[1];
      *iter_b = color[2];

      }

    }
   }
    pc2_publisher_.publish(cloud_msg);
  }

}