<commit_before>/*
   Copyright 2014 Software Reliability Lab, ETH Zurich

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
 */

#include "graph_inference.h"

#include <stdio.h>
#include <math.h>
#include <queue>
#include <unordered_set>
#include <algorithm>
#include <random>
#include <iterator>

#include "stringprintf.h"
#include "glog/logging.h"
#include "gflags/gflags.h"

#include "base.h"
#include "maputil.h"
#include "nbest.h"
#include "simple_histogram.h"
#include "updatable_priority_queue.h"

#include "label_set.h"

DEFINE_bool(initial_greedy_assignment_pass, true, "Whether to run an initial greedy assignment pass.");
DEFINE_bool(duplicate_name_resolution, true, "Whether to attempt a duplicate name resultion on conflicts.");
DEFINE_int32(graph_per_node_passes, 8, "Number of per-node passes for inference");
DEFINE_int32(graph_per_arc_passes, 5, "Number of per-arc passes for inference");
DEFINE_int32(graph_per_factor_passes, 1, "Number of per-factor passes for inference");
DEFINE_int32(graph_loopy_bp_passes, 0, "Number of loopy belief propagation passes for inference");
DEFINE_int32(graph_loopy_bp_steps_per_pass, 3, "Number of loopy belief propagation steps in each inference pass");
DEFINE_int32(skip_per_arc_optimization_for_nodes_above_degree, 32,
    "Skip the per-arc optimization pass if an edge is connected to a node with the in+out degree more than the given value");

DEFINE_bool(use_factors, true, "Flag that enable the use of the factors in training and MAP inference.");
DEFINE_int32(maximum_depth, 2, "Maximum depth of the multi-level map used to store the factor features");
DEFINE_int32(factors_limit, 128, "Maximum number of factor candidates considered for inference using factor features");
DEFINE_uint64(permutations_beam_size, 64, "Maximum number of permutations of the assignments to unknown labels in a factor");

DEFINE_string(valid_labels, "valid_names.txt", "A file describing valid names");
DEFINE_string(unknown_label,
    "", "A special label that denotes that a label is of low frequency (and thus unknown).");
DEFINE_int32(min_freq_known_label, 0,
    "Minimum number of graphs a label must appear it in order to not be declared unknown");

static const size_t kInitialAssignmentBeamSize = 4;

static const size_t kStartPerArcBeamSize = 64;
static const size_t kMaxPerArcBeamSize = 64;

static const size_t kStartPerNodeBeamSize = 64;
static const size_t kMaxPerNodeBeamSize = 64;
static const size_t kLoopyBPBeamSize = 32;

static const size_t kFactorsLimitBeforeGoingDepperMultiLevelMap = 16;

namespace std {
  template <> struct hash<Json::Value> {
    size_t operator()(const Json::Value& v) const {
      if (v.isInt()) {
        return hash<int>()(v.asInt());
      }
      if (v.isString()) {
        const char* str = v.asCString();
        size_t r = 1;
        while (*str) {
          r = (r * 17) + (*str);
          ++str;
        }
        return r;
      }
      LOG(INFO) << v;
      return 0;
    }
  };
}

// Returns -1 if the result will overflow.
uint64 CalculateFactorial(int n) {
  uint64 result = 1;
  for (size_t i = n; i > 0; i--) {
    // Check if result will overflow.
    if (result > (UINT64_MAX / i)) {
      return -1;
    }
    result *= i;
  }
  return result;
}

// http://stackoverflow.com/a/6867612
// https://github.com/aappleby/smhasher/blob/master/src/MurmurHash3.cpp
uint64 HashInt(uint64 x) {
  x++;
  x ^= x >> 33;
  x *= 0xff51afd7ed558ccd;
  x ^= x >> 33;
  x *= 0xc4ceb9fe1a85ec53;
  x ^= x >> 33;
  return x;
}

class JsonValueNumberer {
public:
  int ValueToNumber(const Json::Value& val) {
    auto ins = data_.insert(std::pair<Json::Value, int>(val, data_.size()));
    if (ins.second) {
      number_to_value_.push_back(val);
    }
    return ins.first->second;
  }

  int ValueToNumberOrDie(const Json::Value& val) const {
    auto it = data_.find(val);
    CHECK(it != data_.end());
    return it->second;
  }

  int ValueToNumberWithDefault(const Json::Value& val, int default_number) const {
    auto it = data_.find(val);
    if(it == data_.end()) return default_number;
    return it->second;
  }

  const Json::Value& NumberToValue(int number) const {
    return number_to_value_[number];
  }

  int size() const {
    return data_.size();
  }

private:
  std::unordered_map<Json::Value, int> data_;
  std::vector<Json::Value> number_to_value_;
};


class GraphQuery : public Nice2Query {
public:
  explicit GraphQuery(const StringSet* ss, const LabelChecker* checker) : label_set_(ss, checker) {
    arcs_connecting_node_pair_.set_empty_key(IntPair(-1, -1));
    arcs_connecting_node_pair_.set_deleted_key(IntPair(-2, -2));
  }
  virtual ~GraphQuery() {
  }

  virtual void FromJSON(const Json::Value& query) override {
    arcs_.clear();
    factors_.clear();

    CHECK(query.isArray());
    for (const Json::Value& arc : query) {
      if (arc.isMember("f2")) {
        // A factor connecting two facts (an arc).
        Arc a;
        a.node_a = numberer_.ValueToNumber(arc["a"]);
        a.node_b = numberer_.ValueToNumber(arc["b"]);
        a.type = label_set_.ss()->findString(arc["f2"].asCString());
        if (a.type < 0) continue;
        arcs_.push_back(a);
      }
      if (arc.isMember("cn")) {
        // A scope that lists names that cannot be assigned to the same value.
        const Json::Value& v = arc["n"];
        if (v.isArray()) {
          std::vector<int> scope_vars;
          scope_vars.reserve(v.size());
          for (const Json::Value& item : v) {
            scope_vars.push_back(numberer_.ValueToNumber(item));
          }
          std::sort(scope_vars.begin(), scope_vars.end());
          scope_vars.erase(std::unique(scope_vars.begin(), scope_vars.end()), scope_vars.end());
          nodes_in_scope_.push_back(std::move(scope_vars));
        }
      }
      if (FLAGS_use_factors) {
        if (arc.isMember("group")) {
          const Json::Value& v = arc["group"];
          if (v.isArray()) {
            Factor factor_vars;
            for (const Json::Value& item : v) {
              factor_vars.insert(numberer_.ValueToNumber(item));
            }
            factors_.push_back(factor_vars);
          }
        }
      }
    }
    std::sort(arcs_.begin(), arcs_.end());
    arcs_adjacent_to_node_.assign(numberer_.size(), std::vector<Arc>());
    for (const Arc& a : arcs_) {
      arcs_adjacent_to_node_[a.node_a].push_back(a);
      arcs_adjacent_to_node_[a.node_b].push_back(a);
    }
    for (std::vector<Arc>& v : arcs_adjacent_to_node_) {
      std::sort(v.begin(), v.end());
      v.erase(std::unique(v.begin(), v.end()), v.end());
    }

    arcs_connecting_node_pair_.clear();
    for (const Arc& a : arcs_) {
      arcs_connecting_node_pair_[IntPair(a.node_a, a.node_b)].push_back(a);
      arcs_connecting_node_pair_[IntPair(a.node_b, a.node_a)].push_back(a);
    }

    scopes_per_nodes_.assign(numberer_.size(), std::vector<int>());
    for (size_t scope = 0; scope < nodes_in_scope_.size(); ++scope) {
      for (int node : nodes_in_scope_[scope]) {
        scopes_per_nodes_[node].push_back(scope);
      }
    }

    factors_of_a_node_.assign(numberer_.size(), std::vector<int>());
    for (size_t i = 0; i < factors_.size(); ++i) {
      for (auto var = factors_[i].begin(); var != factors_[i].end(); ++var) {
        factors_of_a_node_[*var].push_back(i);
      }
    }
  }

private:
  struct Arc {
    int node_a, node_b, type;
    bool operator==(const Arc& o) const {
      return node_a == o.node_a && node_b == o.node_b && type == o.type;
    }
    bool operator<(const Arc& o) const {
      if (node_a != o.node_a) return node_a < o.node_a;
      if (node_b != o.node_b) return node_b < o.node_b;
      return type < o.type;
    }
  };

  std::vector<std::vector<Arc> > arcs_adjacent_to_node_;
  std::vector<std::vector<int> > factors_of_a_node_;
  std::vector<Arc> arcs_;
  std::vector<Factor> factors_;
  google::dense_hash_map<IntPair, std::vector<Arc> > arcs_connecting_node_pair_;
  JsonValueNumberer numberer_;

  LabelSet label_set_;

  std::vector<std::vector<int> > nodes_in_scope_;
  std::vector<std::vector<int> > scopes_per_nodes_;

  friend class GraphNodeAssignment;
  friend class LoopyBPInference;
  friend class GraphInference;
};

#ifdef GRAPH_INFERENCE_STATS
struct GraphInferenceStats {
  GraphInferenceStats()
      : position_of_best_per_node_label(10),
        position_of_best_per_arc_label(10),
        label_candidates_per_node(10) {
  }

  SimpleHistogram position_of_best_per_node_label;
  SimpleHistogram position_of_best_per_arc_label;
  SimpleHistogram label_candidates_per_node;

  std::string ToString() const {
    std::string result;
    result += "Candidates per node:\n";
    result += label_candidates_per_node.ToString();
    result += "Improving position @ node:\n";
    result += position_of_best_per_node_label.ToString();
    result += "Improving position @ arc:\n";
    result += position_of_best_per_arc_label.ToString();
    return result;
  }
};
#endif


class GraphNodeAssignment : public Nice2Assignment {
public:
  GraphNodeAssignment(const GraphQuery* query, LabelSet* label_set, int unknown_label)
    : query_(query), label_set_(label_set), unknown_label_(unknown_label) {
    assignments_.assign(query_->numberer_.size(), Assignment());
    penalties_.assign(assignments_.size(), LabelPenalty());
  }
  virtual ~GraphNodeAssignment() {
  }

  virtual void SetUpEqualityPenalty(double penalty) override {
    ClearPenalty();
    for (size_t i = 0; i < assignments_.size(); ++i) {
      if (assignments_[i].must_infer) {
        penalties_[i].label = assignments_[i].label;
        penalties_[i].penalty = penalty;
      }
    }
  }

  virtual void ClearPenalty() override {
    penalties_.assign(assignments_.size(), LabelPenalty());
  }

  virtual void FromJSON(const Json::Value& assignment) override {
    CHECK(assignment.isArray());
    assignments_.assign(query_->numberer_.size(), Assignment());
    for (const Json::Value& a : assignment) {
      Assignment aset;
      if (a.isMember("inf")) {
        aset.label = label_set_->AddLabelName(a["inf"].asCString());
        aset.must_infer = true;
      } else {
        CHECK(a.isMember("giv"));
        aset.label = label_set_->AddLabelName(a["giv"].asCString());
        aset.must_infer = false;
      }
      int number = query_->numberer_.ValueToNumberWithDefault(a.get("v", Json::Value::null), -1);
      if (number != -1) {
        assignments_[number] = aset;
      }
    }

    ClearPenalty();
  }

  virtual void ToJSON(Json::Value* assignment) const override {
    *assignment = Json::Value(Json::arrayValue);
    for (size_t i = 0; i < assignments_.size(); ++i) {
      if (assignments_[i].label < 0) continue;

      Json::Value obj(Json::objectValue);
      obj["v"] = query_->numberer_.NumberToValue(i);
      const char* str_value = label_set_->GetLabelName(assignments_[i].label);
      if (assignments_[i].must_infer) {
        obj["inf"] = Json::Value(str_value);
      } else {
        obj["giv"] = Json::Value(str_value);
      }
      assignment->append(obj);
    }
  }
  
  std::vector<std::pair<int, double>> GetCandidatesForNode(Nice2Inference& inference, int node, int n) {
        GraphInference& graphInference = static_cast<GraphInference&>(inference);
        std::vector<int> candidates;
        candidates.clear();
        GetLabelCandidates(graphInference, node, &candidates, kMaxPerArcBeamSize);

        std::vector<std::pair<int, double>> scoredCandidates;
        Assignment& nodea = assignments_[node];
        int originalLabel = nodea.label;
        for (size_t i = 0; i < candidates.size() ; i++) {
          int candidate = candidates[i];
          nodea.label = candidates[i];
          double score = GetNodeScore(graphInference, node);
          scoredCandidates.push_back(std::pair<int, double>(candidate, score));
        }
        nodea.label = originalLabel;

        std::sort(scoredCandidates.begin(), scoredCandidates.end(), [](const std::pair<int,double> &left, const std::pair<int,double> &right) {
          return right.second < left.second;
        });
        return scoredCandidates;
  }
 
  virtual void GetCandidates(
      Nice2Inference& inference,
      int node,
      int n,
      Json::Value* response) {
    *response = Json::Value(Json::arrayValue);
    
    if (node != -1) {
        std::vector<std::pair<int, double>> scoredCandidates = GetCandidatesForNode(inference, node, n);
        for (size_t i = 0; i < scoredCandidates.size() && i < (size_t)((unsigned)n) ; i++) {
          Json::Value obj(Json::objectValue);
          obj["candidate"] = label_set_->GetLabelName(scoredCandidates[i].first);
          obj["score"] = scoredCandidates[i].second;
          response->append(obj);
        }
    }
    else {
        for (size_t i = 0; i < assignments_.size(); ++i) {
        //for (int i = assignments_.size() - 1; i>= 0; --i) {
          if (assignments_[i].must_infer) {
            //LOG(INFO) << "Adding node " << i << " to response";
            std::vector<std::pair<int, double>> scoredCandidates = GetCandidatesForNode(inference, i, n);
            Json::Value nodeResults(Json::objectValue);
            nodeResults["node"] = (int)i;
            Json::Value nodeCandidates(Json::arrayValue);
            for (size_t j = 0; j < scoredCandidates.size() && j < (size_t)((unsigned)n) ; j++) {
              Json::Value obj(Json::objectValue);
              obj["candidate"] = label_set_->GetLabelName(scoredCandidates[j].first);
              obj["score"] = scoredCandidates[j].second;
              nodeCandidates.append(obj);
            }
            nodeResults["candidates"] = nodeCandidates;
            response->append(nodeResults);
          }
        }
    }
  }

  virtual void ClearInferredAssignment() override {
    for (size_t i = 0; i < assignments_.size(); ++i) {
      if (assignments_[i].must_infer) {
        assignments_[i].label = -1;<commit_msg>fix node Id bug in candidates<commit_after>/*
   Copyright 2014 Software Reliability Lab, ETH Zurich

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
 */

#include "graph_inference.h"

#include <stdio.h>
#include <math.h>
#include <queue>
#include <unordered_set>
#include <algorithm>
#include <random>
#include <iterator>

#include "stringprintf.h"
#include "glog/logging.h"
#include "gflags/gflags.h"

#include "base.h"
#include "maputil.h"
#include "nbest.h"
#include "simple_histogram.h"
#include "updatable_priority_queue.h"

#include "label_set.h"

DEFINE_bool(initial_greedy_assignment_pass, true, "Whether to run an initial greedy assignment pass.");
DEFINE_bool(duplicate_name_resolution, true, "Whether to attempt a duplicate name resultion on conflicts.");
DEFINE_int32(graph_per_node_passes, 8, "Number of per-node passes for inference");
DEFINE_int32(graph_per_arc_passes, 5, "Number of per-arc passes for inference");
DEFINE_int32(graph_per_factor_passes, 1, "Number of per-factor passes for inference");
DEFINE_int32(graph_loopy_bp_passes, 0, "Number of loopy belief propagation passes for inference");
DEFINE_int32(graph_loopy_bp_steps_per_pass, 3, "Number of loopy belief propagation steps in each inference pass");
DEFINE_int32(skip_per_arc_optimization_for_nodes_above_degree, 32,
    "Skip the per-arc optimization pass if an edge is connected to a node with the in+out degree more than the given value");

DEFINE_bool(use_factors, true, "Flag that enable the use of the factors in training and MAP inference.");
DEFINE_int32(maximum_depth, 2, "Maximum depth of the multi-level map used to store the factor features");
DEFINE_int32(factors_limit, 128, "Maximum number of factor candidates considered for inference using factor features");
DEFINE_uint64(permutations_beam_size, 64, "Maximum number of permutations of the assignments to unknown labels in a factor");

DEFINE_string(valid_labels, "valid_names.txt", "A file describing valid names");
DEFINE_string(unknown_label,
    "", "A special label that denotes that a label is of low frequency (and thus unknown).");
DEFINE_int32(min_freq_known_label, 0,
    "Minimum number of graphs a label must appear it in order to not be declared unknown");

static const size_t kInitialAssignmentBeamSize = 4;

static const size_t kStartPerArcBeamSize = 64;
static const size_t kMaxPerArcBeamSize = 64;

static const size_t kStartPerNodeBeamSize = 64;
static const size_t kMaxPerNodeBeamSize = 64;
static const size_t kLoopyBPBeamSize = 32;

static const size_t kFactorsLimitBeforeGoingDepperMultiLevelMap = 16;

namespace std {
  template <> struct hash<Json::Value> {
    size_t operator()(const Json::Value& v) const {
      if (v.isInt()) {
        return hash<int>()(v.asInt());
      }
      if (v.isString()) {
        const char* str = v.asCString();
        size_t r = 1;
        while (*str) {
          r = (r * 17) + (*str);
          ++str;
        }
        return r;
      }
      LOG(INFO) << v;
      return 0;
    }
  };
}

// Returns -1 if the result will overflow.
uint64 CalculateFactorial(int n) {
  uint64 result = 1;
  for (size_t i = n; i > 0; i--) {
    // Check if result will overflow.
    if (result > (UINT64_MAX / i)) {
      return -1;
    }
    result *= i;
  }
  return result;
}

// http://stackoverflow.com/a/6867612
// https://github.com/aappleby/smhasher/blob/master/src/MurmurHash3.cpp
uint64 HashInt(uint64 x) {
  x++;
  x ^= x >> 33;
  x *= 0xff51afd7ed558ccd;
  x ^= x >> 33;
  x *= 0xc4ceb9fe1a85ec53;
  x ^= x >> 33;
  return x;
}

class JsonValueNumberer {
public:
  int ValueToNumber(const Json::Value& val) {
    auto ins = data_.insert(std::pair<Json::Value, int>(val, data_.size()));
    if (ins.second) {
      number_to_value_.push_back(val);
    }
    return ins.first->second;
  }

  int ValueToNumberOrDie(const Json::Value& val) const {
    auto it = data_.find(val);
    CHECK(it != data_.end());
    return it->second;
  }

  int ValueToNumberWithDefault(const Json::Value& val, int default_number) const {
    auto it = data_.find(val);
    if(it == data_.end()) return default_number;
    return it->second;
  }

  const Json::Value& NumberToValue(int number) const {
    return number_to_value_[number];
  }

  int size() const {
    return data_.size();
  }

private:
  std::unordered_map<Json::Value, int> data_;
  std::vector<Json::Value> number_to_value_;
};


class GraphQuery : public Nice2Query {
public:
  explicit GraphQuery(const StringSet* ss, const LabelChecker* checker) : label_set_(ss, checker) {
    arcs_connecting_node_pair_.set_empty_key(IntPair(-1, -1));
    arcs_connecting_node_pair_.set_deleted_key(IntPair(-2, -2));
  }
  virtual ~GraphQuery() {
  }

  virtual void FromJSON(const Json::Value& query) override {
    arcs_.clear();
    factors_.clear();

    CHECK(query.isArray());
    for (const Json::Value& arc : query) {
      if (arc.isMember("f2")) {
        // A factor connecting two facts (an arc).
        Arc a;
        a.node_a = numberer_.ValueToNumber(arc["a"]);
        a.node_b = numberer_.ValueToNumber(arc["b"]);
        a.type = label_set_.ss()->findString(arc["f2"].asCString());
        if (a.type < 0) continue;
        arcs_.push_back(a);
      }
      if (arc.isMember("cn")) {
        // A scope that lists names that cannot be assigned to the same value.
        const Json::Value& v = arc["n"];
        if (v.isArray()) {
          std::vector<int> scope_vars;
          scope_vars.reserve(v.size());
          for (const Json::Value& item : v) {
            scope_vars.push_back(numberer_.ValueToNumber(item));
          }
          std::sort(scope_vars.begin(), scope_vars.end());
          scope_vars.erase(std::unique(scope_vars.begin(), scope_vars.end()), scope_vars.end());
          nodes_in_scope_.push_back(std::move(scope_vars));
        }
      }
      if (FLAGS_use_factors) {
        if (arc.isMember("group")) {
          const Json::Value& v = arc["group"];
          if (v.isArray()) {
            Factor factor_vars;
            for (const Json::Value& item : v) {
              factor_vars.insert(numberer_.ValueToNumber(item));
            }
            factors_.push_back(factor_vars);
          }
        }
      }
    }
    std::sort(arcs_.begin(), arcs_.end());
    arcs_adjacent_to_node_.assign(numberer_.size(), std::vector<Arc>());
    for (const Arc& a : arcs_) {
      arcs_adjacent_to_node_[a.node_a].push_back(a);
      arcs_adjacent_to_node_[a.node_b].push_back(a);
    }
    for (std::vector<Arc>& v : arcs_adjacent_to_node_) {
      std::sort(v.begin(), v.end());
      v.erase(std::unique(v.begin(), v.end()), v.end());
    }

    arcs_connecting_node_pair_.clear();
    for (const Arc& a : arcs_) {
      arcs_connecting_node_pair_[IntPair(a.node_a, a.node_b)].push_back(a);
      arcs_connecting_node_pair_[IntPair(a.node_b, a.node_a)].push_back(a);
    }

    scopes_per_nodes_.assign(numberer_.size(), std::vector<int>());
    for (size_t scope = 0; scope < nodes_in_scope_.size(); ++scope) {
      for (int node : nodes_in_scope_[scope]) {
        scopes_per_nodes_[node].push_back(scope);
      }
    }

    factors_of_a_node_.assign(numberer_.size(), std::vector<int>());
    for (size_t i = 0; i < factors_.size(); ++i) {
      for (auto var = factors_[i].begin(); var != factors_[i].end(); ++var) {
        factors_of_a_node_[*var].push_back(i);
      }
    }
  }

private:
  struct Arc {
    int node_a, node_b, type;
    bool operator==(const Arc& o) const {
      return node_a == o.node_a && node_b == o.node_b && type == o.type;
    }
    bool operator<(const Arc& o) const {
      if (node_a != o.node_a) return node_a < o.node_a;
      if (node_b != o.node_b) return node_b < o.node_b;
      return type < o.type;
    }
  };

  std::vector<std::vector<Arc> > arcs_adjacent_to_node_;
  std::vector<std::vector<int> > factors_of_a_node_;
  std::vector<Arc> arcs_;
  std::vector<Factor> factors_;
  google::dense_hash_map<IntPair, std::vector<Arc> > arcs_connecting_node_pair_;
  JsonValueNumberer numberer_;

  LabelSet label_set_;

  std::vector<std::vector<int> > nodes_in_scope_;
  std::vector<std::vector<int> > scopes_per_nodes_;

  friend class GraphNodeAssignment;
  friend class LoopyBPInference;
  friend class GraphInference;
};

#ifdef GRAPH_INFERENCE_STATS
struct GraphInferenceStats {
  GraphInferenceStats()
      : position_of_best_per_node_label(10),
        position_of_best_per_arc_label(10),
        label_candidates_per_node(10) {
  }

  SimpleHistogram position_of_best_per_node_label;
  SimpleHistogram position_of_best_per_arc_label;
  SimpleHistogram label_candidates_per_node;

  std::string ToString() const {
    std::string result;
    result += "Candidates per node:\n";
    result += label_candidates_per_node.ToString();
    result += "Improving position @ node:\n";
    result += position_of_best_per_node_label.ToString();
    result += "Improving position @ arc:\n";
    result += position_of_best_per_arc_label.ToString();
    return result;
  }
};
#endif


class GraphNodeAssignment : public Nice2Assignment {
public:
  GraphNodeAssignment(const GraphQuery* query, LabelSet* label_set, int unknown_label)
    : query_(query), label_set_(label_set), unknown_label_(unknown_label) {
    assignments_.assign(query_->numberer_.size(), Assignment());
    penalties_.assign(assignments_.size(), LabelPenalty());
  }
  virtual ~GraphNodeAssignment() {
  }

  virtual void SetUpEqualityPenalty(double penalty) override {
    ClearPenalty();
    for (size_t i = 0; i < assignments_.size(); ++i) {
      if (assignments_[i].must_infer) {
        penalties_[i].label = assignments_[i].label;
        penalties_[i].penalty = penalty;
      }
    }
  }

  virtual void ClearPenalty() override {
    penalties_.assign(assignments_.size(), LabelPenalty());
  }

  virtual void FromJSON(const Json::Value& assignment) override {
    CHECK(assignment.isArray());
    assignments_.assign(query_->numberer_.size(), Assignment());
    for (const Json::Value& a : assignment) {
      Assignment aset;
      if (a.isMember("inf")) {
        aset.label = label_set_->AddLabelName(a["inf"].asCString());
        aset.must_infer = true;
      } else {
        CHECK(a.isMember("giv"));
        aset.label = label_set_->AddLabelName(a["giv"].asCString());
        aset.must_infer = false;
      }
      int number = query_->numberer_.ValueToNumberWithDefault(a.get("v", Json::Value::null), -1);
      if (number != -1) {
        assignments_[number] = aset;
      }
    }

    ClearPenalty();
  }

  virtual void ToJSON(Json::Value* assignment) const override {
    *assignment = Json::Value(Json::arrayValue);
    for (size_t i = 0; i < assignments_.size(); ++i) {
      if (assignments_[i].label < 0) continue;

      Json::Value obj(Json::objectValue);
      obj["v"] = query_->numberer_.NumberToValue(i);
      const char* str_value = label_set_->GetLabelName(assignments_[i].label);
      if (assignments_[i].must_infer) {
        obj["inf"] = Json::Value(str_value);
      } else {
        obj["giv"] = Json::Value(str_value);
      }
      assignment->append(obj);
    }
  }
  
  std::vector<std::pair<int, double>> GetCandidatesForNode(Nice2Inference& inference, int node, int n) {
        GraphInference& graphInference = static_cast<GraphInference&>(inference);
        std::vector<int> candidates;
        candidates.clear();
        GetLabelCandidates(graphInference, node, &candidates, kMaxPerArcBeamSize);

        std::vector<std::pair<int, double>> scoredCandidates;
        Assignment& nodea = assignments_[node];
        int originalLabel = nodea.label;
        for (size_t i = 0; i < candidates.size() ; i++) {
          int candidate = candidates[i];
          nodea.label = candidates[i];
          double score = GetNodeScore(graphInference, node);
          scoredCandidates.push_back(std::pair<int, double>(candidate, score));
        }
        nodea.label = originalLabel;

        std::sort(scoredCandidates.begin(), scoredCandidates.end(), [](const std::pair<int,double> &left, const std::pair<int,double> &right) {
          return right.second < left.second;
        });
        return scoredCandidates;
  }
 
  virtual void GetCandidates(
      Nice2Inference& inference,
      int node,
      int n,
      Json::Value* response) {
    *response = Json::Value(Json::arrayValue);
    
    if (node != -1) {
        std::vector<std::pair<int, double>> scoredCandidates = GetCandidatesForNode(inference, node, n);
        for (size_t i = 0; i < scoredCandidates.size() && i < (size_t)((unsigned)n) ; i++) {
          Json::Value obj(Json::objectValue);
          obj["candidate"] = label_set_->GetLabelName(scoredCandidates[i].first);
          obj["score"] = scoredCandidates[i].second;
          response->append(obj);
        }
    }
    else {
        for (size_t i = 0; i < assignments_.size(); ++i) {
        //for (int i = assignments_.size() - 1; i>= 0; --i) {
          if (assignments_[i].must_infer) {
            //LOG(INFO) << "Adding node " << i << " to response, id: " << query_->numberer_.NumberToValue(i) << ", name: " << label_set_->GetLabelName(assignments_[i].label);
            std::vector<std::pair<int, double>> scoredCandidates = GetCandidatesForNode(inference, i, n);
            Json::Value nodeResults(Json::objectValue);
            nodeResults["node"] = query_->numberer_.NumberToValue(i);
            Json::Value nodeCandidates(Json::arrayValue);
            for (size_t j = 0; j < scoredCandidates.size() && j < (size_t)((unsigned)n) ; j++) {
              Json::Value obj(Json::objectValue);
              obj["candidate"] = label_set_->GetLabelName(scoredCandidates[j].first);
              obj["score"] = scoredCandidates[j].second;
              nodeCandidates.append(obj);
            }
            nodeResults["candidates"] = nodeCandidates;
            response->append(nodeResults);
          }
        }
    }
  }

  virtual void ClearInferredAssignment() override {
    for (size_t i = 0; i < assignments_.size(); ++i) {
      if (assignments_[i].must_infer) {
        assignments_[i].label = -1;
{'fraction_non_alphanumeric': 0.13367989809618894, 'fraction_numerical': 0.009295280063345613, 'mean_word_length': 3.326481977956509, 'pattern_counts': {'":': 0, '<': 147, '<?xml version=': 0, '>': 142, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5375398', 'n_tokens_mistral': 9714, 'n_tokens_neox': 9350, 'n_words': 2144}