/**
 * Copyright 2004-present Facebook. All Rights Reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include "source/calibration/GeometricCalibration.h"

#include <future>
#include <iostream>
#include <random>
#include <unordered_map>

#include <boost/algorithm/string.hpp>
#include <boost/timer/timer.hpp>
#include <gflags/gflags.h>
#include <glog/logging.h>
#include <opencv2/opencv.hpp>

#include <folly/FileUtil.h>
#include <folly/Format.h>
#include <folly/dynamic.h>
#include <folly/json.h>

#include "source/calibration/Calibration.h"
#include "source/util/Camera.h"
#include "source/util/CvUtil.h"
#include "source/util/MathUtil.h"
#include "source/util/ThreadPool.h"

using namespace fb360_dep;
using namespace fb360_dep::calibration;

DEFINE_int32(cap_traces, 0, "speed up solver by capping the number of traces");
DEFINE_double(ceres_function_tolerance, 1e-6, "ceres function tolerance");
DEFINE_int32(
    ceres_threads,
    -1,
    "number of threads used by ceres. requires compiled support for multithreading (default 1)");
DEFINE_string(debug_dir, "", "path to debug output");
DEFINE_double(debug_error_scale, 0, "show scaled reprojection errors");
DEFINE_double(debug_matches_overlap, 1, "show matches if overlap exceeds this fraction");
DEFINE_bool(
    dir_per_frame,
    false,
    "is there a directory per frame?\n"
    "i.e. is an image path of the form: \n"
    "    <frame index>/ ... /<camera id>.<extension>\n"
    "    e.g. 1/cam2.bmp or 000001/isp_out/cam14.png\n"
    "the default is a directory per camera. i.e. an image is of the form:\n"
    "    .../<camera id>/<frame index>.<extension>\n"
    "    e.g. cam2/123.bmp or rgb/cam14/000123.png");
DEFINE_bool(discard_outside_fov, true, "discard matches outside fov");
DEFINE_string(errors_dir, "", "directory where errors will be saved");
DEFINE_int32(experiments, 1, "calibrate multiple times");
DEFINE_bool(force_in_front, true, "no intersections behind camera");
DEFINE_bool(keep_invalid_traces, false, "keep traces with multiple points from the same camera");
DEFINE_bool(lock_distortion, true, "lock the distorion");
DEFINE_bool(lock_focal, false, "lock the focal");
DEFINE_bool(lock_positions, true, "don't calibrate position");
DEFINE_bool(lock_principals, false, "don't calibrate principals");
DEFINE_bool(lock_rotations, false, "don't calibrate rotation");
DEFINE_double(max_error, 0.5, "maximum allowable error for calibration to be valid");
DEFINE_int32(min_traces, 10, "minimum number of traces for camera to be sufficiently constrained");
DEFINE_double(outlier_factor, 5, "reject if error is factor * median");
DEFINE_double(
    outlier_z_threshold,
    3,
    "z score threshold on traces to consider a camera an outlier");
DEFINE_int32(pass_count, 10, "number of passes");
DEFINE_double(perturb_focals, 0, "pertub focals (pixels / radian)");
DEFINE_double(perturb_positions, 0, "perturb positions (m)");
DEFINE_double(perturb_principals, 0, "pertub principals (pixels)");
DEFINE_double(perturb_rotations, 0, "perturb rotations (radians)");
DEFINE_int32(point_count, 10000, "artificial points to generate");
DEFINE_double(point_error_stddev, 0.5, "error added to artificial points");
DEFINE_double(point_min_dist, 1, "minimum distance of artificial points");
DEFINE_string(points_file, "", "path to output calibration points file, default next to output");
DEFINE_string(
  points_file_json,
  "",
  "path to output calibration points file including reference points, default next to output");
DEFINE_string(reference_camera, "", "reference camera to lock if positions are unlocked");
DEFINE_double(
    remove_sparse_overlaps,
    0,
    "reject overlaps with fewer than this fraction of the average match count");
DEFINE_bool(report_per_camera_errors, false, "per camera reprojection error statistics");
DEFINE_bool(robust, true, "use Huber loss function");
DEFINE_int32(seed, -1, "seed for random number generator");
DEFINE_bool(shared_distortion, true, "all cameras in a group share the same distortion");
DEFINE_bool(
    shared_principal_and_focal,
    false,
    "all cameras in a group share the same focal, principal");
DEFINE_bool(
    weight_by_trace_count,
    false,
    "weight the residual error by the number of traces per camera");
DEFINE_bool(weighted_statistics, false, "compute statistics of weighted residuals");

std::unordered_map<std::string, int> cameraIdToIndex;
std::unordered_map<std::string, int> cameraGroupToIndex;

std::string imageIdFormat() {
  return FLAGS_dir_per_frame ? "<frame index>/ ... /<camera id>.<extension>"
                             : ".../<camera id>/<frame index>.<extension>";
}

void buildCameraIndexMaps(const Camera::Rig& rig) {
  for (int i = 0; i < int(rig.size()); ++i) {
    cameraIdToIndex[rig[i].id] = i;
    cameraGroupToIndex[rig[i].group] = i; // last camera in group wins
  }
}

using ImageId = std::string;

std::string getCameraId(const ImageId& image) {
  // image is actually a path
  filesystem::path path = image;
  if (FLAGS_dir_per_frame) {
    return path.stem().native();
  }
  return path.parent_path().filename().native();
}

int getFrameIndex(const ImageId& image) {
  // image is actually a path
  filesystem::path path = image;
  if (FLAGS_dir_per_frame) {
    return std::stoi(path.begin()->native());
  }
  return std::stoi(path.stem().native());
}

bool hasCameraIndex(const ImageId& image) {
  return cameraIdToIndex.count(getCameraId(image));
}

int getCameraIndex(const ImageId& image) {
  return cameraIdToIndex.at(getCameraId(image));
}

// create a string that adheres to the format of an image path
ImageId makeArtificialPath(int frame, const std::string& cameraId) {
  if (FLAGS_dir_per_frame) {
    return std::to_string(frame) + "/" + cameraId;
  }
  return cameraId + "/" + std::to_string(frame);
}

// input path includes basename and extension
cv::Mat_<cv::Vec3w> loadImage(const filesystem::path& path) {
  const filesystem::path colorDir = FLAGS_color;
  return cv_util::loadImage<cv::Vec3w>(colorDir / path);
}

folly::dynamic parseJsonFile(const std::string& path) {
  std::string json;
  folly::readFile(path.c_str(), json);
  CHECK(!json.empty()) << "could not read JSON file: " << path;
  return folly::parseJson(json);
}

// a feature is a point in an image
struct Feature {
  Camera::Vector2 position; // position of the feature in its image, in pixels
  int trace; // index of trace that feature belongs to (or -1 if none)

  Feature(const Camera::Vector2& position) : position(position), trace(-1) {}
};

// a featuremap holds, for each image, a vector of its features
using FeatureMap = std::unordered_map<ImageId, std::vector<Feature>>;

// an overlap is a pair of images and the matches between their features
struct Overlap {
  std::array<ImageId, 2> images;
  std::vector<std::array<size_t, 2>> matches;

  Overlap(const ImageId& image0, const ImageId& image1) {
    images[0] = image0;
    images[1] = image1;
  }

  bool isIntraFrame() const {
    return getFrameIndex(images[0]) == getFrameIndex(images[1]);
  }
};

// a trace is a world coordinate and a list of observations that reference it
struct Trace {
  Camera::Vector3 position;

  std::vector<std::pair<ImageId, int>> references;

  void add(const ImageId& image, const int index) {
    references.emplace_back(image, index);
  }

  // inherit trace's references
  void inherit(Trace& trace, FeatureMap& featureMap, int me) {
    for (const auto& ref : trace.references) {
      featureMap[ref.first][ref.second].trace = me;
    }
    references.insert(references.end(), trace.references.begin(), trace.references.end());
    trace.references.clear();
  }

  void clear(FeatureMap& featureMap) {
    for (const auto& ref : references) {
      featureMap[ref.first][ref.second].trace = -1;
    }
    references.clear();
  }
};

/* Parse features from <parsed> JSON which has structure
    {
        "images": {
            image_name: [{"x": x, "y": y, ...}]
        },
        ...
    }
    where image_name is defined as in imageIdFormat, above.
*/
FeatureMap loadFeatureMap(const folly::dynamic& parsed) {
  FeatureMap result;

  for (const auto& image : parsed["images"].items()) {
    const ImageId path = image.first.getString();
    if (!hasCameraIndex(path)) {
      LOG(INFO) << folly::sformat("ignoring image id {}", path);
      continue;
    }
    std::vector<Feature>& features = result[path];
    for (const auto& feature : image.second) {
      features.emplace_back(Camera::Vector2(feature["x"].asDouble(), feature["y"].asDouble()));
    }
  }

  CHECK(!result.empty()) << "verify image id format: " << imageIdFormat();
  LOG(INFO) << folly::sformat("{} images loaded", result.size());

  return result;
}

/* Parse matches from <parsed> JSON which has structure
    {
        "all_matches": [{
            "image1": image_name1,
            "image2": image_name1,
            "matches": [{
                "idx1": idx1,
                "idx2": idx2
            }]
        }],
        ...
    }
*/
std::vector<Overlap> loadOverlaps(const folly::dynamic& parsed) {
  std::vector<Overlap> result;

  size_t count = 0;
  for (const auto& overlap : parsed["all_matches"]) {
    ImageId path0 = overlap["image1"].getString();
    ImageId path1 = overlap["image2"].getString();
    if (!hasCameraIndex(path0) || !hasCameraIndex(path1)) {
      continue;
    }
    result.emplace_back(path0, path1);
    for (const auto& match : overlap["matches"]) {
      // A threshold of 0 indicates that score should be ignored
      // Check if score should be ignored before attempting to access
      // match["score"] since it might not be defined in the json
      if (FLAGS_match_score_threshold == 0 ||
          FLAGS_match_score_threshold <= match["score"].asDouble()) {
        result.back().matches.push_back(
            {{size_t(match["idx1"].asInt()), size_t(match["idx2"].asInt())}});
      }
    }
    count += 2 * result.back().matches.size();
  }

  LOG(INFO) << folly::sformat("{} feature observations loaded", count);

  return result;
}

Overlap& findOrAddOverlap(std::vector<Overlap>& overlaps, const ImageId& i0, const ImageId& i1) {
  for (Overlap& overlap : overlaps) {
    if (overlap.images[0] == i0 && overlap.images[1] == i1) {
      return overlap;
    }
    // make sure we don't have the image pair (i1, i0) in there
    CHECK(overlap.images[0] != i1 || overlap.images[1] != i0);
  }
  // image pair (i0, i1) not found: add it
  overlaps.emplace_back(i0, i1);
  return overlaps.back();
}

template <typename T>
Camera::Vector2 keypointError(T& gen) {
  std::normal_distribution<> rng(0, FLAGS_point_error_stddev);
  return {rng(gen), rng(gen)};
}

void generateArtificalPoints(
    FeatureMap& featureMap,
    std::vector<Overlap>& overlaps,
    const std::vector<Camera>& cameras) {
  std::mt19937 mt;
  for (int p = 0; p < FLAGS_point_count; ++p) {
    // create a random unit vector
    double longitude = std::uniform_real_distribution<>(-M_PI, +M_PI)(mt);
    double z = std::uniform_real_distribution<>(-1, 1)(mt);
    Camera::Vector3 rig(sqrt(1 - z * z) * cos(longitude), sqrt(1 - z * z) * sin(longitude), z);
    CHECK_NEAR(rig.squaredNorm(), 1, 0.001);

    // divide unit vector by random disparity
    rig /= std::uniform_real_distribution<>(0, 1 / FLAGS_point_min_dist)(mt);

    // add keypoint to every camera that sees rig
    std::vector<ImageId> images;
    for (const Camera& camera : cameras) {
      if (camera.sees(rig)) {
        ImageId image = makeArtificialPath(0, camera.id);
        featureMap[image].emplace_back(camera.pixel(rig) + keypointError(mt));
        images.push_back(image);
      }
    }

    // add a match for every pair of cameras that see rig
    for (int index1 = 0; index1 < int(images.size()); ++index1) {
      const ImageId& i1 = images[index1];
      for (int index0 = 0; index0 < index1; ++index0) {
        const ImageId& i0 = images[index0];
        Overlap& overlap = findOrAddOverlap(overlaps, i0, i1);
        overlap.matches.push_back({{featureMap[i0].size() - 1, featureMap[i1].size() - 1}});
      }
    }
  }
}

Camera::Vector3 triangulate(const Observations& observations) {
  return triangulateNonlinear(observations, FLAGS_force_in_front);
}

// return reprojection errors for each camera
std::vector<std::vector<Camera::Real>> reprojectionErrors(
    const std::vector<Overlap>& overlaps,
    const FeatureMap& featureMap,
    const std::vector<Trace>& traces,
    const std::vector<Camera>& cameras) {
  std::vector<std::vector<Camera::Real>> errors(ssize(cameras));
  for (const Overlap& overlap : overlaps) {
    if (!overlap.isIntraFrame()) {
      continue;
    }
    const std::reference_wrapper<const ImageId> images[2] = {overlap.images[0], overlap.images[1]};
    const int idxs[2] = {getCameraIndex(images[0]), getCameraIndex(images[1])};
    const std::reference_wrapper<const std::vector<Feature>> features[2] = {
        featureMap.at(images[0]), featureMap.at(images[1])};
    for (const auto& match : overlap.matches) {
      std::array<const Feature, 2> kps = {
          {features[0].get()[match[0]], features[1].get()[match[1]]}};
      CHECK_EQ(kps[0].trace, kps[1].trace) << "matching features belong to different traces";
      Camera::Vector3 rig = kps[0].trace < 0
          ? triangulate({{cameras[idxs[0]], kps[0].position}, {cameras[idxs[1]], kps[1].position}})
          : traces[kps[0].trace].position;
      for (ssize_t i = 0; i < ssize(match); ++i) {
        const Camera::Vector2 pixel = cameras[idxs[i]].pixel(rig);
        errors[idxs[i]].push_back((pixel - kps[i].position).squaredNorm());
      }
    }
  }

  return errors;
}

// report reprojection errors for each camera
void reportReprojectionErrors(
    const std::vector<Overlap>& overlaps,
    const FeatureMap& featureMap,
    const std::vector<Trace>& traces,
    const std::vector<Camera>& cameras) {
  if (!FLAGS_report_per_camera_errors) {
    return;
  }
  std::vector<std::vector<Camera::Real>> errors =
      reprojectionErrors(overlaps, featureMap, traces, cameras);
  for (ssize_t i = 0; i < ssize(cameras); ++i) {
    std::sort(errors[i].begin(), errors[i].end());
    std::ostringstream line;
    for (int percentile : {50, 90, 99}) {
      int index = percentile * (ssize(errors[i]) - 1) / 100.0 + 0.5;
      line << folly::format("{}%: {:.2f} ", percentile, sqrt(errors[i][index]));
    }
    LOG(INFO) << folly::sformat(
        "{}: {} reproj. percentile {}", cameras[i].id, ssize(errors[i]), line.str());
  }
}

void removeOutliersFromCameras(
    std::vector<Overlap>& overlaps,
    const FeatureMap& featureMap,
    const std::vector<Trace>& traces,
    const std::vector<Camera>& cameras,
    const Camera::Real outlierFactor) {
  // compute reprojection errors for each camera
  std::vector<std::vector<Camera::Real>> errors =
      reprojectionErrors(overlaps, featureMap, traces, cameras);

  // compute median for each camera
  std::vector<Camera::Real> medians(ssize(errors));
  for (ssize_t i = 0; i < ssize(errors); ++i) {
    medians[i] = calcPercentile(errors[i]);
  }

  std::unordered_map<ImageId, int> outliers;

  // remove matches that neither endpoint wants to keep around
  std::vector<ssize_t> errorIdxs(ssize(errors), 0);
  int total = 0;
  int inlierTotal = 0;
  for (Overlap& overlap : overlaps) {
    if (!overlap.isIntraFrame()) {
      continue;
    }
    const std::reference_wrapper<const ImageId> images[2] = {overlap.images[0], overlap.images[1]};
    const int idxs[2] = {getCameraIndex(images[0]), getCameraIndex(images[1])};

    // move the inliers to front of overlap.matches and resize to fit
    int inliers = 0;
    for (const auto& match : overlap.matches) {
      bool inlier = false;
      for (ssize_t i = 0; i < ssize(match); ++i) {
        int cameraIndex = idxs[i];
        if (errors[cameraIndex][errorIdxs[cameraIndex]++] < medians[cameraIndex] * outlierFactor) {
          inlier = true;
        }
      }
      if (inlier) {
        overlap.matches[inliers++] = match;
      }
    }
    outliers[overlap.images[0]] += ssize(overlap.matches) - inliers;
    outliers[overlap.images[1]] += ssize(overlap.matches) - inliers;
    total += ssize(overlap.matches);
    overlap.matches.resize(inliers);
    inlierTotal += inliers;
  }

  // sanity check that we consumed all the errors
  for (ssize_t i = 0; i < ssize(errors); ++i) {
    CHECK_EQ(errorIdxs[i], ssize(errors[i]));
  }
  if (FLAGS_log_verbose) {
    for (const auto& p : outliers) {
      LOG(INFO) << folly::sformat("Removed {} outliers from {}", p.second, p.first);
    }
  }
  LOG(INFO) << folly::sformat("{} of {} matches were inliers", inlierTotal, total);
}

void removeInvalidTraces(std::vector<Trace>& traces, FeatureMap& featureMap) {
  int total = 0;
  int removed = 0;
  for (Trace& trace : traces) {
    if (!trace.references.empty()) {
      ++total;
    }
    std::unordered_set<ImageId> uniqueCameras;
    for (const auto& ref : trace.references) {
      if (!uniqueCameras.insert(ref.first).second) {
        // image referenced more than once, remove the trace
        trace.clear(featureMap);
        ++removed;
        break;
      }
    }
  }
  LOG(INFO) << folly::sformat("removed {} out of {} traces", removed, total);
}

void triangulateTracesThread(
    std::vector<Trace>& traces,
    const size_t begin,
    const size_t end,
    const FeatureMap& featureMap,
    const std::vector<Camera>& cameras) {
  for (size_t i = begin; i < end; ++i) {
    Trace& trace = traces[i];
    if (!trace.references.empty()) {
      Observations observations;
      for (const auto& ref : trace.references) {
        const Feature& feature = featureMap.at(ref.first)[ref.second];
        const Camera& camera = cameras[getCameraIndex(ref.first)];
        observations.emplace_back(camera, feature.position);
      }
      trace.position = triangulate(observations);
    }
  }
}

void triangulateTraces(
    std::vector<Trace>& traces,
    const FeatureMap& featureMap,
    const std::vector<Camera>& cameras) {
  const int threadCount = ThreadPool::getThreadCountFromFlag(FLAGS_threads);
  std::vector<std::thread> threads;
  for (int thread = 0; thread < threadCount; ++thread) {
    threads.emplace_back(
        triangulateTracesThread,
        std::ref(traces),
        thread * traces.size() / threadCount,
        (thread + 1) * traces.size() / threadCount,
        std::cref(featureMap),
        std::cref(cameras));
  }
  for (std::thread& thread : threads) {
    thread.join();
  }
}

std::vector<Trace> assembleTraces(FeatureMap& featureMap, const std::vector<Overlap>& overlaps) {
  // mark all features as unreferenced
  for (auto& features : featureMap) {
    for (Feature& feature : features.second) {
      feature.trace = -1;
    }
  }
  std::vector<Trace> traces;
  int nonemptyTraceCount = 0;
  for (const Overlap& overlap : overlaps) {
    const std::reference_wrapper<std::vector<Feature>> features[2] = {
        featureMap[overlap.images[0]], featureMap[overlap.images[1]]};
    for (const auto& match : overlap.matches) {
      const std::reference_wrapper<int> indexes[2] = {features[0].get()[match[0]].trace,
                                                      features[1].get()[match[1]].trace};
      if (indexes[0] < 0 && indexes[1] < 0) { // neither belongs to a trace, start new trace
        traces.emplace_back();
        nonemptyTraceCount++;
        for (ssize_t i = 0; i < ssize(indexes); ++i) {
          indexes[i].get() = traces.size() - 1;
          traces[indexes[i]].add(overlap.images[i], match[i]);
        }
      } else if (indexes[0] < 0) { // 0 does not belong to a trace, add to 1's trace
        indexes[0].get() = indexes[1];
        traces[indexes[0]].add(overlap.images[0], match[0]);
      } else if (indexes[1] < 0) { // 1 does not belong to a trace, add to 0's trace
        indexes[1].get() = indexes[0];
        traces[indexes[1]].add(overlap.images[1], match[1]);
      } else if (indexes[0] != indexes[1]) { // merge two traces, 0 inherits 1's references
        traces[indexes[0]].inherit(traces[indexes[1]], featureMap, indexes[0]);
        --nonemptyTraceCount;
      }
    }
  }

  LOG(INFO) << folly::sformat("found {} nonempty traces", nonemptyTraceCount);

  return traces;
}

cv::Mat_<cv::Vec3w> blend(const cv::Mat_<cv::Vec3w>& mat0, const cv::Mat_<cv::Vec3w>& mat1) {
  if (mat0.empty()) {
    return 0.5 * mat1;
  }
  cv::Mat_<cv::Vec3w> result;
  cv::addWeighted(mat0, 0.5, mat1, 0.5, 0, result);
  return result;
}

// draw a line that starts out red, ends up green
void drawRedGreenLine(
    cv::Mat_<cv::Vec3w>& dst,
    const Camera::Vector2& r,
    const Camera::Vector2& g,
    const Camera::Vector2& m) {
  const cv::Scalar red(0, 0, 255);
  const cv::Scalar green(0, 255, 0);
  cv::line(dst, cv::Point2f(r.x(), r.y()), cv::Point2f(m.x(), m.y()), red, 2);
  cv::line(dst, cv::Point2f(g.x(), g.y()), cv::Point2f(m.x(), m.y()), green, 2);
}

template <typename T>
cv::Mat_<T> projectImageBetweenCamerasNearest(
    const Camera& dst,
    const Camera& src,
    const cv::Mat_<T>& srcImage) {
  cv::Mat_<T> dstImage(cv::Size(dst.resolution.x(), dst.resolution.y()));
  for (int y = 0; y < dstImage.rows; ++y) {
    for (int x = 0; x < dstImage.cols; ++x) {
      Camera::Vector3 rig = dst.rigNearInfinity({x + 0.5, y + 0.5});
      Camera::Vector2 srcPixel;
      if (src.sees(rig, srcPixel)) {
        dstImage(y, x) = srcImage.empty() ? cv_util::createBGR<T>(255, 255, 255)
                                          : srcImage(srcPixel.y(), srcPixel.x());
      } else {
        dstImage(y, x) = cv_util::createBGR<T>(0, 0, 0);
      }
    }
  }
  return dstImage;
}

cv::Mat_<cv::Vec3w> renderOverlap(
    const Overlap& overlap,
    const FeatureMap& featureMap,
    const std::vector<Trace>& traces,
    const std::vector<Camera>& cameras) {
  // transform image 1 into image 0's space and overlay the two
  const ImageId& image0 = overlap.images[0];
  const ImageId& image1 = overlap.images[1];
  const Camera& camera0 = cameras[getCameraIndex(image0)];
  const Camera& camera1 = cameras[getCameraIndex(image1)];
  const std::vector<Feature>& features0 = featureMap.at(image0);
  const std::vector<Feature>& features1 = featureMap.at(image1);
  cv::Mat_<cv::Vec3w> result = blend(
      loadImage(image0), projectImageBetweenCamerasNearest(camera0, camera1, loadImage(image1)));
  for (const auto& match : overlap.matches) {
    Camera::Vector2 p0 = features0[match[0]].position;
    Camera::Vector2 p1 = features1[match[1]].position;
    int trace = features0[match[0]].trace;
    CHECK_EQ(trace, features1[match[1]].trace);
    Camera::Vector3 rig =
        trace < 0 ? triangulate({{camera0, p0}, {camera1, p1}}) : traces[trace].position;
    drawRedGreenLine(
        result,
        p0, // p0 in red
        camera0.pixel(camera1.rigNearInfinity(p1)), // transformed p1 in green
        camera0.pixel(rig)); // via transformed triangulation
  }

  return result;
}

cv::Mat_<cv::Vec3w> renderReprojections(
    const ImageId& image,
    const Camera& camera,
    const std::vector<Feature>& features,
    const std::vector<Trace>& traces,
    const Camera::Real scale) {
  cv::Mat_<cv::Vec3w> result = 0.5f * loadImage(image);
  cv::Mat_<cv::Vec3f> errors = cv::Mat::zeros(result.size(), CV_32FC3);
  const cv::Scalar green(0, 255, 0);
  const cv::Scalar red(0, 0, 255);
  for (const Feature& feature : features) {
    if (feature.trace >= 0) {
      // draw red line from image feature to reprojected world point
      // then continue in green in the same direction but scale x as far
      Camera::Vector2 proj = camera.pixel(traces[feature.trace].position);
      Camera::Vector2 error = proj - feature.position;

      // OpenCV can only save 1, 3, 4 channels so the third is simply an artifact
      errors(cv::Point(feature.position.x(), feature.position.y())) =
          cv::Vec3f(error.x(), error.y(), 0.0f);
      drawRedGreenLine(result, feature.position, proj + scale * error, proj);
    }
  }

  if (FLAGS_errors_dir != "") {
    filesystem::create_directories(FLAGS_errors_dir);
    std::string errorsFile = folly::sformat("{}/{}.exr", FLAGS_errors_dir, getCameraId(image));
    cv_util::imwriteExceptionOnFail(errorsFile, errors);
  }

  return result;
}

std::string getReprojectionReport(
    const ceres::Problem& problem,
    const double* parameter = nullptr) {
  std::vector<double> norms =
      getReprojectionErrorNorms(problem, parameter, FLAGS_weighted_statistics);
  double total = 0;
  double totalSq = 0;
  for (double norm : norms) {
    total += norm;
    totalSq += norm * norm;
  }

  std::ostringstream result;
  result << "reprojections " << norms.size() << " "
         << "RMSE " << sqrt(totalSq / norms.size()) << " "
         << "average " << total / norms.size() << " "
         << "median " << calcPercentile(norms, 0.5) << " "
         << "90% " << calcPercentile(norms, 0.9) << " "
         << "99% " << calcPercentile(norms, 0.99) << " ";

  result << "worst 3: ";
  std::sort(norms.begin(), norms.end());
  for (int i = norms.size() - 3; i < int(norms.size()); ++i) {
    result << norms[i] << " ";
  }

  return result.str();
}

double acosClamp(double x) {
  return std::acos(std::min(std::max(-1.0, x), 1.0));
}

std::string getCameraRmseReport(
    const std::vector<Camera>& cameras,
    const std::vector<Camera>& groundTruth) {
  Camera::Real position = 0;
  Camera::Real rotation = 0;
  Camera::Real principal = 0;
  Camera::Real distortion = 0;
  Camera::Real focal = 0;

  for (int i = 0; i < int(cameras.size()); ++i) {
    {
      auto before = groundTruth[i].position;
      auto after = cameras[i].position;
      position += (after - before).squaredNorm();
    }
    for (int v = 0; v < 3; ++v) {
      auto before = groundTruth[i].rotation.row(v);
      auto after = cameras[i].rotation.row(v);
      rotation += (after - before).squaredNorm();
    }
    {
      auto before = groundTruth[i].principal;
      auto after = cameras[i].principal;
      principal += (after - before).squaredNorm();
    }
    {
      auto before = groundTruth[i].getDistortion();
      auto after = cameras[i].getDistortion();
      distortion += (after - before).squaredNorm();
    }
    {
      auto before = groundTruth[i].focal;
      auto after = cameras[i].focal;
      focal += (after - before).squaredNorm();
    }
  }

  Camera::Real angle = 0;
  int angleCount = 0;
  for (int i = 0; i < int(cameras.size()); ++i) {
    for (int j = 0; j < i; ++j) {
      for (int v = 2; v < 3; ++v) {
        // angle between camera i and j
        auto before = acosClamp(groundTruth[i].rotation.row(v).dot(groundTruth[j].rotation.row(v)));
        if (before > 1) {
          continue; // only count angles less than a radian
        }
        auto after = acosClamp(cameras[i].rotation.row(v).dot(cameras[j].rotation.row(v)));
        angle += (after - before) * (after - before);
        ++angleCount;
      }
    }
  }

  // average
  position /= cameras.size();
  rotation /= 3 * cameras.size();
  principal /= cameras.size();
  distortion /= cameras.size();
  focal /= cameras.size();
  angle /= angleCount;

  std::ostringstream result;
  result << "RMSEs: "
         << "Pos " << sqrt(position) << " "
         << "Rot " << sqrt(rotation) << " "
         << "Principal " << sqrt(principal) << " "
         << "Distortion " << sqrt(distortion) << " "
         << "Focal " << sqrt(focal) << " "
         << "Angle " << sqrt(angle) << " ";

  return result.str();
}

template <typename T>
double* parameterBlock(T& t) {
  return t.data();
}

template <>
double* parameterBlock(Camera::Real& t) {
  return &t;
}

template <>
double* parameterBlock(Trace& t) {
  return t.position.data();
}

template <typename T>
void lockParameter(ceres::Problem& problem, T& param, const bool lock = true) {
  if (lock) {
    problem.SetParameterBlockConstant(parameterBlock(param));
  } else {
    problem.SetParameterBlockVariable(parameterBlock(param));
  }
}

template <typename T>
void lockParameters(ceres::Problem& problem, std::vector<T>& params, const bool lock = true) {
  for (T& param : params) {
    lockParameter(problem, param, lock);
  }
}

void reasonableResize(cv::Mat_<cv::Vec3w>& mat) {
  const double kWidth = 1200;
  const double kHeight = 800;
  double factor = std::min(kWidth / mat.cols, kHeight / mat.rows);
  if (factor < 1) {
    cv::resize(mat, mat, {}, factor, factor, cv::INTER_AREA);
  }
}

void showMatches(
    const std::vector<Camera>& cameras,
    const FeatureMap& featureMap,
    const std::vector<Overlap>& overlaps,
    const std::vector<Trace>& traces,
    const int pass) {
  // visualization for debugging
  for (const Overlap& overlap : overlaps) {
    const int idx0 = getCameraIndex(overlap.images[0]);
    const int idx1 = getCameraIndex(overlap.images[1]);
    if (cameras[idx0].overlap(cameras[idx1]) > FLAGS_debug_matches_overlap) {
      cv::Mat_<cv::Vec3w> image = renderOverlap(overlap, featureMap, traces, cameras);
      reasonableResize(image);

      if (!FLAGS_debug_dir.empty()) {
        std::string filename = FLAGS_debug_dir + "/" + "pass" + std::to_string(pass) + "_" +
            getCameraId(overlap.images[0]) + "-" + getCameraId(overlap.images[1]) + ".png";
        imwrite(filename, image);
      } else {
        cv::imshow("overlap", image);
        cv::waitKey();
      }
    }
  }
}

void showReprojections(
    const std::vector<Camera>& cameras,
    const FeatureMap& featureMap,
    const std::vector<Trace>& traces,
    const Camera::Real scale) {
  for (const auto& entry : featureMap) {
    const ImageId& image = entry.first;
    const std::vector<Feature>& features = entry.second;
    const Camera& camera = cameras[getCameraIndex(image)];
    cv::Mat_<cv::Vec3w> render = renderReprojections(image, camera, features, traces, scale);
    reasonableResize(render);
    if (!FLAGS_debug_dir.empty()) {
      std::string filename = FLAGS_debug_dir + "/" + camera.id + ".png";
      imwrite(filename, render);
    } else {
      cv::imshow("reprojections", render);
      cv::waitKey();
    }
  }
}

// returns true with a probability of numerator / denominator
bool randomSample(int numerator, int denominator) {
  static std::default_random_engine e;
  return numerator > std::uniform_int_distribution<>(0, denominator - 1)(e);
}

void solve(ceres::Problem& problem) {
  ceres::Solver::Options options;
  options.use_inner_iterations = true;
  options.max_num_iterations = 500;
  options.minimizer_progress_to_stdout = false;
  options.num_threads = ThreadPool::getThreadCountFromFlag(FLAGS_ceres_threads);
  if (options.num_threads == 0) {
    options.num_threads = 1;
  }
  options.function_tolerance = FLAGS_ceres_function_tolerance;

  ceres::Solver::Summary summary;

  LOG(INFO) << getReprojectionReport(problem);

  int previousFLAGS_v = FLAGS_v; // save FLAGS_v
  if (FLAGS_log_verbose) {
    FLAGS_v = std::max(1, FLAGS_v); // overwrite FLAGS_v
  }
  Solve(options, &problem, &summary);
  FLAGS_v = previousFLAGS_v; // restore FLAGS_v

  LOG(INFO) << summary.BriefReport();
  if (FLAGS_log_verbose) {
    LOG(INFO) << summary.FullReport();
  }

  if (summary.termination_type == ceres::NO_CONVERGENCE) {
    throw std::runtime_error("Failed to converge");
  }

  LOG(INFO) << getReprojectionReport(problem);
}

void validateMatchCount(const std::vector<Camera>& cameras, const std::vector<int>& counts) {
  const double sum = std::accumulate(counts.begin(), counts.end(), 0.0);
  const double mean = sum / counts.size();

  const double sqSum = std::inner_product(counts.begin(), counts.end(), counts.begin(), 0.0);
  const double stdev = std::sqrt(sqSum / counts.size() - mean * mean);

  std::vector<std::string> lowTraceErrors;
  for (int i = 0; i < int(counts.size()); ++i) {
    if (FLAGS_log_verbose) {
      LOG(INFO) << folly::sformat("Camera: {} Traces: {}", cameras[i].id, counts[i]);
    }
    const double z = (counts[i] - mean) / stdev;
    if (-z > FLAGS_outlier_z_threshold || counts[i] < FLAGS_min_traces) {
      lowTraceErrors.push_back(
          folly::sformat("Too few matches in camera {}: {}", cameras[i].id, counts[i]));
    }
  }

  if (!lowTraceErrors.empty()) {
    std::string errorString = boost::algorithm::join(lowTraceErrors, "\n");
    throw std::runtime_error(errorString);
  }
}

void savePointsFileJson(FeatureMap& featureMap, const std::vector<Trace>& traces) {
  folly::dynamic arrayOfTraces = folly::dynamic::array();
  for (const Trace& trace : traces) {
    if (trace.references.empty()) {
      continue; // don't output zombie traces, a different trace has the references now
    }
    folly::dynamic arrayOfFeatures = folly::dynamic::array();
    for (const auto& ref : trace.references) {
      const ImageId& image = ref.first;
      const Feature& feature = featureMap[image][ref.second];
      folly::dynamic featureSerialized = folly::dynamic::object
        ("y", feature.position.y())
        ("x", feature.position.x())
        ("image_id", ref.first);
      arrayOfFeatures.push_back(featureSerialized);
    }
    folly::dynamic traceSerialized = folly::dynamic::object
      ("features", arrayOfFeatures)
      ("number of references", trace.references.size())
      ("z", trace.position.z())
      ("y", trace.position.y())
      ("x", trace.position.x());
    arrayOfTraces.push_back(traceSerialized);
  }

  folly::dynamic points = folly::dynamic::object("points", arrayOfTraces);
  CHECK(folly::writeFile(folly::toPrettyJson(points), FLAGS_points_file_json.c_str()));
}

void savePointsFile(FeatureMap& featureMap, const std::vector<Trace>& traces) {
  std::ofstream file(FLAGS_points_file);
  for (const Trace& trace : traces) {
    if (trace.references.empty()) {
      continue; // don't output zombie traces, a different trace has the references now
    }
    file << folly::format("{} {} {} ", trace.position.x(), trace.position.y(), trace.position.z());
    file << folly::format("1 "); // delimiter
    file << folly::format("0 0 0"); // RGB value for the point
    file << "\n";
  }
}

std::vector<int> calculateCameraWeights(
    const std::vector<Camera>& cameras,
    const std::vector<Trace>& traces) {
  std::vector<int> weights(cameras.size(), 1);
  if (FLAGS_weight_by_trace_count) {
    for (ssize_t i = 0; i < ssize(cameras); ++i) {
      int cameraTraces = 0;
      for (const Trace& trace : traces) {
        for (const auto& reference : trace.references) {
          if (cameras[i].id == cameras[getCameraIndex(reference.first)].id) {
            cameraTraces++;
            continue;
          }
        }
      }
      weights[i] = cameraTraces;
    }
  }
  return weights;
}

bool positionsUnlocked(int pass) {
  return !FLAGS_lock_positions && pass != 0;
}

double refine(
    std::vector<Camera>& cameras,
    const std::vector<Camera>& groundTruth,
    FeatureMap featureMap,
    std::vector<Overlap> overlaps,
    const int pass) {
  boost::timer::cpu_timer timer;
  // remove outlier matches
  LOG(INFO) << "Removing outlier matches...";
  removeOutliersFromCameras(overlaps, featureMap, {}, cameras, FLAGS_outlier_factor);

  // assemble and remove outlier traces
  LOG(INFO) << "Assembling traces and removing outlier traces...";
  std::vector<Trace> traces = assembleTraces(featureMap, overlaps);
  triangulateTraces(traces, featureMap, cameras);
  removeOutliersFromCameras(overlaps, featureMap, traces, cameras, FLAGS_outlier_factor);

  // final triangulation
  LOG(INFO) << "Reassembling traces with outliers removed and removing invalid traces...";
  traces = assembleTraces(featureMap, overlaps);
  if (!FLAGS_keep_invalid_traces) {
    removeInvalidTraces(traces, featureMap);
  }
  triangulateTraces(traces, featureMap, cameras);

  std::vector<int> weights = calculateCameraWeights(cameras, traces);

  // visualization for debugging
  showMatches(cameras, featureMap, overlaps, traces, pass);

  // read camera parameters from cameras
  std::vector<Camera::Vector3> positions;
  std::vector<Camera::Vector3> rotations;
  std::vector<Camera::Vector2> principals;
  std::vector<Camera::Real> focals;
  std::vector<Camera::Distortion> distortions;
  for (Camera& camera : cameras) {
    positions.push_back(camera.position);
    rotations.push_back(camera.getRotation());
    principals.push_back(camera.principal);
    focals.push_back(camera.getScalarFocal());
    distortions.push_back(camera.getDistortion());
  }

  int referenceCameraIdx = -1;
  int relativeCameraIdx = -1;
  Camera::Real theta;
  Camera::Real phi;
  Camera::Real radius = 0.0; // initialized to silence compiler but this will never get used

  // If positions are unlocked, define a locked reference camera and lock the baseline between
  // the reference camera and relative camera
  if (positionsUnlocked(pass)) {
    if (FLAGS_reference_camera.empty()) {
      referenceCameraIdx = 0;
    } else {
      CHECK(cameraIdToIndex.count(FLAGS_reference_camera))
          << "bad reference_camera: " << FLAGS_reference_camera;
      referenceCameraIdx = cameraIdToIndex[FLAGS_reference_camera];
    }
    relativeCameraIdx = (referenceCameraIdx + 1) % ssize(cameras);

    Camera::Vector3 relativePosition = positions[relativeCameraIdx] - positions[referenceCameraIdx];
    cartesianToSpherical(radius, theta, phi, relativePosition);
  }

  // create the problem: add a residual for each observation
  ceres::Problem problem;
  std::vector<int> counts(cameras.size());
  for (const Trace& trace : traces) {
    if (FLAGS_cap_traces && !randomSample(FLAGS_cap_traces, traces.size())) {
      continue;
    }
    for (const auto& ref : trace.references) {
      const ImageId& image = ref.first;
      const Feature& feature = featureMap[image][ref.second];
      const int camera = getCameraIndex(image);
      ++counts[camera];
      const int group = cameraGroupToIndex[cameras[camera].group];
      if (camera == relativeCameraIdx) {
        SphericalReprojectionFunctor::addResidual(
            problem,
            theta,
            phi,
            rotations[camera],
            principals[FLAGS_shared_principal_and_focal ? group : camera],
            focals[FLAGS_shared_principal_and_focal ? group : camera],
            distortions[FLAGS_shared_distortion ? group : camera],
            traces[feature.trace].position,
            radius,
            cameras[referenceCameraIdx].position,
            cameras[camera],
            feature.position,
            FLAGS_robust,
            weights[camera]);
      } else {
        ReprojectionFunctor::addResidual(
            problem,
            positions[camera],
            rotations[camera],
            principals[FLAGS_shared_principal_and_focal ? group : camera],
            focals[FLAGS_shared_principal_and_focal ? group : camera],
            distortions[FLAGS_shared_distortion ? group : camera],
            traces[feature.trace].position,
            cameras[camera],
            feature.position,
            FLAGS_robust,
            weights[camera]);
      }
    }
  }

  validateMatchCount(cameras, counts);

  // lock focal and distortion
  if (pass == 0 || FLAGS_lock_focal) {
    if (FLAGS_shared_principal_and_focal) {
      for (const auto& mapping : cameraGroupToIndex) {
        lockParameter(problem, focals[mapping.second]);
      }
    } else {
      lockParameters(problem, focals);
    }
  }
  if (pass == 0 || FLAGS_lock_distortion) {
    if (FLAGS_shared_distortion) {
      for (const auto& mapping : cameraGroupToIndex) {
        lockParameter(problem, distortions[mapping.second]);
      }
    } else {
      lockParameters(problem, distortions);
    }
  }
  if (FLAGS_lock_principals) {
    lockParameters(problem, principals);
  }

  // lock position
  LOG(INFO) << folly::sformat("Pass: {}", pass);
  // If positions are unlocked, only lock the position and rotation of the reference camera
  if (positionsUnlocked(pass)) {
    problem.SetParameterBlockConstant(positions[referenceCameraIdx].data());
    problem.SetParameterBlockConstant(rotations[referenceCameraIdx].data());
  } else {
    lockParameters(problem, positions);
  }

  if (FLAGS_lock_rotations) {
    lockParameters(problem, rotations);
  }

  if (FLAGS_robust) {
    std::vector<calibration::ReprojectionErrorOutlier> errorsIgnored =
        getReprojectionErrorOutliers(problem);
    LOG(INFO) << folly::sformat("Number of down-weighted outliers: {}", errorsIgnored.size());
    std::sort(errorsIgnored.begin(), errorsIgnored.end(), math_util::sortdescPair<double, double>);
    LOG(INFO) << folly::sformat("Highest 3 (true/weighted): {}/{}, {}/{}, {}/{}",
        errorsIgnored[2].first, errorsIgnored[2].second,
        errorsIgnored[1].first, errorsIgnored[1].second,
        errorsIgnored[0].first, errorsIgnored[0].second);
  }
  reportReprojectionErrors(overlaps, featureMap, traces, cameras);
  solve(problem);
  if (positionsUnlocked(pass)) {
    positions[relativeCameraIdx] = sphericalToCartesian(radius, theta, phi);
    positions[relativeCameraIdx] += positions[referenceCameraIdx];
  }

  std::vector<double> norms =
      getReprojectionErrorNorms(problem, nullptr, FLAGS_weighted_statistics);
  double median = calcPercentile(norms, 0.5);
  if (pass == FLAGS_pass_count - 1 && median > FLAGS_max_error) {
    LOG(INFO) << folly::sformat("Warning: Final pass median error too high: {}", median);
  }

  // write optimized camera parameters back into cameras
  for (int i = 0; i < int(cameras.size()); ++i) {
    const int group = cameraGroupToIndex[cameras[i].group];
    cameras[i] = makeCamera(
        cameras[i],
        positions[i],
        rotations[i],
        principals[FLAGS_shared_principal_and_focal ? group : i],
        focals[FLAGS_shared_principal_and_focal ? group : i],
        distortions[FLAGS_shared_distortion ? group : i]);
  }

  reportReprojectionErrors(overlaps, featureMap, traces, cameras);

  if (FLAGS_points_file != "" && pass == FLAGS_pass_count - 1) {
    savePointsFile(featureMap, traces);
  }
  if (FLAGS_points_file_json != "" && pass == FLAGS_pass_count - 1) {
    savePointsFileJson(featureMap, traces);
  }

  // visualization for debugging
  if (FLAGS_debug_error_scale && pass == FLAGS_pass_count - 1) {
    showReprojections(cameras, featureMap, traces, FLAGS_debug_error_scale);
  }

  if (FLAGS_enable_timing) {
    LOG(INFO) << folly::sformat("Pass {} timing :{}", pass, timer.format());
  }
  return median;
}

double geometricCalibration() {
  CHECK_NE(FLAGS_rig_in, "");
  CHECK_NE(FLAGS_rig_out, "");

  if (FLAGS_debug_error_scale || FLAGS_debug_matches_overlap < 1) {
    CHECK_NE(FLAGS_color, "");
  }

  if (!FLAGS_debug_dir.empty()) {
    filesystem::create_directories(FLAGS_debug_dir);
  }

  const Camera::Rig groundTruth = Camera::loadRig(FLAGS_rig_in);
  buildCameraIndexMaps(groundTruth);
  double medianError = 0;

  if (FLAGS_seed != -1) {
    std::srand(FLAGS_seed);
  }

  for (int experiment = 0; experiment < FLAGS_experiments; ++experiment) {
    Camera::Rig cameras = groundTruth;

    Camera::perturbCameras(
        cameras,
        FLAGS_perturb_positions,
        FLAGS_perturb_rotations,
        FLAGS_perturb_principals,
        FLAGS_perturb_focals);

    FeatureMap featureMap;
    std::vector<Overlap> overlaps;

    if (!FLAGS_matches.empty()) {
      folly::dynamic parsed = parseJsonFile(FLAGS_matches);
      featureMap = loadFeatureMap(parsed);
      overlaps = loadOverlaps(parsed);
    } else {
      generateArtificalPoints(featureMap, overlaps, groundTruth);
    }

    LOG(INFO) << getCameraRmseReport(cameras, groundTruth);
    boost::timer::cpu_timer timer;

    for (int pass = 0; pass < FLAGS_pass_count; ++pass) {
      medianError = refine(cameras, groundTruth, featureMap, overlaps, pass);
      std::cout << "pass " << pass << ": " << getCameraRmseReport(cameras, groundTruth)
                << std::endl;
    }
    if (FLAGS_enable_timing) {
      LOG(INFO) << folly::sformat("Aggregate timing: {}", timer.format());
    }
    Camera::saveRig(FLAGS_rig_out, cameras);
  }

  return medianError;
}


#pragma once

#include <functional>
#include <set>
#include <unordered_set>
#include <boost/container/flat_set.hpp>

#include "nifty/tools/changable_priority_queue.hxx"
#include "nifty/graph/edge_contraction_graph.hxx"
#include "nifty/graph/agglo/cluster_policies/cluster_policies_common.hxx"

namespace nifty{
namespace graph{
namespace agglo{

template<
    class GRAPH,bool ENABLE_UCM
>
class NewPolicy{

    typedef NewPolicy<
        GRAPH, ENABLE_UCM
    > SelfType;

private:
    typedef typename GRAPH:: template EdgeMap<uint8_t> UInt8EdgeMap;
    typedef typename GRAPH:: template EdgeMap<double> FloatEdgeMap;
    typedef typename GRAPH:: template NodeMap<double> FloatNodeMap;

    typedef boost::container::flat_set<uint64_t> SetType;
    //typedef std::set<uint64_t> SetType;
    //typedef std::unordered_set<uint64_t> SetType;

public:
    // input types
    typedef GRAPH                                       GraphType;
    typedef FloatEdgeMap                                EdgePrioType;
    typedef FloatEdgeMap                                EdgeSizesType;
    typedef FloatNodeMap                                NodeSizesType;

    struct SettingsType{

        enum UpdateRule{
            SMOOTH_MAX,
            GENERALIZED_MEAN
            //HISTOGRAM_RANK
        };

        UpdateRule updateRule0{GENERALIZED_MEAN};
        UpdateRule updateRule1{GENERALIZED_MEAN};

        bool zeroInit = false;
        double p0{1.0};
        double p1{1.0};
        double gamma{0.9}
        uint64_t numberOfNodesStop{1};
        //uint64_t numberOfBins{40};
    };

    typedef EdgeContractionGraph<GraphType, SelfType>   EdgeContractionGraphType;

    friend class EdgeContractionGraph<GraphType, SelfType, ENABLE_UCM> ;
private:

    // internal types

    typedef nifty::tools::ChangeablePriorityQueue< double , std::greater<double> > QueueType;

public:

    template<class MERGE_PRIOS, class NOT_MERGE_PRIOS, class IS_LOCAL_EDGE, class EDGE_SIZES>
    NewPolicy(const GraphType &,
                              const MERGE_PRIOS & ,
                              const NOT_MERGE_PRIOS &,
                              const IS_LOCAL_EDGE &,
                              const EDGE_SIZES & ,
                              const SettingsType & settings = SettingsType());

    std::pair<uint64_t, double> edgeToContractNext() const;
    bool isDone();

    // callback called by edge contraction graph

    EdgeContractionGraphType & edgeContractionGraph();

private:
    double pqMergePrio(const uint64_t edge) const;

public:
    // callbacks called by edge contraction graph
    void contractEdge(const uint64_t edgeToContract);
    void mergeNodes(const uint64_t aliveNode, const uint64_t deadNode);
    void mergeEdges(const uint64_t aliveEdge, const uint64_t deadEdge);
    void contractEdgeDone(const uint64_t edgeToContract);

    bool isMergeAllowed(const uint64_t edge){
        if(isLocalEdge_[edge]){
            // todo this isPureLocal_ seems to be legacy
            // check if needed
           return isPureLocal_[edge] ? true : mergePrios_[edge] > notMergePrios_[edge];
        }
        else{
            return false;
        }
    }

    const EdgePrioType & mergePrios() const {
        return mergePrios_;
    }
    const EdgePrioType & notMergePrios() const {
        return notMergePrios_;
    }
    const EdgeSizesType & edgeSizes() const {
        return edgeSizes_;
    }

private:

    const double getMergePrio(const uint64_t edge)const{
        return mergePrios_[edge];
    }

    const double notMergePrio(const uint64_t edge)const{
        return notMergePrios_[edge];
    }

    // INPUT
    const GraphType &   graph_;

    EdgePrioType mergePrios_;
    EdgePrioType notMergePrios_;

    UInt8EdgeMap isLocalEdge_;
    UInt8EdgeMap isPureLocal_;
    UInt8EdgeMap isPureLifted_;
    EdgeSizesType       edgeSizes_;
    SettingsType        settings_;

    // INTERNAL
    EdgeContractionGraphType edgeContractionGraph_;
    QueueType pq_;

    uint64_t edgeToContractNext_;
    double   edgeToContractNextMergePrio_;
};

template<class GRAPH, bool ENABLE_UCM>
template<class MERGE_PRIOS, class NOT_MERGE_PRIOS, class IS_LOCAL_EDGE,class EDGE_SIZES>
inline NewPolicy<GRAPH, ENABLE_UCM>::
NewPolicy(
    const GraphType & graph,
    const MERGE_PRIOS & mergePrios,
    const NOT_MERGE_PRIOS & notMergePrios,
    const IS_LOCAL_EDGE & isLocalEdge,
    const EDGE_SIZES      & edgeSizes,
    const SettingsType & settings
)
:   graph_(graph),
    mergePrios_(graph),
    notMergePrios_(graph),
    isLocalEdge_(graph),
    isPureLocal_(graph),
    isPureLifted_(graph),
    edgeSizes_(graph),
    pq_(graph.edgeIdUpperBound()+1),
    settings_(settings),
    edgeContractionGraph_(graph, *this)
{

    graph_.forEachEdge([&](const uint64_t edge){
        isLocalEdge_[edge] = isLocalEdge[edge];

        notMergePrios_[edge] = notMergePrios[edge];
        mergePrios_[edge] = mergePrios[edge];

        if(settings_.zeroInit){
            if(isLocalEdge_[edge])
                notMergePrios_[edge] = 0.0;
            else
                mergePrios_[edge] = 0.0;
        }

        isPureLocal_[edge] = isLocalEdge[edge];
        isPureLifted_[edge] = !isLocalEdge[edge];

        edgeSizes_[edge] = edgeSizes[edge];
        pq_.push(edge, this->pqMergePrio(edge));
    });
}

template<class GRAPH, bool ENABLE_UCM>
inline std::pair<uint64_t, double>
NewPolicy<GRAPH, ENABLE_UCM>::
edgeToContractNext() const {
    return std::pair<uint64_t, double>(edgeToContractNext_,edgeToContractNextMergePrio_) ;
}

template<class GRAPH, bool ENABLE_UCM>
inline bool
NewPolicy<GRAPH, ENABLE_UCM>::
isDone()     {
    if(edgeContractionGraph_.numberOfNodes() <= settings_.numberOfNodesStop){
        //std::cout<<"done a1\n";
        return  true;
    }
    else if(pq_.empty() || pq_.topPriority() <  -0.0000001){
        //std::cout<<"done a2\n";
        return  true;
    }
    else{
        while(pq_.topPriority() > -0.0000001 ){
            const auto nextActioneEdge = pq_.top();
            if(isLocalEdge_[nextActioneEdge]){
                if(this->isMergeAllowed(nextActioneEdge)){
                    edgeToContractNext_ = nextActioneEdge;
                    edgeToContractNextMergePrio_ = pq_.topPriority();
                    //std::cout<<"not done\n";
                    return false;
                }
                else{
                    pq_.push(nextActioneEdge, -1.0);
                }
            }
            else{
                pq_.push(nextActioneEdge, -1.0);
            }
        }
        //std::cout<<"done b\n";
        return true;
    }
}

template<class GRAPH, bool ENABLE_UCM>
inline double
NewPolicy<GRAPH, ENABLE_UCM>::
pqMergePrio(
    const uint64_t edge
) const {
    return isLocalEdge_[edge] ?  mergePrios_[edge] : -1.0;
}

template<class GRAPH, bool ENABLE_UCM>
inline void
NewPolicy<GRAPH, ENABLE_UCM>::
contractEdge(
    const uint64_t edgeToContract
){
    //std::cout<<"contract edge: "<<edgeToContract<<"\n";
    pq_.deleteItem(edgeToContract);
}

template<class GRAPH, bool ENABLE_UCM>
inline typename NewPolicy<GRAPH, ENABLE_UCM>::EdgeContractionGraphType &
NewPolicy<GRAPH, ENABLE_UCM>::
edgeContractionGraph(){
    return edgeContractionGraph_;
}

template<class GRAPH, bool ENABLE_UCM>
inline void
NewPolicy<GRAPH, ENABLE_UCM>::
mergeNodes(
    const uint64_t aliveNode,
    const uint64_t deadNode
){

}

template<class GRAPH, bool ENABLE_UCM>
inline void
NewPolicy<GRAPH, ENABLE_UCM>::
mergeEdges(
    const uint64_t aliveEdge,
    const uint64_t deadEdge
){
    //std::cout<<"    merge edges: a/d "<<aliveEdge<<" "<<deadEdge<<" \n";
    NIFTY_CHECK_OP(aliveEdge,!=,deadEdge,"");
    NIFTY_CHECK(pq_.contains(aliveEdge),"");
    NIFTY_CHECK(pq_.contains(deadEdge),"");

    pq_.deleteItem(deadEdge);

    auto generalized_mean = [](
        const long double a,
        const long double d,
        const long double wa,
        const long double wd,
        const long double p
    ){
        const long double  eps = 0.000000001;
        if(std::isinf(p)){
            // max
            if(p>0){
                return std::max(a,d);
            }
            // min
            else{
                return std::min(a,d);
            }
        }
        else if(p > 1.0-eps && p< 1.0+ eps){
            return (wa*a + wd*d)/(wa+wd);
        }
        else{
            const auto wad = wa+wd;
            const auto nwa = wa/wad;
            const auto nwd = wd/wad;
            const auto sa = nwa * std::pow(a, p);
            const auto sd = nwd * std::pow(d, p);
            return std::pow(sa+sd, 1.0/p);
        }
    };

    auto smooth_max = [](
        const long double a,
        const long double d,
        const long double wa,
        const long double wd,
        const long double p
    ){
        const long double  eps = 0.000000001;
        if(std::isinf(p)){
            // max
            if(p>0){
                return std::max(a,d);
            }
            // min
            else{
                return std::min(a,d);
            }
        }
        else if(p > 0.0-eps && p< 0.0+ eps){
            return (wa*a + wd*d)/(wa+wd);
        }
        else{

            const auto eaw = wa * std::exp(a*p);
            const auto edw = wd * std::exp(d*p);
            return (a*eaw + d*edw)/(eaw + edw);
        }
    };

    //  sizes
    const auto sa = edgeSizes_[aliveEdge];
    const auto sd = edgeSizes_[deadEdge];
    const auto zi = settings_.zeroInit ;

    // update merge prio
    if(zi && isPureLifted_[aliveEdge] && !isPureLifted_[deadEdge]){
        mergePrios_[aliveEdge] = mergePrios_[deadEdge];
    }
    else if(zi && !isPureLifted_[aliveEdge] && isPureLifted_[deadEdge]){
        mergePrios_[deadEdge] = mergePrios_[aliveEdge];
    }
    else{
        if(settings_.updateRule0 == SettingsType::GENERALIZED_MEAN){
            mergePrios_[aliveEdge]    = generalized_mean(mergePrios_[aliveEdge],     mergePrios_[deadEdge],    sa, sd, settings_.p0);
        }
        else if(settings_.updateRule0 == SettingsType::SMOOTH_MAX){
            mergePrios_[aliveEdge]    = smooth_max(mergePrios_[aliveEdge],     mergePrios_[deadEdge],    sa, sd, settings_.p0);
        }
        else{
            NIFTY_CHECK(false,"not yet implemented");
        }
    }

    // update notMergePrio
    if(zi && isPureLocal_[aliveEdge] && !isPureLocal_[deadEdge]){
        notMergePrios_[aliveEdge] = notMergePrios_[deadEdge];
    }
    else if(zi && !isPureLocal_[aliveEdge] && isPureLocal_[deadEdge]){
        notMergePrios_[aliveEdge] = notMergePrios_[deadEdge];
    }
    else{
        if(settings_.updateRule0 == SettingsType::GENERALIZED_MEAN){
            notMergePrios_[aliveEdge] = generalized_mean(notMergePrios_[aliveEdge] , notMergePrios_[deadEdge], sa, sd, settings_.p1);
        }
        else if(settings_.updateRule0 == SettingsType::SMOOTH_MAX){
            notMergePrios_[aliveEdge] = smooth_max(notMergePrios_[aliveEdge] , notMergePrios_[deadEdge], sa, sd, settings_.p1);
        }
        else{
            NIFTY_CHECK(false,"not yet implemented");
        }
    }

    edgeSizes_[aliveEdge] = sa + sd;

    const auto deadIsLocalEdge = isLocalEdge_[deadEdge];
    auto & aliveIsLocalEdge = isLocalEdge_[aliveEdge];
    aliveIsLocalEdge = deadIsLocalEdge || aliveIsLocalEdge;

    isPureLocal_[aliveEdge] = isPureLocal_[aliveEdge] && isPureLocal_[deadEdge];
    isPureLifted_[aliveEdge] = isPureLifted_[aliveEdge] && isPureLifted_[deadEdge];

    pq_.push(aliveEdge, this->pqMergePrio(aliveEdge));

}

template<class GRAPH, bool ENABLE_UCM>
inline void
NewPolicy<GRAPH, ENABLE_UCM>::
contractEdgeDone(
    const uint64_t edgeToContract
){
    //std::cout<<"contract edge done: "<<edgeToContract<<"\n\n";
}

} // namespace agglo
} // namespace nifty::graph
} // namespace nifty


//
// Copyright (c) 2010 Dariusz Gadomski <dgadomski@gmail.com>
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
//  * Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
//  * Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation and/or
// other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS''
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
// ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//

#include <predictionserver.h>

#include <comm/protocol.h>
#include <arima/arima.h>
#include <chaos/chaos.h>
#include <grey/grey.h>
#include <neural/neuralnet.h>
#include <util.h>

#include <boost/asio.hpp>
#include <boost/bind.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/thread/mutex.hpp>

// Must come before boost/serialization headers.
#include <comm/connection.h>
#include <comm/protocol.h>

namespace prediction
{

namespace server
{

using namespace models::dataprovider;
using namespace comm;
using namespace debug;

PredictionServer::PredictionServer(boost::asio::io_service & io_service,
        const ParsedOptions& opts) :
    _acceptor(io_service, boost::asio::ip::tcp::endpoint(
            boost::asio::ip::tcp::v4(), opts.ListenPort)), _algorithm(
            opts.Algorithm), _dataProvider(new DataProvider(opts.InputFile)),
            _opts(opts)
//  _connection(io_service), _algorithm(opts.Algorithm), _stopFlag(false),
//          _predictionStarted(false), _opts(opts)
{
    connection_ptr new_conn(new connection(_acceptor.get_io_service()));
    _acceptor.async_accept(new_conn->socket(), boost::bind(
            &PredictionServer::handle_accept, this,
            boost::asio::placeholders::error, new_conn));

    createPredictionModel(_algorithm);

}

void PredictionServer::handle_read(const boost::system::error_code& e,
        connection_ptr conn)
{
    if (!e)
    {
        dbg(debug::Informational) << "Handle read: " << std::endl;
        dbg(debug::Informational) << _inBuffer << std::endl;

        unsigned dataStart = _inBuffer.DataOffset;
        unsigned dataLength = _inBuffer.DataLength;
        unsigned horizon = _inBuffer.Horizon;

        std::vector<double> inputBuffer(_dataProvider->getDataVector(dataStart,
                dataLength));

        _predictionModel->provideInput(inputBuffer, horizon);

        double prediction = _predictionModel->getPrediction(horizon);

//      if( _algorithm == std::string("neural") )
//      {
//          printSeq("!!! Input: ", inputBuffer, debug::High);
//          dbg(debug::High) << "!!! Prediction: " << prediction << std::endl;
//      }

        _outBuffer = _inBuffer;
        _outBuffer.Result = prediction;
        _outBuffer.Algorithm = _algorithm;
        _outBuffer.DataOffset = _inBuffer.DataOffset;
        _outBuffer.DataLength = _inBuffer.DataLength;

        conn->async_write(_outBuffer, boost::bind(
                &PredictionServer::handle_write, this,
                boost::asio::placeholders::error, conn));
    }
    else
    {
        dbg(debug::High) << e.message() << std::endl;
    }
}

void PredictionServer::handle_write(const boost::system::error_code& e,
        connection_ptr conn)
{
    if (!e)
    {
        dbg(debug::Informational) << "Handle write: " << std::endl;
        dbg(debug::Informational) << _outBuffer << std::endl;

        conn->async_read(_inBuffer, boost::bind(&PredictionServer::handle_read,
                        this, boost::asio::placeholders::error, conn));
    }
    else
    {
        dbg(debug::High) << e.message() << std::endl;
    }
}

void PredictionServer::handle_accept(const boost::system::error_code& e,
        connection_ptr conn)
{
    if (!e)
    {
        dbg() << "Accepted connection!" << std::endl;

        conn->async_read(_inBuffer, boost::bind(&PredictionServer::handle_read,
                this, boost::asio::placeholders::error, conn));

        connection_ptr new_conn(new connection(_acceptor.get_io_service()));
        _acceptor.async_accept(new_conn->socket(), boost::bind(
                &PredictionServer::handle_accept, this,
                boost::asio::placeholders::error, new_conn));
    }
    else
    {
        // An error occurred. Log it and return. Since we are not starting a new
        // accept operation the io_service will run out of work to do and the
        // server will exit.
        dbg(debug::High) << e.message() << std::endl;
    }
}

void PredictionServer::createPredictionModel(const std::string& algorithm)
{
    models::AbstractModel *model = 0;

    if (algorithm == "arima")
    {
        models::arima::Arima *arima = new models::arima::Arima();
        std::vector<int> order;
        order.push_back(1);
        order.push_back(2);
        order.push_back(1);
        arima->setOrder(order);
        model = arima;
    }
    else if (algorithm == "grey")
    {
        model = new models::grey::Grey();
    }
    else if (algorithm == "chaos")
    {
        model = new models::chaos::Chaos(3, 1);
    }
    else if (algorithm == "neural")
    {
        models::neural::NeuralNet *net = models::neural::NeuralNet::load("learning.net");
        net->setScale(1.0 / _dataProvider->getMaxValue());
        model = net;
    }

    if (model)
    {
        _predictionModel = boost::shared_ptr<models::AbstractModel>(model);
    }
}

double PredictionServer::getPrediction(size_t offset, size_t length,
        size_t horizon, size_t progress)
{
    std::vector<double> inputVec = _dataProvider->getDataVector(offset, length);

    if (progress == 0)
    {
        _predictionModel->provideInput(inputVec, horizon);
    }

    // get prediction - progress=0 -> prediction for t+1 etc.
    return _predictionModel->getPrediction(progress + 1);
}

}
}


#include <fstream>  // NOLINT(readability/streams)
#include <iostream>  // NOLINT(readability/streams)
#include <string>
#include <utility>
#include <vector>

#include <boost/iostreams/device/array.hpp>
#include <boost/iostreams/stream.hpp>
#include <boost/iostreams/copy.hpp>
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/algorithm/string.hpp>

#include "caffe/data_transformer.hpp"
#include "caffe/layers/base_data_layer.hpp"
#include "caffe/layers/molgrid_data_layer.hpp"
#include "caffe/util/benchmark.hpp"
#include "caffe/util/io.hpp"
#include "caffe/util/math_functions.hpp"

#include <openbabel/mol.h>
#include <openbabel/obconversion.h>
#include <openbabel/obiter.h>
#include <boost/timer/timer.hpp>

//allocate and initialize atom type data
namespace smina_atom_type
{
  info data[NumTypes] = {{},};

  struct atom_data_initializer {
    atom_data_initializer() {
      for(size_t i = 0u; i < smina_atom_type::NumTypes; ++i)
        smina_atom_type::data[i] = smina_atom_type::default_data[i];
    }
  };

  atom_data_initializer initialize_defaults;
}

namespace caffe {

template <typename Dtype>
MolGridDataLayer<Dtype>::~MolGridDataLayer<Dtype>() {
  //this->StopInternalThread();

  if(gpu_gridatoms) {
    cudaFree(gpu_gridatoms);
    gpu_gridatoms = NULL;
  }
  if(gpu_gridwhich) {
    cudaFree(gpu_gridwhich);
    gpu_gridwhich = NULL;
  }

  if(data) delete data;
  if(data2) delete data2;
}

template <typename Dtype>
MolGridDataLayer<Dtype>::example::example(MolGridDataLayer<Dtype>::string_cache& cache, string line, bool hasaffinity, bool hasrmsd)
  : label(0), affinity(0.0), rmsd(0.0)
{
  stringstream stream(line);
  string tmp;
  //first the label
  stream >> label;
  if(hasaffinity)
   stream >> affinity;
  if(hasrmsd)
   stream >> rmsd;
  //receptor
  stream >> tmp;
  CHECK(tmp.length() > 0) << "Empty receptor, missing affinity/rmsd? Line:\n" << line;
  receptor = cache.get(tmp);
  //ligand
  tmp.clear();
  stream >> tmp;
  CHECK(tmp.length() > 0) << "Empty ligand, missing affinity/rmsd? Line:\n" << line;

  ligand = cache.get(tmp);
}

//for in-memory inputs, set the desired label (for gradient computation)
//note that zero affinity/rmsd means to ignore these
template<typename Dtype>
void MolGridDataLayer<Dtype>::setLabels(Dtype pose, Dtype affinity, Dtype rmsd)
{
  labels.clear();
  affinities.clear();
  rmsds.clear();
  labels.push_back(pose);
  affinities.push_back(affinity);
  rmsds.push_back(rmsd);
}

//the following really shouldn't be recalculated each evaluation (not including gradients)
template<typename Dtype>
void MolGridDataLayer<Dtype>::getReceptorAtoms(int batch_idx, vector<float4>& atoms)
{
  atoms.resize(0);
  mol_info& mol = batch_transform[batch_idx].mol;
  for (unsigned i = 0, n = mol.atoms.size(); i < n; ++i)
    if (mol.whichGrid[i] < numReceptorTypes)
      atoms.push_back(mol.atoms[i]);
}

template<typename Dtype>
void MolGridDataLayer<Dtype>::getLigandAtoms(int batch_idx, vector<float4>& atoms)
{
  atoms.resize(0);
  mol_info& mol = batch_transform[batch_idx].mol;
  for (unsigned i = 0, n = mol.atoms.size(); i < n; ++i)
    if (mol.whichGrid[i] >= numReceptorTypes)
      atoms.push_back(mol.atoms[i]);
}

template<typename Dtype>
void MolGridDataLayer<Dtype>::getReceptorChannels(int batch_idx, vector<short>& whichGrid)
{
  whichGrid.resize(0);
  mol_info& mol = batch_transform[batch_idx].mol;
  for (unsigned i = 0, n = mol.atoms.size(); i < n; ++i)
    if (mol.whichGrid[i] < numReceptorTypes)
      whichGrid.push_back(mol.whichGrid[i]);
}

template<typename Dtype>
void MolGridDataLayer<Dtype>::getLigandChannels(int batch_idx, vector<short>& whichGrid)
{
  whichGrid.resize(0);
  mol_info& mol = batch_transform[batch_idx].mol;
  for (unsigned i = 0, n = mol.atoms.size(); i < n; ++i)
    if (mol.whichGrid[i] >= numReceptorTypes)
      whichGrid.push_back(mol.whichGrid[i]);
}

template<typename Dtype>
void MolGridDataLayer<Dtype>::getReceptorGradient(int batch_idx, vector<float3>& gradient)
{
  gradient.resize(0);
  CHECK(compute_atom_gradients) << "Gradients requested but not computed";
  mol_info& mol = batch_transform[batch_idx].mol;
  for (unsigned i = 0, n = mol.atoms.size(); i < n; ++i)
    if (mol.whichGrid[i] < numReceptorTypes)
    {
      gradient.push_back(mol.gradient[i]);
    }
}

/*
 * Compute the transformation gradient of a rigid receptor around the center.
 * The first three numbers are the translation.  The next are the torque.
 */
template<typename Dtype>
void MolGridDataLayer<Dtype>::getReceptorTransformationGradient(int batch_idx, vec& force, vec& torque)
{
  force = vec(0,0,0);
  torque = vec(0,0,0);

  CHECK(compute_atom_gradients) << "Gradients requested but not computed";
  mol_info& mol = batch_transform[batch_idx].mol;

  CHECK(mol.center == mem_lig.center) << "Centers not equal; receptor transformation gradient only supported in-mem";

  for (unsigned i = 0, n = mol.atoms.size(); i < n; ++i)
  {
    if (mol.whichGrid[i] < numReceptorTypes)
    {
      float3 g = mol.gradient[i];
      float4 a = mol.atoms[i];
      vec v(g.x,g.y,g.z);
      vec pos(a.x,a.y,a.z);

      force += v;
      torque += cross_product(pos - mol.center, v);
    }
  }
}

template<typename Dtype>
void MolGridDataLayer<Dtype>::getMappedReceptorGradient(int batch_idx, unordered_map<string, float3>& gradient)
{
  CHECK(compute_atom_gradients) << "Gradients requested but not computed";
  mol_info& mol = batch_transform[batch_idx].mol;
  for (unsigned i = 0, n = mol.atoms.size(); i < n; ++i)
  {
    if (mol.whichGrid[i] < numReceptorTypes)
    {
        string xyz = xyz_to_string(mol.atoms[i].x, mol.atoms[i].y, mol.atoms[i].z);
        gradient[xyz] = mol.gradient[i];
    }
  }
}

template<typename Dtype>
void MolGridDataLayer<Dtype>::getLigandGradient(int batch_idx, vector<float3>& gradient)
{
  CHECK(compute_atom_gradients) << "Gradients requested but not computed";
  gradient.resize(0);
  mol_info& mol = batch_transform[batch_idx].mol;
  for (unsigned i = 0, n = mol.atoms.size(); i < n; ++i)
    if (mol.whichGrid[i] >= numReceptorTypes)
    {
      gradient.push_back(mol.gradient[i]);
    }
}

template<typename Dtype>
void MolGridDataLayer<Dtype>::getMappedLigandGradient(int batch_idx, unordered_map<string, float3>& gradient)
{
  CHECK(compute_atom_gradients) << "Gradients requested but not computed";
  mol_info& mol = batch_transform[batch_idx].mol;
  for (unsigned i = 0, n = mol.atoms.size(); i < n; ++i)
  {
    if (mol.whichGrid[i] >= numReceptorTypes)
    {
        string xyz = xyz_to_string(mol.atoms[i].x, mol.atoms[i].y, mol.atoms[i].z);
        gradient[xyz] = mol.gradient[i];
    }
  }
}

//modify examples to remove any without both actives an inactives
//factored this into its own function due to the need to fully specialize setup below
template<typename Dtype>
void MolGridDataLayer<Dtype>::remove_missing_and_setup(vector<typename MolGridDataLayer<Dtype>::balanced_example_provider>& examples)
{
  vector<balanced_example_provider> tmp;
  for(unsigned i = 0, n = examples.size(); i < n; i++)
  {
    if(examples[i].num_actives() > 0 && examples[i].num_decoys() > 0) {
      //eliminate empty buckets
      tmp.push_back(examples[i]);
      tmp.back().setup();
    }
    else if(examples[i].num_actives() > 0)
    {
      example tmp;
      examples[i].next_active(tmp);
      LOG(INFO) << "Dropping receptor " << tmp.receptor << " with no decoys.";
    }
    else if(examples[i].num_decoys() > 0)
    {
      example tmp;
      examples[i].next_decoy(tmp);
      LOG(INFO) << "Dropping receptor " << tmp.receptor << " with no decoys.";
    }
  }

  swap(examples,tmp);
}

//specialized version for balanced data that remove receptors without any actives or decoys
//annoyingly, have to specialize Dtype
template<>
template<>
void MolGridDataLayer<float>::receptor_stratified_example_provider<typename MolGridDataLayer<float>::balanced_example_provider, 2>::setup()
{
  currenti = 0; currentk = 0;
  remove_missing_and_setup(examples);
  //also shuffle receptors
  if(randomize) shuffle(examples.begin(), examples.end(), caffe::caffe_rng());
}

template<>
template<>
void MolGridDataLayer<double>::receptor_stratified_example_provider<typename MolGridDataLayer<double>::balanced_example_provider, 2>::setup()
{
  currenti = 0; currentk = 0;
  remove_missing_and_setup(examples);
  //also shuffle receptors
  if(randomize) shuffle(examples.begin(), examples.end(), caffe::caffe_rng());
}

//ensure gpu memory is of sufficient size
template <typename Dtype>
void MolGridDataLayer<Dtype>::allocateGPUMem(unsigned sz)
{
  if(sz > gpu_alloc_size) {
    //deallocate
    if(gpu_gridatoms) {
      cudaFree(gpu_gridatoms);
    }
    if(gpu_gridwhich) {
      cudaFree(gpu_gridwhich);
    }
    //allocate larger
    CUDA_CHECK(cudaMalloc(&gpu_gridatoms, sz*sizeof(float4)));
    CUDA_CHECK(cudaMalloc(&gpu_gridwhich, sz*sizeof(short)));
  }
}

//allocate and return an example provider to the specifications of the parm object
template <typename Dtype>
typename MolGridDataLayer<Dtype>::example_provider* MolGridDataLayer<Dtype>::create_example_data(const MolGridDataParameter& parm)
{
  bool balanced  = parm.balanced();
  bool strat_receptor  = parm.stratify_receptor();
  bool strat_aff = parm.stratify_affinity_max() != parm.stratify_affinity_min();

  //strat_aff > strat_receptor > balanced
  if(strat_aff)
  {
    if(strat_receptor)
    {
      if(balanced) // sample 2 from each receptor
      {
        return new affinity_stratified_example_provider<receptor_stratified_example_provider<balanced_example_provider, 2> >(parm);
      }
      else //sample 1 from each receptor
      {
        return new affinity_stratified_example_provider<receptor_stratified_example_provider<uniform_example_provider, 1> >(parm);
      }
    }
    else
    {
      if(balanced)
      {
        return new affinity_stratified_example_provider<balanced_example_provider>(parm);
      }
      else //sample 1 from each receptor
      {
        return new affinity_stratified_example_provider<uniform_example_provider>(parm);
      }
    }
  }
  else if(strat_receptor)
  {
    if(balanced) // sample 2 from each receptor
    {
      return new receptor_stratified_example_provider<balanced_example_provider, 2>(parm);
    }
    else //sample 1 from each receptor
    {
      return new receptor_stratified_example_provider<uniform_example_provider, 1>(parm);
    }
  }
  else if(balanced)
  {
    return new balanced_example_provider(parm);
  }
  else
  {
    return new uniform_example_provider(parm);
  }
}

//make sure can append to root_folder path
static string sanitize_path(const string& p)
{
  //make sure root folder(s) have trailing slash
  if (p.length() > 0 && p[p.length()-1] != '/')
    return p + "/";
  else
    return p;
}

//fill in training examples
template <typename Dtype>
void MolGridDataLayer<Dtype>::populate_data(const string& root_folder, const string& source,
    MolGridDataLayer<Dtype>::example_provider* data, bool hasaffinity, bool hasrmsd)
{
  LOG(INFO) << "Opening file " << source;
  std::ifstream infile(source.c_str());
  CHECK((bool)infile) << "Could not open " << source;
  string line;
  while (getline(infile, line))
  {
    example ex(scache, line, hasaffinity, hasrmsd);
    data->add(ex);
  }
  CHECK_GT(data->size(),0) << "No examples provided in source: " << source;

  data->setup(); //done adding

}

//read in structure input and atom type maps
template <typename Dtype>
void MolGridDataLayer<Dtype>::DataLayerSetUp(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {

  const MolGridDataParameter& param = this->layer_param_.molgrid_data_param();
  root_folder = param.root_folder();
  num_rotations = param.rotate();
  inmem = param.inmemory();
  dimension = param.dimension();
  resolution = param.resolution();
  binary = param.binary_occupancy();
  bool spherize = param.spherical_mask();
  randtranslate = param.random_translate();
  randrotate = param.random_rotation();
  ligpeturb = param.peturb_ligand();
  ligpeturb_translate = param.peturb_ligand_translate();
  ligpeturb_rotate = param.peturb_ligand_rotate();
  jitter = param.jitter();
  ignore_ligand = param.ignore_ligand();
  radiusmultiple = param.radius_multiple();
  fixedradius = param.fixed_radius();
  bool hasaffinity = param.has_affinity();
  bool hasrmsd = param.has_rmsd();
  data_ratio = param.source_ratio();
  root_folder2 = param.root_folder2();

  if(binary) radiusmultiple = 1.0;

  gmaker.initialize(resolution, dimension, radiusmultiple, binary, spherize);

  dim = round(dimension/resolution)+1; //number of grid points on a side
  numgridpoints = dim*dim*dim;
  if(numgridpoints % 512 != 0)
    LOG(INFO) << "Total number of grid points (" << numgridpoints << ") is not evenly divisible by 512.";

  //shape must come from parameters
  int batch_size = param.batch_size();

  if(!inmem)
  {

    const string& source = param.source();
    const string& source2 = param.source2();
    root_folder = sanitize_path(param.root_folder());
    root_folder2 = param.root_folder2();
    if(root_folder2.length() > 0) root_folder2 = sanitize_path(root_folder2);
    else root_folder2 = root_folder; //fall back on first

    CHECK_GT(source.length(), 0) << "No data source file provided";

    // Read source file(s) with labels and structures,
    // each line is label [affinity] [rmsd] receptor_file ligand_file
    data = create_example_data(param);
    populate_data(root_folder, source, data, hasaffinity, hasrmsd);

    if(source2.length() > 0)
    {
      CHECK_GE(data_ratio, 0) << "Must provide non-negative ratio for two data sources";
      data2 = create_example_data(param);
      populate_data(root_folder2, source2, data2, hasaffinity, hasrmsd);
    }

    LOG(INFO) << "Total examples: " << data->size() + (data2 ? data2->size() : 0);

    // Check if we would need to randomly skip a few data points
    if (param.rand_skip())
    {
      unsigned int skip = caffe_rng_rand() %  param.rand_skip();

      LOG(INFO) << "Skipping first " << skip << " data points from each source.";

      example dummy;
      for(unsigned i = 0; i < skip; i++) {
        data->next(dummy);
      }
      if(data2)
      {
        for(unsigned i = 0; i < skip; i++) {
          data2->next(dummy);
        }
      }
    }
  }
  else //in memory always batch size of 1
  {
    batch_size = 1;
  }

  CHECK_GT(batch_size, 0) << "Positive batch size required";
  //keep track of atoms and transformations for each example in batch
  batch_transform.resize(batch_size);

  //initialize atom type maps
  string recmap = param.recmap();
  string ligmap = param.ligmap();

  if (recmap.size() == 0)
    numReceptorTypes = GridMaker::createDefaultRecMap(rmap);
  else
    numReceptorTypes = GridMaker::createAtomTypeMap(recmap, rmap);

  if (ligmap.size() == 0)
    numLigandTypes = GridMaker::createDefaultLigMap(lmap);
  else
    numLigandTypes = GridMaker::createAtomTypeMap(ligmap, lmap);

  //setup shape of layer
  top_shape.clear();
  top_shape.push_back(batch_size);
  top_shape.push_back(numReceptorTypes+numLigandTypes);
  top_shape.push_back(dim);
  top_shape.push_back(dim);
  top_shape.push_back(dim);

  example_size = (numReceptorTypes+numLigandTypes)*numgridpoints;

  // Reshape prefetch_data and top[0] according to the batch_size.
  top[0]->Reshape(top_shape);

  // Reshape label, affinity, rmsds
  vector<int> label_shape(1, batch_size); // [batch_size]

  top[1]->Reshape(label_shape);

  if (hasaffinity)
  {
    top[2]->Reshape(label_shape);
    if (hasrmsd)
    {
      top[3]->Reshape(label_shape);
    }
  }
  else if(hasrmsd)
  {
    top[2]->Reshape(label_shape);
  }

  if(ligpeturb) {
    vector<int> peturbshape(2);
    peturbshape[0] = batch_size;
    peturbshape[1] = 6; //trans+orient
    top.back()->Reshape(peturbshape);
  }
}

//return quaternion representing one of 24 distinct axial rotations
template <typename Dtype>
typename MolGridDataLayer<Dtype>::quaternion MolGridDataLayer<Dtype>::axial_quaternion()
{
  unsigned rot = current_rotation;
  qt ret;
  //first rotate to a face
  switch(rot%6) {
    case 0:
      ret = quaternion(1,0,0,0); //identity
      break;
    case 1: //rotate z 90
      ret = quaternion(sqrt(0.5),0,0,sqrt(0.5));
      break;
    case 2: //rotate z 180
      ret = quaternion(0,0,0,1);
      break;
    case 3: //rotate z 270 (-90)
      ret = quaternion(sqrt(0.5),0,0,-sqrt(0.5));
      break;
    case 4: //rotate y 90
      ret = quaternion(sqrt(0.5),0,sqrt(0.5),0);
      break;
    case 5: //rotate y -90
      ret = quaternion(sqrt(0.5),0,-sqrt(0.5),0);
      break;
  }

  //now four rotations around x axis
  rot /= 6;
  switch(rot%4) {
    case 0:
      break; //identity
    case 1: //90 degrees
      ret *= quaternion(sqrt(0.5),sqrt(0.5),0,0);
      break;
    case 2: //180
      ret *= quaternion(0,1,0,0);
      break;
    case 3: //270
      ret *= quaternion(sqrt(0.5),-sqrt(0.5),0,0);
      break;
  }
  return ret;
}

template <typename Dtype>
void MolGridDataLayer<Dtype>::set_mol_info(const string& file, const vector<int>& atommap,
    unsigned mapoffset, mol_info& minfo)
{
  //read mol info from file
  //OpenBabel is SLOW, especially for the receptor, so we cache the result
  //if this gets too annoying, can add support for spawning a thread for openbabel
  //but since this gets amortized across many hits to the same example, not a high priority
  using namespace OpenBabel;

  vec center(0,0,0);
  minfo.atoms.clear();
  minfo.whichGrid.clear();
  minfo.gradient.clear();

  //also, implemented a custom gninatypes files to precalc this info
  if(boost::algorithm::ends_with(file,".gninatypes"))
  {
    struct info {
      float x,y,z;
      int type;
    } atom;

    ifstream in(file.c_str());
    CHECK(in) << "Could not read " << file;

    int cnt = 0;
    while(in.read((char*)&atom, sizeof(atom)))
    {
      smt t = (smt)atom.type;
      int index = atommap[t];

      if(index >= 0)
      {
        cnt++;
        float4 ainfo;
        ainfo.x = atom.x;
        ainfo.y = atom.y;
        ainfo.z = atom.z;
        if(fixedradius <= 0)
        	ainfo.w = xs_radius(t);
        else
        	ainfo.w = fixedradius;
        float3 gradient(0,0,0);

        minfo.atoms.push_back(ainfo);
        minfo.whichGrid.push_back(index+mapoffset);
        minfo.gradient.push_back(gradient);
        center += vec(atom.x,atom.y,atom.z);
      }
      else if(t > 1) //silence on hydrogens
      {
       std::cerr << "WARNING: Unknown atom type " << t << " in " << file << ".  This atom will be discarded\n";
      }
    }

    if(cnt == 0) {
      std::cerr << "WARNING: No atoms in " << file <<"\n";
    }
    else {
      center /= cnt;
    }
  }
  else if(!boost::algorithm::ends_with(file,"none")) //reserved word
  {
    //read mol from file and set mol info (atom coords and grid positions)
    //types are mapped using atommap values plus offset
    OpenBabel::OBConversion conv;
    OBMol mol;
    CHECK(conv.ReadFile(&mol, file)) << "Could not read " << file;

    if(this->layer_param_.molgrid_data_param().addh()) {
      mol.AddHydrogens();
    }

    minfo.atoms.reserve(mol.NumHvyAtoms());
    minfo.whichGrid.reserve(mol.NumHvyAtoms());
    minfo.gradient.reserve(mol.NumHvyAtoms());

    int cnt = 0;
    FOR_ATOMS_OF_MOL(a, mol)
    {
      smt t = obatom_to_smina_type(*a);
      int index = atommap[t];

      if(index >= 0)
      {
        cnt++;
        float4 ainfo;
        ainfo.x = a->x();
        ainfo.y = a->y();
        ainfo.z  = a->z();
        if(fixedradius <= 0)
        	ainfo.w = xs_radius(t);
        else
        	ainfo.w = fixedradius;
        float3 gradient(0,0,0);

        minfo.atoms.push_back(ainfo);
        minfo.whichGrid.push_back(index+mapoffset);
        minfo.gradient.push_back(gradient);
        center += vec(a->x(),a->y(),a->z());
      }
      else
      {
        std::cerr << "WARNING: Unknown atom type in " << file << ".  This atom will be discarded\n";
      }
    }
    center /= cnt;
  }

  minfo.center = center;

}

template <typename Dtype>
void MolGridDataLayer<Dtype>::set_grid_ex(Dtype *data, const MolGridDataLayer<Dtype>::example& ex,
    const string& root_folder, MolGridDataLayer<Dtype>::mol_transform& transform, output_transform& peturb, bool gpu)
{
  //set grid values for example
  //cache atom info
  bool docache = this->layer_param_.molgrid_data_param().cache_structs();

  if(docache)
  {
    if(molcache.count(ex.receptor) == 0)
    {
      set_mol_info(root_folder+ex.receptor, rmap, 0, molcache[ex.receptor]);
    }
    if(molcache.count(ex.ligand) == 0)
    {
      set_mol_info(root_folder+ex.ligand, lmap, numReceptorTypes, molcache[ex.ligand]);
    }

    set_grid_minfo(data, molcache[ex.receptor], molcache[ex.ligand], transform, peturb, gpu);
  }
  else
  {
    mol_info rec;
    mol_info lig;
    set_mol_info(root_folder+ex.receptor, rmap, 0, rec);
    set_mol_info(root_folder+ex.ligand, lmap, numReceptorTypes, lig);
    set_grid_minfo(data, rec, lig, transform, peturb, gpu);
  }
}

template <typename Dtype>
void MolGridDataLayer<Dtype>::set_grid_minfo(Dtype *data, const MolGridDataLayer<Dtype>::mol_info& recatoms,
  const MolGridDataLayer<Dtype>::mol_info& ligatoms, MolGridDataLayer<Dtype>::mol_transform& transform,
  output_transform& peturb, bool gpu)
{
  //set grid values from mol info
  //first clear transform from the previous batch
  rng_t* rng = caffe_rng();
  transform = mol_transform();
  mol_transform ligtrans;

  //include receptor and ligand atoms
  transform.mol.append(recatoms);
  //set center to ligand center
  transform.mol.center = ligatoms.center;

  if(ligpeturb) {
    if(ligpeturb_rotate)
    {
      ligtrans.set_random_quaternion(rng);
    }
    else
    {
      ligtrans.Q = quaternion(1,0,0,0); //identity
    }
    ligtrans.add_random_displacement(rng, ligpeturb_translate);
    transform.mol.transform_and_append(ligatoms, ligtrans);

    //store the inverse transformation
    peturb.x = -ligtrans.center[0];
    peturb.y = -ligtrans.center[1];
    peturb.z = -ligtrans.center[2];

    qt qinv = conj(ligtrans.Q)/norm(ligtrans.Q); //not Cayley, not euclidean norm - already squared
    peturb.set_from_quaternion(qinv);

    //set the center to the translated value
    transform.mol.center = ligatoms.center + ligtrans.center;
  } else if(ignore_ligand) {
    //do nothing - ligand is only used to set center
  } else {
    transform.mol.append(ligatoms);
  }

  //figure out transformation
  transform.Q = quaternion(1,0,0,0);

  if(current_rotation == 0 && !randrotate)
    transform.Q = quaternion(1,0,0,0); //check real part to avoid mult

  if (randrotate)
  {
    transform.set_random_quaternion(rng);
  }

  transform.center[0] = transform.mol.center[0];
  transform.center[1] = transform.mol.center[1];
  transform.center[2] = transform.mol.center[2];
  if (randtranslate)
  {
    double radius = ligatoms.radius();
    //don't let ligand atoms translate out of sphere inscribed in box
    if(ignore_ligand) radius = 0;
    double maxtrans = max(dimension/2.0 - radius,0.0);
    transform.add_random_displacement(rng, min(randtranslate,maxtrans));
  }

  if(current_rotation > 0) {
    transform.Q *= axial_quaternion();
  }

  //TODO move this into gridmaker.setAtoms, have it just take the mol_transform as input - separate receptor transform as well
  gmaker.setCenter(transform.center[0], transform.center[1], transform.center[2]);

  if(transform.mol.atoms.size() == 0) {
     std::cerr << "ERROR: No atoms in molecule.  I can't deal with this.\n";
     exit(-1); //presumably you never actually want this and it results in a cuda error
  }
  if(jitter > 0) {
    //add small random displacement (in-place) to atoms
    for(unsigned i = 0, n = transform.mol.atoms.size(); i < n; i++) {
      float4& atom = transform.mol.atoms[i];
      float xdiff = jitter*(unit_sample(rng)*2.0-1.0);
      atom.x += xdiff;
      float ydiff = jitter*(unit_sample(rng)*2.0-1.0);
      atom.y += ydiff;
      float zdiff = jitter*(unit_sample(rng)*2.0-1.0);
      atom.z += zdiff;
    }
  }

  //compute grid from atom info arrays
  if(gpu)
  {
    unsigned natoms = transform.mol.atoms.size();
    allocateGPUMem(natoms);
    CUDA_CHECK(cudaMemcpy(gpu_gridatoms, &transform.mol.atoms[0], natoms*sizeof(float4), cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(gpu_gridwhich, &transform.mol.whichGrid[0], natoms*sizeof(short), cudaMemcpyHostToDevice));

    gmaker.setAtomsGPU<Dtype>(natoms, gpu_gridatoms, gpu_gridwhich, transform.Q, numReceptorTypes+numLigandTypes, data);
  }
  else
  {
    Grids grids(data, boost::extents[numReceptorTypes+numLigandTypes][dim][dim][dim]);
    gmaker.setAtomsCPU(transform.mol.atoms, transform.mol.whichGrid, transform.Q.boost(), grids);
  }
}

//return a string representation of the atom type(s) represented by index
//in map - this isn't particularly efficient, but is only for debug purposes
template <typename Dtype>
string MolGridDataLayer<Dtype>::getIndexName(const vector<int>& map, unsigned index) const
		{
	stringstream ret;
	stringstream altret;
	for (unsigned at = 0; at < smina_atom_type::NumTypes; at++)
	{
		if (map[at] == index)
		{
			ret << smina_type_to_string((smt) at);
			altret << "_" << at;
		}
	}

	if (ret.str().length() > 32) //there are limits on file name lengths
		return altret.str();
	else
		return ret.str();
}

//output a grid the file in dx format (for debug)
template<typename Dtype>
void MolGridDataLayer<Dtype>::outputDXGrid(std::ostream& out, Grids& grid, unsigned g, double scale) const
{
  unsigned n = dim;
  out.precision(5);
  setprecision(5);
  out << fixed;
  out << "object 1 class gridpositions counts " << n << " " << n << " " << " " << n << "\n";
  out << "origin";
  for (unsigned i = 0; i < 3; i++)
  {
    out << " " << mem_lig.center[i]-dimension/2.0;
  }
  out << "\n";
  out << "delta " << resolution << " 0 0\ndelta 0 " << resolution << " 0\ndelta 0 0 " << resolution << "\n";
  out << "object 2 class gridconnections counts " << n << " " << n << " " << " " << n << "\n";
  out << "object 3 class array type double rank 0 items [ " << n*n*n << "] data follows\n";
  //now coordinates - x,y,z
  out << scientific;
  out.precision(6);
  unsigned total = 0;
  for (unsigned i = 0; i < n; i++)
  {
    for (unsigned j = 0; j < n; j++)
    {
      for (unsigned k = 0; k < n; k++)
      {
        out << grid[g][i][j][k]*scale;
        total++;
        if(total % 3 == 0) out << "\n";
        else out << " ";
      }
    }
  }
}

//dump dx files for every atom type, with files names starting with prefix
//only does the very first grid for now
template<typename Dtype>
void MolGridDataLayer<Dtype>::dumpDiffDX(const std::string& prefix,
		Blob<Dtype>* top, double scale) const
{
	Grids grids(top->mutable_cpu_diff(),
			boost::extents[numReceptorTypes + numLigandTypes][dim][dim][dim]);
    CHECK_GT(mem_lig.atoms.size(),0) << "DX dump only works with in-memory ligand";
    CHECK_EQ(randrotate, false) << "DX dump requires no rotation";
	for (unsigned a = 0, na = numReceptorTypes; a < na; a++) {
		string name = getIndexName(rmap, a);
		string fname = prefix + "_rec_" + name + ".dx";
		ofstream out(fname.c_str());
		outputDXGrid(out, grids, a, scale);
	}
	for (unsigned a = 0, na = numLigandTypes; a < na; a++) {
			string name = getIndexName(lmap, a);
			string fname = prefix + "_lig_" + name + ".dx";
			ofstream out(fname.c_str());
			outputDXGrid(out, grids, numReceptorTypes+a, scale);
	}

}

template <typename Dtype>
void MolGridDataLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top)
{
  forward(bottom, top, false);
}

template <typename Dtype>
void MolGridDataLayer<Dtype>::forward(const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top, bool gpu)
{
  bool hasaffinity = this->layer_param_.molgrid_data_param().has_affinity();
  bool hasrmsd = this->layer_param_.molgrid_data_param().has_rmsd();

  Dtype *top_data = NULL;
  if(gpu)
    top_data = top[0]->mutable_gpu_data();
  else
    top_data = top[0]->mutable_cpu_data();

  perturbations.clear();
  unsigned batch_size = top_shape[0];
  output_transform peturb;

  //if in memory must be set programmatically
  if(inmem)
  {
    CHECK_GT(mem_rec.atoms.size(),0) << "Receptor not set in MolGridDataLayer";
    CHECK_GT(mem_lig.atoms.size(),0) << "Ligand not set in MolGridDataLayer";
    //memory is now available
    set_grid_minfo(top_data, mem_rec, mem_lig, batch_transform[0], peturb, gpu); //TODO how do we know what batch position?
    perturbations.push_back(peturb);

    if (num_rotations > 0) {
      current_rotation = (current_rotation+1)%num_rotations;
    }

    CHECK_GT(labels.size(),0) << "Did not set labels in memory based molgrid";

  }
  else
  {
    //clear batch labels
    labels.clear();
    affinities.clear();
    rmsds.clear();

    //percent of batch from first data source
    unsigned dataswitch = batch_size;
    if (data2)
      dataswitch = batch_size*data_ratio/(data_ratio+1);

    for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx)
    {
      example ex;
      string *root;
      if (batch_idx < dataswitch)
      {
        data->next(ex);
        root = &root_folder;
      }
      else
      {
        data2->next(ex);
        root = &root_folder2;
      }

      labels.push_back(ex.label);
      affinities.push_back(ex.affinity);
      rmsds.push_back(ex.rmsd);

      int offset = batch_idx*example_size;
      set_grid_ex(top_data+offset, ex, *root, batch_transform[batch_idx], peturb, gpu);
      perturbations.push_back(peturb);
      //NOTE: num_rotations not actually implemented!
    }

  }

  if(gpu) {
    caffe_copy(labels.size(), &labels[0], top[1]->mutable_gpu_data());
    if(hasaffinity) {
      caffe_copy(affinities.size(), &affinities[0], top[2]->mutable_gpu_data());
      if(hasrmsd) {
        caffe_copy(rmsds.size(), &rmsds[0], top[3]->mutable_gpu_data());
      }
    } else if(hasrmsd) {
      caffe_copy(rmsds.size(), &rmsds[0], top[2]->mutable_gpu_data());
    }

    if(ligpeturb) {
      //trusting struct layout is normal
      caffe_copy(perturbations.size()*6, (Dtype*)&perturbations[0], top.back()->mutable_gpu_data());
    }

  }
  else {
    caffe_copy(labels.size(), &labels[0], top[1]->mutable_cpu_data());
    if(hasaffinity) {
      caffe_copy(affinities.size(), &affinities[0], top[2]->mutable_cpu_data());
      if(hasrmsd) {
        caffe_copy(rmsds.size(), &rmsds[0], top[3]->mutable_cpu_data());
      }
    } else if(hasrmsd) {
      caffe_copy(rmsds.size(), &rmsds[0], top[2]->mutable_cpu_data());
    }
    if(ligpeturb) {
      //trusting struct layout is normal
      caffe_copy(perturbations.size()*6, (Dtype*)&perturbations[0], top.back()->mutable_cpu_data());
    }
  }

}

template <typename Dtype>
void MolGridDataLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom)
{
  backward(top, bottom, false);
}

/* backpropagates gradients onto atoms (note there is not actual bottom)
 * Only performed when compute_atom_gradients is true.
 */
template <typename Dtype>
void MolGridDataLayer<Dtype>::backward(const vector<Blob<Dtype>*>& top, const vector<Blob<Dtype>*>& bottom, bool gpu)
{
  //propagate gradient grid onto atom positions
  if(compute_atom_gradients) {
    unsigned batch_size = top_shape[0];
    Dtype *diff = NULL;
    if(gpu) {
      diff = top[0]->mutable_gpu_diff();
      setAtomGradientsGPU(gmaker, diff, batch_size);
    }
    else {
      diff = top[0]->mutable_cpu_diff();
      for (int item_id = 0; item_id < batch_size; ++item_id) {

        int offset = item_id*example_size;
        Grids grids(diff+offset, boost::extents[numReceptorTypes+numLigandTypes][dim][dim][dim]);

        mol_transform& transform = batch_transform[item_id];
        gmaker.setCenter(transform.center[0], transform.center[1], transform.center[2]);
        gmaker.setAtomGradientsCPU(transform.mol.atoms, transform.mol.whichGrid,
                transform.Q.boost(), grids, transform.mol.gradient);
      }
    }
  }
}

template <typename Dtype>
void MolGridDataLayer<Dtype>::Backward_relevance(const vector<Blob<Dtype>*>& top, const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom)
{

  /*
  float top_sum = 0.0;

  std::cout << "MOLGRID TOP:";
  for(int i = 0; i < top[0]->count(); i++)
  {
          top_sum += top[0]->cpu_diff()[i];
  }
  std::cout << "MOLGRID TOP: " << top_sum << '\n';
  */

  Dtype *diff = top[0]->mutable_cpu_diff(); //TODO: implement gpu

  //propagate gradient grid onto atom positions
  unsigned batch_size = top_shape[0];
  for (int item_id = 0; item_id < batch_size; ++item_id) {

    int offset = item_id*example_size;
    Grids grids(diff+offset, boost::extents[numReceptorTypes+numLigandTypes][dim][dim][dim]);

    mol_transform& transform = batch_transform[item_id];
    gmaker.setCenter(transform.center[0], transform.center[1], transform.center[2]);
    gmaker.setAtomRelevanceCPU(transform.mol.atoms, transform.mol.whichGrid, transform.Q.boost(), grids,
        transform.mol.gradient);
  }

  //float bottom_sum = 0.0;
  //for(int i = 0; i < bottom[0]->count(); i++)
  //{
  //        bottom_sum += bottom[0]->cpu_diff()[i];
  //}
  //std::cout << "MOLGRID BOTTOM: " << bottom_sum << '\n';

}

INSTANTIATE_CLASS(MolGridDataLayer);
REGISTER_LAYER_CLASS(MolGridData);

}  // namespace caffe