#include "neighbor_attack.h"

#include <algorithm>
#include <boost/asio.hpp>
#include <boost/bind.hpp>
#include <cstdlib>
#include <functional>
#include <iostream>
#include <list>
#include <random>

#include <boost/lexical_cast.hpp>

#include "bounding_box.h"
#include "decision_forest.h"
#include "decision_tree.h"
#include "interval.h"
#include "timing.h"

using std::cout;
using std::endl;

namespace cz {

namespace {

class DualNormScore : public cz::NeighborScore {
 public:
  DualNormScore(int norm_type_1)
      : norm_type_1_(norm_type_1), norm_type_2_(norm_type_1 == 2 ? -1 : 2) {}

  ScoreType Slow(const Point& point,
                 const Point& ref_point) const override {
    Timing::Instance()->StartTimer("DualNormScore::Slow");
    double norm1 = point.Norm(ref_point, norm_type_1_);
    double norm2 = point.Norm(ref_point, norm_type_2_);
    Timing::Instance()->EndTimer("DualNormScore::Slow");
    return {-norm1, -norm2};
  }

  ScoreType Fast(const ScoreType& old_score,
                 const Point& old_point,
                 const Point& ref_point,
                 const Patch& new_patch) const override {
    Timing::Instance()->StartTimer("DualNormScore::Fast");
    double norm1 = NormFast(-old_score.first, old_point, ref_point, new_patch,
                            norm_type_1_);
    double norm2 = NormFast(-old_score.second, old_point, ref_point, new_patch,
                            norm_type_2_);
    Timing::Instance()->EndTimer("DualNormScore::Fast");
    return {-norm1, -norm2};
  }

 private:
  int norm_type_1_;
  int norm_type_2_;
};

template <class CIter>
std::list<Patch> GeneratePatches(CIter& iter,
                                 const CIter& end,
                                 const std::vector<FeatureDir>& is_bounded,
                                 const BoundingBox* cached_intersection) {
  FeatureDir effective_dir = FeatureDir::None;
  while (iter != end) {
    if ((is_bounded[iter->first] & FeatureDir::Upper) &&
        iter->second.upper.has_value()) {
      if (!cached_intersection ||
          iter->second.HasSameUpper((*cached_intersection)[iter->first]))
        effective_dir |= FeatureDir::Upper;
    }

    if ((is_bounded[iter->first] & FeatureDir::Lower) &&
        iter->second.lower.has_value()) {
      if (!cached_intersection ||
          iter->second.HasSameLower((*cached_intersection)[iter->first]))
        effective_dir |= FeatureDir::Lower;
    }

    if (effective_dir != FeatureDir::None)
      break;

    ++iter;
  }

  if (iter == end)
    return {Patch()};

  int feature_id = iter->first;
  const Interval& inter = iter->second;
  std::vector<double> new_values;
  if (effective_dir & FeatureDir::Upper) {
    if (cached_intersection) {
      new_values.push_back(cached_intersection->Upper(feature_id) + eps);
    } else {
      new_values.push_back(inter.upper.value() + eps);
    }
  }
  if (effective_dir & FeatureDir::Lower) {
    if (cached_intersection) {
      new_values.push_back(cached_intersection->Lower(feature_id) - eps);
    } else {
      new_values.push_back(inter.lower.value() - eps);
    }
  }

  ++iter;
  std::list<Patch> patches =
      GeneratePatches(iter, end, is_bounded, cached_intersection);
  std::list<Patch> new_patches;

  for (double new_value : new_values) {
    for (const auto& p : patches) {
      Patch tmp(p);
      tmp[feature_id] = new_value;
      new_patches.emplace_back(std::move(tmp));
    }
  }
  new_patches.splice(new_patches.end(), patches);

  return std::move(new_patches);
}

void GenerateNaivePatches(SearchMode search_mode,
                          const Point& starting_point,
                          const BoundingBox* current_box,
                          const std::vector<FeatureDir>& is_bounded,
                          const BoundingBox* cached_intersection,
                          std::vector<Patch>* patches) {
  if (search_mode == SearchMode::NaiveLeaf) {
    auto leaves = current_box->OwnerTree()->GetLeaves();
    for (const auto* leaf : leaves) {
      patches->emplace_back(leaf->ClosestPatchTo(starting_point));
    }
    return;
  }

  assert(search_mode == SearchMode::NaiveFeature);
  for (const auto& iter : current_box->Intervals()) {
    int feature_id = iter.first;
    const auto& inter = iter.second;

    if ((is_bounded[feature_id] & FeatureDir::Upper) &&
        inter.HasSameUpper((*cached_intersection)[feature_id])) {
      patches->push_back({{feature_id, inter.upper.value() + eps}});
    } else if ((is_bounded[feature_id] & FeatureDir::Lower) &&
               inter.HasSameLower((*cached_intersection)[feature_id])) {
      patches->push_back({{feature_id, inter.lower.value() - eps}});
    }
  }
}

Patch GenRandomDir(int feature_size) {
  Patch dir;
  const int kNumDir = 5;
  while (dir.size() < kNumDir) {
    int feature_id = rand() % feature_size;
    int feature_dir = (rand() % 2) * 2 - 1;
    if (dir.find(feature_id) == dir.end())
      dir[feature_id] = feature_dir;
  }

  return std::move(dir);
}

Patch ComputeRecoveryPatch(const Point& point, const BoundingBox& box) {
  Patch p;
  for (const auto& iter : box.Intervals()) {
    p[iter.first] = point[iter.first];
  }
  return std::move(p);
}

Patch ComputeRecoveryPatch(const Point& point, const Direction& dir) {
  Patch p;
  for (auto d : dir) {
    p[abs(d)] = point[abs(d)];
  }
  return std::move(p);
}

bool HasSameLabel(int l1, int l2) {
  return l1 == l2;
}

void OptimizeLinearSearch(LayeredBoundingBox* layered_box,
                          const Point& victim_point,
                          int victim_label) {
  while (true) {
    Point adv_point = victim_point;

    const auto* box = layered_box->GetCachedIntersection();
    for (const auto& iter : box->Intervals()) {
      int feature_id = iter.first;
      if (iter.second.upper.has_value()) {
        adv_point[feature_id] =
            fmin(adv_point[feature_id], iter.second.upper.value() + eps);
      }

      if (iter.second.lower.has_value()) {
        adv_point[feature_id] =
            fmax(adv_point[feature_id], iter.second.lower.value() - eps);
      }
    }

    Point original_point = layered_box->Location();
    layered_box->ShiftPoint(adv_point);
    if (HasSameLabel(layered_box->PredictionLabel(), victim_label)) {
      layered_box->ShiftPoint(original_point);
      return;
    }
  }
}

std::pair<double, double> NormFirstScore(
    int victim_label,
    double norm,
    const std::vector<double>& label_scores) {
  int adv_index = MaxIndex(label_scores, victim_label);
  return std::make_pair(-norm,
                        label_scores[adv_index] - label_scores[victim_label]);
}

std::pair<double, double> LabelFirstScore(
    int victim_label,
    double norm,
    const std::vector<double>& label_scores) {
  int adv_index = MaxIndex(label_scores, victim_label);
  return std::make_pair(label_scores[adv_index] - label_scores[victim_label],
                        -norm);
}

std::pair<double, double> WeightedScore(
    int victim_label,
    double norm_weight,
    double norm,
    const std::vector<double>& label_scores) {
  int adv_index = MaxIndex(label_scores, victim_label);
  double label_score = label_scores[adv_index] - label_scores[victim_label];
  double score = norm_weight * label_score + label_score / (norm + 2.0);
  // double score = norm_weight * (-norm) + (1 - norm_weight) * label_score;
  return std::make_pair(score, 0);
}

bool True() {
  return true;
}

std::vector<const BoundingBox*> GetIncompatibleBoxes(
    const Patch& patch,
    const LayeredBoundingBox* layered_box,
    std::map<std::pair<int, double>, std::vector<const BoundingBox*>>*
        cached_incompatible_boxes_per_feature) {
  Timing::Instance()->StartTimer(
      "GetIncompatibleBoxes::cached_incompatible_boxes_per_feature");
  std::vector<const BoundingBox*> incompatible_boxes;
  for (const auto& feature_value : patch) {
    auto iter = cached_incompatible_boxes_per_feature->find(feature_value);
    if (iter != cached_incompatible_boxes_per_feature->end()) {
      incompatible_boxes.insert(incompatible_boxes.end(), iter->second.begin(),
                                iter->second.end());
    } else {
      std::vector<const BoundingBox*> boxes;
      layered_box->FillIncompatibleBoxes(feature_value.first,
                                         feature_value.second, &boxes);
      incompatible_boxes.insert(incompatible_boxes.end(), boxes.begin(),
                                boxes.end());
      cached_incompatible_boxes_per_feature->emplace(feature_value,
                                                     std::move(boxes));
    }
  }
  Timing::Instance()->EndTimer(
      "GetIncompatibleBoxes::cached_incompatible_boxes_per_feature");

  Timing::Instance()->StartTimer("GetIncompatibleBoxes::sort");

  std::sort(incompatible_boxes.begin(), incompatible_boxes.end());
  auto last = std::unique(incompatible_boxes.begin(), incompatible_boxes.end());
  incompatible_boxes.erase(last, incompatible_boxes.end());
  Timing::Instance()->EndTimer("GetIncompatibleBoxes::sort");

  Timing::Instance()->BinCount("incompatible_boxes.size()",
                               incompatible_boxes.size());

  return std::move(incompatible_boxes);
}

}  // namespace

const std::vector<int> NeighborAttack::kAllowedNormTypes = {-1, 1, 2};

NeighborAttack::NeighborAttack(const Config& config) : config_(config) {}

NeighborAttack::~NeighborAttack() {}

void NeighborAttack::LoadForestFromJson(const std::string& path) {
  assert(!forest_);
  forest_ = DecisionForest::CreateFromJson(
      path, config_.num_classes,
      config_.num_features + config_.feature_start - 1);

  // Also load train data if it's RBA-Appr (Yang et al. 2019).
  if (config_.search_mode == SearchMode::Region) {
    auto train_list =
        cz::LoadSVMFile(config_.train_path.c_str(), config_.num_features,
                        config_.feature_start);
    // Move to vector for better parallelize.
    train_data_.insert(train_data_.end(), train_list.begin(), train_list.end());
  }
}

NeighborAttack::Result NeighborAttack::FindAdversarialPoint(
    const Point& victim_point) const {

  int victim_label = forest_->PredictLabel(victim_point);
  printf("NeighborAttack::FindAdversarialPoint Spawning %d threads\n",
         config_.num_threads);
  boost::asio::thread_pool pool(config_.num_threads);

  int num_work = config_.num_attack_per_point;
  if (config_.search_mode == SearchMode::Region)
    num_work = config_.num_threads;

  Result results[num_work];
  for (int i = 0; i < num_work; ++i) {
    boost::asio::post(
        pool, boost::bind(&NeighborAttack::FindAdversarialPoint_ThreadRun, this,
                          i, victim_point, victim_label, &results[i]));
  }
  pool.join();

  Result best_result;
  for (int i = 0; i < num_work; ++i) {
    if (!results[i].success())
      continue;

    best_result.hist_points.push_back(
        results[i].best_points[config_.norm_type]);
    for (const int norm_type : kAllowedNormTypes) {
      double new_norm = results[i].best_norms[norm_type];
      if (new_norm < best_result.best_norms[norm_type]) {
        best_result.best_norms[norm_type] = new_norm;
        best_result.best_points[norm_type] = results[i].best_points[norm_type];
      }
    }
  }

  return std::move(best_result);
}

void NeighborAttack::FindAdversarialPoint_ThreadRun(int task_id,
                                                    const Point& victim_point,
                                                    int victim_label,
                                                    Result* result) const {
  printf(
      "NeighborAttack::FindAdversarialPoint_ThreadRun Trying %d/%d random "
      "starting points... thread_id:%s\n",
      task_id + 1, config_.num_attack_per_point,
      boost::lexical_cast<std::string>(boost::this_thread::get_id()).c_str());

  if (config_.search_mode == SearchMode::Region) {
    return RegionBasedAttackAppr_ThreadRun(task_id, victim_point, victim_label,
                                           result);
  }

  // TODO: Improve the quality of |GenFromVictim| points.
  bool use_random = true;
  Point p;

  Timing::Instance()->StartTimer("GenInitialPoint");
  if (use_random) {
    auto res = NormalRandomPoint(victim_label, victim_point, task_id);
    if (res.first) {
      p = res.second;
    } else {
      Timing::Instance()->EndTimer("GenInitialPoint");
      printf("  Failed to generate %d/%d random starting points... thread_id:%s\n",
        task_id + 1, config_.num_attack_per_point,
        boost::lexical_cast<std::string>(boost::this_thread::get_id()).c_str());
      return;
    }
    // p = FeatureSplitsRandomPoint(victim_point, victim_label, task_id);
    // p = PureRandomPoint(victim_point, victim_label);
  } else {
    assert(false);
    // double norm_weight =
    //     task_id * 1.0 / fmax(1, config_.num_attack_per_point - 1);
    // cout << "norm_weight:" << norm_weight << endl;
    // auto opt_p = GenFromVictim(victim_point, victim_label, norm_weight);
    // if (opt_p.has_value()) {
    //   p = std::move(opt_p.value());
    // } else {
    //   cout << "Fall back to regular random points..." << endl;
    //   p = PureRandomPoint(victim_point, victim_label, task_id);
    // }
  }
  Timing::Instance()->EndTimer("GenInitialPoint");

  int target_label = forest_->PredictLabel(p);

  printf("Initial point label:%d\n",
         forest_->PredictLabelBetween(p, target_label, victim_label));
  DCHECK(
      !HasSameLabel(forest_->PredictLabelBetween(p, target_label, victim_label),
                    victim_label));

  p = OptimizeAdversarialPoint(p, victim_point, target_label, victim_label,
                               result);
  DCHECK(
      !HasSameLabel(forest_->PredictLabelBetween(p, target_label, victim_label),
                    victim_label));

  printf("Norms for random point %d:%s\n", task_id + 1,
         GetNormStringForLogging(p, victim_point).c_str());

  UpdateResult(p, victim_point, result);
}

int NeighborAttack::PredictLabel(const Point& point) const {
  return forest_->PredictLabel(point);
}

int NeighborAttack::HammingDistanceBetween(const Point& p1, const Point& p2, int class1, int class2) const {
  return forest_->HammingDistanceBetween(p1, p2, class1, class2);
}

int NeighborAttack::NeighborDistanceBetween(const Point& start_point,
                                            const Point& end_point,
                                            int adv_class,
                                            int victim_class,
                                            const Point& victim_point) const {
  int num_trials = 200;

  boost::asio::thread_pool pool(config_.num_threads);
  int results[num_trials];
  for (int i = 0; i < num_trials; ++i) {
    boost::asio::post(
        pool, boost::bind(&NeighborAttack::NeighborDistanceBetween_ThreadRun,
                          this, i, start_point, end_point, adv_class,
                          victim_class, victim_point, &results[i]));
  }
  pool.join();

  int neighbor_dist = *std::min_element(results, results + num_trials);

  return neighbor_dist;
}

void NeighborAttack::NeighborDistanceBetween_ThreadRun(
    int task_id,
    const Point& start_point,
    const Point& end_point,
    int adv_class,
    int victim_class,
    const Point& victim_point,
    int* out_neighbor_dist) const {
  int norm_type1 = config_.norm_type;
  int norm_type2 = norm_type1 == 2 ? -1 : 2;

  auto start_scores = forest_->ComputeScores(start_point);
  auto end_scores = forest_->ComputeScores(end_point);
  // End point must be valid and better adv points.
  assert(end_scores[adv_class] > end_scores[victim_class]);

  ScoreType current_norm{victim_point.Norm(start_point, norm_type1),
                         victim_point.Norm(start_point, norm_type2)};

  auto end_box =
      forest_->GetLayeredBoundingBox(end_point, adv_class, victim_class);
  ScoreType end_norm{
      end_box->GetCachedIntersection()->NormTo(victim_point, norm_type1),
      end_box->GetCachedIntersection()->NormTo(victim_point, norm_type2)};

  if (current_norm <= end_norm) {
    // The |start_point| is already better.
    *out_neighbor_dist = 0;
    return;
  }

  // assert(victim_point.Norm(start_point, norm_type) >
  // victim_point.Norm(end_point, norm_type));

  // Positive: victim_class; Negative: adv_class.
  // |initial_label| may be negative in multi-class case.
  double initial_label = start_scores[victim_class] - start_scores[adv_class];

  // Adv boxes goes first.
  using TupleType = std::tuple<double, const BoundingBox*, const BoundingBox*>;
  std::vector<TupleType> negative_tuples;
  std::vector<TupleType> positive_tuples;

  BoundingBox shared_box;

  for (const auto& t : forest_->trees_) {
    if (t->ClassId() != adv_class && t->ClassId() != victim_class)
      continue;

    const auto* sb = t->GetBoundingBox(start_point);
    const auto* eb = t->GetBoundingBox(end_point);

    if (sb != eb) {
      double label_diff = eb->Label() - sb->Label();
      if (t->ClassId() == adv_class)
        label_diff = -label_diff;
      if (label_diff < 0) {
        negative_tuples.push_back({label_diff, sb, eb});
      } else {
        positive_tuples.push_back({label_diff, sb, eb});
      }
    } else {
      shared_box.Intersect(*sb);
    }
  }

  double current_label = initial_label;

  if (task_id == 0) {
    std::sort(negative_tuples.begin(), negative_tuples.end());
    std::sort(positive_tuples.begin(), positive_tuples.end());
  } else {
    auto rng = std::default_random_engine(task_id);
    std::shuffle(negative_tuples.begin(), negative_tuples.end(), rng);
    std::shuffle(positive_tuples.begin(), positive_tuples.end(), rng);
  }

  std::list<TupleType> changed_list;
  std::list<TupleType> remaining_list;
  remaining_list.insert(remaining_list.end(), negative_tuples.begin(),
                        negative_tuples.end());
  remaining_list.insert(remaining_list.end(), positive_tuples.begin(),
                        positive_tuples.end());

  int max_neighbor_dist = 0;
  while (!remaining_list.empty()) {
    int current_neighbor_dist = 0;
    bool should_add_first_node = true;
    bool added_new_node = false;
    BoundingBox current_box;

    while (should_add_first_node || current_label > 0 || added_new_node) {
      added_new_node = false;
      if (current_neighbor_dist == 0 || current_label > 0)
        assert(!remaining_list.empty());

      if (should_add_first_node || current_label > 0) {
        should_add_first_node = false;
        auto node = remaining_list.front();
        remaining_list.pop_front();
        assert(current_label < 0 || std::get<0>(node) < 0);
        current_label += std::get<0>(node);
        current_box.Intersect(*std::get<2>(node));
        added_new_node = true;
        ++current_neighbor_dist;
        changed_list.emplace_back(std::move(node));
      }

      auto iter = remaining_list.begin();
      while (iter != remaining_list.end()) {
        if (!current_box.Overlaps(*std::get<1>(*iter))) {
          auto node = *iter;
          iter = remaining_list.erase(iter);
          current_label += std::get<0>(node);
          assert(current_box.Overlaps(*std::get<2>(node)));
          current_box.Intersect(*std::get<2>(node));
          added_new_node = true;
          ++current_neighbor_dist;
          changed_list.emplace_back(std::move(node));
        } else {
          ++iter;
        }
      }

      if (!added_new_node) {
        auto merged_box = shared_box;
        for (const auto& node : remaining_list)
          merged_box.Intersect(*std::get<1>(node));
        for (const auto& node : changed_list)
          merged_box.Intersect(*std::get<2>(node));

        ScoreType new_norm{merged_box.NormTo(victim_point, norm_type1),
                           merged_box.NormTo(victim_point, norm_type2)};
        if ((new_norm >= current_norm || new_norm <= end_norm) &&
            !remaining_list.empty()) {
          should_add_first_node = true;
        }
      }
    }

    auto merged_box = shared_box;
    for (const auto& node : remaining_list)
      merged_box.Intersect(*std::get<1>(node));
    for (const auto& node : changed_list)
      merged_box.Intersect(*std::get<2>(node));
    ScoreType new_norm{merged_box.NormTo(victim_point, norm_type1),
                       merged_box.NormTo(victim_point, norm_type2)};
    assert(new_norm < current_norm && new_norm >= end_norm);
    current_norm = new_norm;

    max_neighbor_dist = std::max(max_neighbor_dist, current_neighbor_dist);
  }

  *out_neighbor_dist = max_neighbor_dist;
}

void NeighborAttack::RegionBasedAttackAppr_ThreadRun(int task_id,
                                                     const Point& victim_point,
                                                     int victim_label,
                                                     Result* result) const {
  int step = config_.num_threads;
  for (int i = task_id; i < train_data_.size(); i += step) {
    auto train = train_data_[i];
    // int y_train = train.first;
    const auto& x_train = train.second;
    // |y_pred| may be different from |y_train|.
    int y_pred = PredictLabel(x_train);
    if (y_pred == victim_label)
      continue;
    auto joint_box = forest_->GetBoundingBox(x_train);
    auto optimze_adv = OptimizeLocalSearch(&joint_box, victim_point);

    UpdateResult(optimze_adv, victim_point, result);
  }
}

const DecisionForest* NeighborAttack::ForestForTesting() const {
  return forest_.get();
}

Point NeighborAttack::OptimizeAdversarialPoint(Point adv_point,
                                               const Point& victim_point,
                                               int target_label,
                                               int victim_label,
                                               Result* result) const {
  adv_point =
      OptimizeBinarySearch(adv_point, victim_point, target_label, victim_label);

  auto layered_box =
      forest_->GetLayeredBoundingBox(adv_point, target_label, victim_label);
  // TODO: Could remove?
  adv_point =
      OptimizeLocalSearch(layered_box->GetCachedIntersection(), victim_point);
  layered_box->ShiftPoint(adv_point);

  // Timing::Instance()->StartTimer("OptimizeLinearSearch1");
  // OptimizeLinearSearch(layered_box.get(), victim_point, victim_label);
  // adv_point = layered_box->Location();
  // Timing::Instance()->EndTimer("OptimizeLinearSearch1");

  bool found_better_adv = true;
  auto best_norm = Norm(adv_point, victim_point);
  int i = 0;

  const int kMaxRandomFallback = 20;
  int used_random_fallback = 0;

  while (found_better_adv) {
    found_better_adv = false;
    if (config_.collect_histogram) {
      printf(
          "NeighborAttack::OptimizeAdversarialPoint iteration: %d best_norm: "
          "%lf\n   %s\n",
          i++, best_norm,
          GetNormStringForLogging(adv_point, victim_point).c_str());
    }

    Timing::Instance()->StartTimer("OptimizeNeighborSearch");
    found_better_adv =
        OptimizeNeighborSearch(layered_box.get(), victim_point, target_label,
                               victim_label, DualNormScore(config_.norm_type),
                               config_.search_mode, config_.max_dist);
    Timing::Instance()->EndTimer("OptimizeNeighborSearch");

    adv_point = layered_box->Location();
    DCHECK(!HasSameLabel(
        forest_->PredictLabelBetween(adv_point, target_label, victim_label),
        victim_label));

    // Add randomness to jump out of local minimum.
    if (!found_better_adv && used_random_fallback < kMaxRandomFallback) {
      ++used_random_fallback;
      auto res = OptimizeRandomSearch(adv_point, victim_point, victim_label);
      found_better_adv = res.first;

      if (found_better_adv) {
        target_label = forest_->PredictLabel(res.second);
        layered_box = forest_->GetLayeredBoundingBox(res.second, target_label,
                                                     victim_label);
        adv_point = OptimizeLocalSearch(layered_box->GetCachedIntersection(),
                                        victim_point);
        layered_box->ShiftPoint(adv_point);
      }
    }

    UpdateResult(adv_point, victim_point, result);

    auto new_norm = Norm(adv_point, victim_point);

    if (found_better_adv) {
      // |new_norm| could be larger than |best_norm| on non-inf norms due to
      // |OptimizeRandomSearch|.
      // assert(config_.norm_type != -1 || new_norm <= best_norm);
      best_norm = fmin(best_norm, new_norm);
    }
  }

  // printf("OptimizeAdversarialPoint::used_random_fallback:%d\n",
         // used_random_fallback);
  Timing::Instance()->BinCount("OptimizeAdversarialPoint::used_random_fallback",
                               used_random_fallback);
  Timing::Instance()->BinCount("OptimizeAdversarialPoint::Iterations", i);

  return std::move(adv_point);
}

Point NeighborAttack::OptimizeLocalSearch(const BoundingBox* box,
                                          const Point& victim_point) const {
  Timing::Instance()->StartTimer("NeighborAttack::OptimizeLocalSearch");

  Point adv_point = victim_point;

  for (const auto& iter : box->Intervals()) {
    int feature_id = iter.first;
    if (iter.second.upper.has_value()) {
      adv_point[feature_id] =
          fmin(adv_point[feature_id], iter.second.upper.value() - eps);
    }

    if (iter.second.lower.has_value()) {
      adv_point[feature_id] =
          fmax(adv_point[feature_id], iter.second.lower.value() + eps);
    }
  }

  Timing::Instance()->EndTimer("NeighborAttack::OptimizeLocalSearch");
  return std::move(adv_point);
}

Point NeighborAttack::OptimizeBinarySearch(const Point& adv_point,
                                           const Point& victim_point,
                                           int target_label,
                                           int victim_label) const {
  Timing::Instance()->StartTimer("NeighborAttack::OptimizeBinarySearch");
  if (config_.collect_histogram) {
    printf("Starting binary search, %s\n",
           GetNormStringForLogging(adv_point, victim_point).c_str());
  }

  int iterations = 0;

  Point small_diff = (adv_point - victim_point) / 100.0;
  Point upper = victim_point;
  while (forest_->PredictLabelBetween(upper, target_label, victim_label) ==
         victim_label) {
    ++iterations;
    upper += small_diff;
  }

  if (config_.collect_histogram) {
    printf("Finished linear search, iterations=%d %s\n", iterations,
           GetNormStringForLogging(adv_point, victim_point).c_str());
  }

  Point lower = victim_point;

  while ((upper - lower).Norm(-1) > config_.binary_search_threshold) {
    ++iterations;
    Point mid = (upper + lower) / 2;
    if (forest_->PredictLabelBetween(mid, target_label, victim_label) ==
        victim_label) {
      lower = std::move(mid);
    } else {
      upper = std::move(mid);
    }
  }

  if (config_.collect_histogram) {
    printf("Finished binary search, iterations=%d %s\n", iterations,
           GetNormStringForLogging(adv_point, victim_point).c_str());
  }
  Timing::Instance()->BinCount("BinarySearchIterations", iterations);

  Timing::Instance()->EndTimer("NeighborAttack::OptimizeBinarySearch");

  return upper;
}

bool NeighborAttack::OptimizeNeighborSearch(LayeredBoundingBox* layered_box,
                                            const Point& victim_point,
                                            int target_label,
                                            int victim_label,
                                            const NeighborScore& neighbor_score,
                                            SearchMode search_mode,
                                            int max_dist) const {
  Timing::Instance()->StartTimer("OptimizeNeighborSearch::Setup");
  const Point& starting_point = layered_box->Location();
  const auto& starting_score =
      neighbor_score.Slow(starting_point, victim_point);

  bool found_better_adv = false;
  Point best_adv;
  auto best_score = starting_score;
  L(Patch best_patch);
  Timing::Instance()->EndTimer("OptimizeNeighborSearch::Setup");

  Timing::Instance()->StartTimer("OptimizeNeighborSearch::bounded_features");
  const int kFeatureSize = victim_point.Size();
  std::vector<FeatureDir> is_bounded(kFeatureSize, FeatureDir::None);
  // <distance, feature_id>
  std::vector<std::pair<double, int>> bounded_features;

  // Compute |is_bounded| and |bounded_features|.
  for (const auto& iter : layered_box->GetCachedIntersection()->Intervals()) {
    int feature_id = iter.first;
    if (victim_point[feature_id] > starting_point[feature_id]) {
      is_bounded[feature_id] = FeatureDir::Upper;
      bounded_features.push_back(std::make_pair(
          std::abs(starting_point[feature_id] - victim_point[feature_id]),
          feature_id));
    } else if (victim_point[feature_id] < starting_point[feature_id]) {
      is_bounded[feature_id] = FeatureDir::Lower;
      bounded_features.push_back(std::make_pair(
          std::abs(starting_point[feature_id] - victim_point[feature_id]),
          feature_id));
    }
  }

  Timing::Instance()->BinCount("bounded_features.size()",
                               bounded_features.size());
  Timing::Instance()->EndTimer("OptimizeNeighborSearch::bounded_features");

  std::set<const BoundingBox*> visited_box;
  std::set<Patch, PatchCompare> visited_patch;

  int early_return_interval;
  int current_iteration = 0;
  if (config_.enable_early_return) {
    Timing::Instance()->StartTimer(
        "OptimizeNeighborSearch::kEnableEarlyReturn");
    early_return_interval = 1;
    std::sort(bounded_features.begin(), bounded_features.end(),
              std::greater<>());
    Timing::Instance()->EndTimer("OptimizeNeighborSearch::kEnableEarlyReturn");
  }

  std::map<std::pair<int, double>, std::vector<const BoundingBox*>>
      cached_incompatible_boxes_per_feature;

  for (const auto& dist_feature : bounded_features) {
    const int feature_id = dist_feature.second;

    Timing::Instance()->StartTimer(
        "OptimizeNeighborSearch::GetEffectiveBoxesForFeature");
    auto effective_boxes =
        layered_box->GetEffectiveBoxesForFeature(feature_id, search_mode);
    Timing::Instance()->EndTimer(
        "OptimizeNeighborSearch::GetEffectiveBoxesForFeature");

    std::vector<Patch> all_patches;
    for (const auto* box : effective_boxes) {
      Timing::Instance()->StartTimer("OptimizeNeighborSearch::visited_box");
      if (visited_box.count(box) > 0) {
        Timing::Instance()->EndTimer("OptimizeNeighborSearch::visited_box");
        continue;
      }
      visited_box.insert(box);
      Timing::Instance()->EndTimer("OptimizeNeighborSearch::visited_box");

      Timing::Instance()->StartTimer("OptimizeNeighborSearch::GeneratePatches");

      if (search_mode == SearchMode::ChangeOne) {
        Timing::Instance()->StartTimer("GetNeighborsWithinDistance");
        GetNeighborsWithinDistance(max_dist, feature_id, starting_score,
                                   starting_point, victim_point, box,
                                   layered_box, is_bounded, &all_patches);
        Timing::Instance()->EndTimer("GetNeighborsWithinDistance");
      } else {
        assert(search_mode == SearchMode::NaiveFeature ||
               search_mode == SearchMode::NaiveLeaf);
        GenerateNaivePatches(search_mode, starting_point, box, is_bounded,
                             layered_box->GetCachedIntersection(),
                             &all_patches);
        // DCHECK(patches.back().empty());
        // patches.pop_back();
      }

      Timing::Instance()->EndTimer("OptimizeNeighborSearch::GeneratePatches");
    }

    Timing::Instance()->BinCount("neighbor1t-all_patches.size",
                                 all_patches.size());

    for (const auto& patch : all_patches) {
      Timing::Instance()->BinCount("BeforeHashPatchSize", patch.size());

      Timing::Instance()->StartTimer("OptimizeNeighborSearch::visited_patch");
      if (visited_patch.count(patch) > 0) {
        Timing::Instance()->EndTimer("OptimizeNeighborSearch::visited_patch");
        continue;
      }
      visited_patch.insert(patch);
      Timing::Instance()->EndTimer("OptimizeNeighborSearch::visited_patch");

      Timing::Instance()->BinCount("AfterHashPatchSize", patch.size());

      Timing::Instance()->StartTimer(
          "OptimizeNeighborSearch::StretchWithinBox");
      const BoundingBox* constrain_box = nullptr;
      if (search_mode == SearchMode::ChangeOne) {
        constrain_box = patch.box;
      }

      Timing::Instance()->StartTimer(
          "OptimizeNeighborSearch::incompatible_boxes");
      std::vector<const BoundingBox*> incompatible_boxes = GetIncompatibleBoxes(
          patch, layered_box, &cached_incompatible_boxes_per_feature);
      Timing::Instance()->EndTimer(
          "OptimizeNeighborSearch::incompatible_boxes");

      Timing::Instance()->StartTimer("OptimizeNeighborSearch::new_boxes");
      std::vector<const BoundingBox*> new_boxes =
          layered_box->GetNewBoxes(patch, incompatible_boxes);
      Timing::Instance()->EndTimer("OptimizeNeighborSearch::new_boxes");

      Timing::Instance()->StartTimer("OptimizeNeighborSearch::new_adv_scores");
      std::vector<double> new_adv_scores =
          layered_box->GetNewScores(incompatible_boxes, new_boxes);
      Timing::Instance()->EndTimer("OptimizeNeighborSearch::new_adv_scores");

      Timing::Instance()->StartTimer("OptimizeNeighborSearch::MaxIndex");
      int new_label =
          MaxIndexBetween(new_adv_scores, target_label, victim_label);

      Timing::Instance()->EndTimer("OptimizeNeighborSearch::MaxIndex");

      if (new_label != target_label) {
        Timing::Instance()->BinCount("Filter_Label", 1);
        continue;
      }
      Timing::Instance()->BinCount("Filter_Label", 0);

      // Note: |new_adv_patch| may not have the same bounding box as |patch|
      // before |TightenPoint|.
      Patch new_adv_patch = layered_box->StretchWithinBox(
          patch, victim_point, constrain_box, incompatible_boxes);
      Timing::Instance()->EndTimer("OptimizeNeighborSearch::StretchWithinBox");

      // Used stretched |new_adv| and |layered_box->Scores()| to do a quick
      // filter.
      // TODO: Verify if it's ok to use |layered_box->Scores()|.
      if (neighbor_score.Fast(starting_score, starting_point, victim_point,
                              new_adv_patch) < best_score) {
        Timing::Instance()->BinCount("Filter_StretchNorm", 1);

        Timing::Instance()->StartTimer(
            "OptimizeNeighborSearch::incompatible_boxes.clear()");
        incompatible_boxes.clear();
        Timing::Instance()->EndTimer(
            "OptimizeNeighborSearch::incompatible_boxes.clear()");
        continue;
      }
      Timing::Instance()->BinCount("Filter_StretchNorm", 0);

      Timing::Instance()->StartTimer("OptimizeNeighborSearch::new_adv");
      Point new_adv(starting_point);
      new_adv.Apply(new_adv_patch);
      Timing::Instance()->EndTimer("OptimizeNeighborSearch::new_adv");

      Timing::Instance()->StartTimer("OptimizeNeighborSearch::TightenPoint");
      layered_box->TightenPoint(&new_adv, new_boxes);
      Timing::Instance()->EndTimer("OptimizeNeighborSearch::TightenPoint");

      Timing::Instance()->StartTimer(
          "OptimizeNeighborSearch::incompatible_boxes.clear()");
      incompatible_boxes.clear();
      Timing::Instance()->EndTimer(
          "OptimizeNeighborSearch::incompatible_boxes.clear()");

      // TODO: Add back hashing.
      // if (visted_boxes_.count(layered_box->Hash()) > 0)
      //   continue;
      // visted_boxes_.insert(layered_box->Hash());

      // Timing::Instance()->IncreaseSample("VisitsPerBox",
      // layered_box->Hash());

      // TODO: Avoid computing full norm each time.
      Timing::Instance()->StartTimer("OptimizeNeighborSearch::score_func");
      auto new_score =
          neighbor_score.Slow(new_adv, victim_point);
      Timing::Instance()->EndTimer("OptimizeNeighborSearch::score_func");

      Timing::Instance()->StartTimer("OptimizeNeighborSearch::SaveBetter");
      if (new_score > best_score) {
        Timing::Instance()->BinCount("Filter_Score", 0);
        found_better_adv = true;
        best_adv = new_adv;
        best_score = new_score;
        L(best_patch = patch);
      } else {
        Timing::Instance()->BinCount("Filter_Score", 1);
      }
      Timing::Instance()->EndTimer("OptimizeNeighborSearch::SaveBetter");
    }

    if (config_.enable_early_return) {
      ++current_iteration;
      if (found_better_adv && current_iteration % early_return_interval == 0)
        break;
    }
  }

  if (config_.enable_early_return) {
    Timing::Instance()->BinCount(
        "EarlyReturnBatch", (int)(current_iteration / early_return_interval));
  }

  Timing::Instance()->BinCount("bound_trees-visited_box.size()",
                               visited_box.size());
  Timing::Instance()->BinCount("neighbor1-visited_patch.size()",
                               visited_patch.size());

  Timing::Instance()->StartTimer(
      "OptimizeNeighborSearch::visited_patch.clear()");
  visited_patch.clear();
  Timing::Instance()->EndTimer("OptimizeNeighborSearch::visited_patch.clear()");

  Timing::Instance()->StartTimer(
      "OptimizeNeighborSearch::cached_incompatible_boxes_per_feature.clear()");
  cached_incompatible_boxes_per_feature.clear();
  Timing::Instance()->EndTimer(
      "OptimizeNeighborSearch::cached_incompatible_boxes_per_feature.clear()");

  L(cout << "best_patch:" << ToDebugString(best_patch) << endl);

  Timing::Instance()->StartTimer("OptimizeNeighborSearch::ShiftPoint");
  if (found_better_adv)
    layered_box->ShiftPoint(best_adv);
  Timing::Instance()->EndTimer("OptimizeNeighborSearch::ShiftPoint");
  return found_better_adv;
}

std::pair<bool, Point> NeighborAttack::OptimizeRandomSearch(
    const Point& adv_point,
    const Point& victim_point,
    int victim_label) const {
  Timing::Instance()->StartTimer("OptimizeRandomSearch");

  double l_inf = adv_point.Norm(victim_point, -1);
  std::default_random_engine generator;
  std::normal_distribution<double> distribution(0.0, 2 * l_inf);
  std::bernoulli_distribution should_change(0.1);

  bool found_better_adv = false;
  Point best_adv = adv_point;

  const int len = adv_point.Size();
  Point p_min(len), p_max(len);
  double l_inf_box = fabs(l_inf - 3 * eps);
  for (int i = 0; i < len; ++i) {
    p_min[i] = victim_point[i] - l_inf_box;
    p_max[i] = victim_point[i] + l_inf_box;
  }

  int ct_used_trials = 0;

  Timing::Instance()->StartTimer("OptimizeRandomSearch::trial_loop");
  const int kMaxTrial = 300;
  Point p(len);
  while (ct_used_trials < kMaxTrial) {
    ++ct_used_trials;
    for (int i = 0; i < len; ++i) {
      p[i] = Clip(adv_point[i] +
                      distribution(generator) * int(should_change(generator)),
                  p_min[i], p_max[i]);
    }

    Timing::Instance()->StartTimer(
        "OptimizeRandomSearch::trial_loop::PredictLabel");
    // |p| is guaranteed to have smaller norm on l-inf. We have
    // |kMaxRandomFallback| to guard other norms.
    if (!HasSameLabel(forest_->PredictLabel(p), victim_label)) {
      best_adv = p;
      found_better_adv = true;
      break;
    }
    Timing::Instance()->EndTimer(
        "OptimizeRandomSearch::trial_loop::PredictLabel");
  }
  Timing::Instance()->EndTimer("OptimizeRandomSearch::trial_loop");

  Timing::Instance()->BinCount("OptimizeRandomSearch::ct_used_trials",
                               ct_used_trials);

  Timing::Instance()->EndTimer("OptimizeRandomSearch");
  return std::make_pair(found_better_adv, best_adv);
}

std::pair<bool, Point> NeighborAttack::ReverseNeighborSearch(
    const Point& initial_point,
    int target_label,
    int victim_label,
    const BoundingBox& p_constrain) const {
  Timing::Instance()->StartTimer("ReverseNeighborSearch");
  Timing::Instance()->StartTimer("ReverseNeighborSearch::Setup");

  Timing::Instance()->StartTimer(
      "ReverseNeighborSearch::Setup::GetLayeredBoundingBox");
  auto layered_box =
      forest_->GetLayeredBoundingBox(initial_point, target_label, victim_label);
  Timing::Instance()->EndTimer(
      "ReverseNeighborSearch::Setup::GetLayeredBoundingBox");

  Timing::Instance()->StartTimer(
      "ReverseNeighborSearch::Setup::GetBoxForAllTree");
  auto all_box = layered_box->GetBoxForAllTree();
  Timing::Instance()->EndTimer(
      "ReverseNeighborSearch::Setup::GetBoxForAllTree");
  assert(all_box.size() > 0);
  std::default_random_engine generator;

  Timing::Instance()->EndTimer("ReverseNeighborSearch::Setup");

  const int kMaxAllBoxIter = 1;
  bool found_adv = false;
  int all_box_iter = 0;
  int per_box_iter = 0;
  while (all_box_iter < kMaxAllBoxIter) {
    ++all_box_iter;
    bool found_better_box = false;

    Timing::Instance()->StartTimer("ReverseNeighborSearch::shuffle");
    std::shuffle(all_box.begin(), all_box.end(), generator);
    Timing::Instance()->EndTimer("ReverseNeighborSearch::shuffle");

    int len = std::min(100, (int)all_box.size());
    for (int i = 0; i < len; ++i) {
      const auto* box_to_replace = all_box[i];
      if (found_adv)
        break;
      ++per_box_iter;
      double old_label =
          (box_to_replace->OwnerTree()->ClassId() == target_label)
              ? -box_to_replace->Label()
              : box_to_replace->Label();

      Timing::Instance()->StartTimer(
          "ReverseNeighborSearch::GetAlternativeBoxes");
      auto alt_boxes = layered_box->GetAlternativeBoxes(
          BoundingBox(), 1, box_to_replace, config_.enable_relaxed_boundary,
          &p_constrain);
      Timing::Instance()->BinCount("ReverseNeighborSearch::alt_boxes.size()",
                                   alt_boxes.size());
      Timing::Instance()->EndTimer(
          "ReverseNeighborSearch::GetAlternativeBoxes");

      double min_label = old_label;
      const BoundingBox* min_box = nullptr;
      for (const auto* new_box : alt_boxes) {
        double new_label = (new_box->OwnerTree()->ClassId() == target_label)
                               ? -new_box->Label()
                               : new_box->Label();
        if (new_label < min_label) {
          min_label = new_label;
          min_box = new_box;
        }
      }

      if (min_box) {
        found_better_box = true;
        Timing::Instance()->StartTimer("ReverseNeighborSearch::ShiftByPatch");
        auto patch = min_box->ClosestPatchTo(layered_box->Location());
        layered_box->ShiftByPatch(patch);
        Timing::Instance()->EndTimer("ReverseNeighborSearch::ShiftByPatch");
        if (layered_box->PredictionLabel() == target_label) {
          found_adv = true;
          break;
        }
      }
    }

    if (!found_better_box)
      break;
  }

  Timing::Instance()->BinCount("ReverseNeighborSearch::all_box_iter",
                               all_box_iter);
  Timing::Instance()->BinCount("ReverseNeighborSearch::per_box_iter",
                               per_box_iter);

  Timing::Instance()->EndTimer("ReverseNeighborSearch");

  return std::make_pair(found_adv, layered_box->Location());
}

void NeighborAttack::GetNeighborsWithinDistance(
    int max_dist,
    int target_feature_id,
    const ScoreType& starting_score,
    const Point& starting_point,
    const Point& victim_point,
    const BoundingBox* box_to_replace,
    const LayeredBoundingBox* layered_box,
    const std::vector<FeatureDir>& is_bounded,
    std::vector<Patch>* candidate_patches) const {
  Timing::Instance()->StartTimer(
      "GetNeighborsWithinDistance::GetAlternativeBoxes");
  BoundingBox target_feature_constrain;
  Interval inter;
  if (is_bounded[target_feature_id] == FeatureDir::Upper) {
    inter.lower = starting_point[target_feature_id] + 2 * eps;
  } else {
    inter.upper = starting_point[target_feature_id] - 2 * eps;
  }
  target_feature_constrain.IntersectFeature(target_feature_id, inter);
  auto alt_boxes = layered_box->GetAlternativeBoxes(
      target_feature_constrain, max_dist, box_to_replace,
      config_.enable_relaxed_boundary);

  Timing::Instance()->BinCount("alt_boxes.size()", alt_boxes.size());
  Timing::Instance()->EndTimer(
      "GetNeighborsWithinDistance::GetAlternativeBoxes");

  Timing::Instance()->StartTimer(
      "GetNeighborsWithinDistance::candidate_patches");

  for (const auto* box : alt_boxes) {
    Patch alt_patch;
    alt_patch.box = box;
    for (const auto& iter : box->Intervals()) {
      int feature_id = iter.first;
      const Interval& inter = iter.second;
      if (!inter.Contains(starting_point[feature_id]))
        alt_patch[feature_id] = inter.ClosestTo(starting_point[feature_id]);
    }

    if (!alt_patch.empty())
      candidate_patches->emplace_back(std::move(alt_patch));
  }
  Timing::Instance()->EndTimer("GetNeighborsWithinDistance::candidate_patches");
}

double NeighborAttack::Norm(const Point& p1, const Point& p2) const {
  Timing::Instance()->StartTimer("NeighborAttack::Norm");
  double n = p1.Norm(p2, config_.norm_type);
  Timing::Instance()->EndTimer("NeighborAttack::Norm");
  return n;
}

double NeighborAttack::NormFast(const Patch& new_patch,
                                double old_norm,
                                const Point& old_point,
                                const Point& ref_point) const {
  return cz::NormFast(old_norm, old_point, ref_point, new_patch,
                      config_.norm_type);
}

std::pair<bool, Point> NeighborAttack::PureRandomPoint(const Point& ref_point,
                                                       int victim_label,
                                                       int rseed) const {
  Timing::Instance()->StartTimer("PureRandomPoint");

  std::default_random_engine generator(rseed);
  std::uniform_real_distribution<double> distribution(-0.5, 1.5);

  Point p(ref_point.Size());

  int max_try = 3000;
  while (max_try > 0) {
    --max_try;
    for (auto& v : p) {
      v = distribution(generator);
    }

    if (!HasSameLabel(forest_->PredictLabel(p), victim_label)) {
      Timing::Instance()->EndTimer("PureRandomPoint");
      return std::make_pair(true, p);
    }
  }

  return std::make_pair(false, Point());
}

std::pair<bool, Point> NeighborAttack::NormalRandomPoint(int victim_label,
                                                         const Point& ref_point,
                                                         int rseed) const {
  Timing::Instance()->StartTimer("NormalRandomPoint");

  std::default_random_engine generator(rseed);
  std::normal_distribution<double> distribution(0.0, 1.0);

  const int len = ref_point.Size();
  Point p(len);

  int max_try = 300;
  while (max_try > 0) {
    --max_try;
    for (int i = 0; i < len; ++i) {
      p[i] = ref_point[i] + distribution(generator);
    }

    if (!HasSameLabel(forest_->PredictLabel(p), victim_label)) {
      Timing::Instance()->EndTimer("NormalRandomPoint");
      return std::make_pair(true, p);
    }
  }
  Timing::Instance()->EndTimer("NormalRandomPoint");

  printf(
      "NormalRandomPoint failed after 300 trials for victim_label: %d. "
      "fallback to FeatureSplitsRandomPoint.\n",
      victim_label);
  return FeatureSplitsRandomPoint(victim_label, ref_point, rseed);
}

std::pair<bool, Point> NeighborAttack::FeatureSplitsRandomPoint(
    int victim_label,
    const Point& ref_point,
    int rseed) const {
  Timing::Instance()->StartTimer("FeatureSplitsRandomPoint");

  std::default_random_engine generator(rseed);
  std::uniform_int_distribution<int> distribution(0, 30000007);

  const auto& feature_splits = forest_->FeatureSplits();
  const int len = ref_point.Size();
  Point p(len);

  int max_try = 10000;
  while (max_try > 0) {
    --max_try;
    for (int i = 0; i < len; ++i) {
      p[i] = feature_splits[i][distribution(generator) %
                               feature_splits[i].size()] +
             eps;
    }

    if (!HasSameLabel(forest_->PredictLabel(p), victim_label)) {
      Timing::Instance()->EndTimer("FeatureSplitsRandomPoint");
      return std::make_pair(true, p);
    }
  }

  Timing::Instance()->EndTimer("FeatureSplitsRandomPoint");

  return std::make_pair(false, Point());
}

bool NeighborAttack::FlipWithDirection(LayeredBoundingBox* layered_box,
                                       const Patch& dir) const {
  Patch p;
  Timing::Instance()->StartTimer("FlipWithDirection::Loops");
  const auto& intervals = layered_box->GetCachedIntersection()->Intervals();

  for (const auto& iter : dir) {
    int feature_id = iter.first;
    int feature_dir = iter.second;

    const auto& inter = intervals.find(feature_id);
    if (inter == intervals.end())
      continue;

    if (feature_dir < 0 && inter->second.lower.has_value()) {
      p[feature_id] = inter->second.lower.value() - eps;
    } else if (feature_dir > 0 && inter->second.upper.has_value()) {
      p[feature_id] = inter->second.upper.value() + eps;
    }
  }
  Timing::Instance()->EndTimer("FlipWithDirection::Loops");

  Timing::Instance()->StartTimer("FlipWithDirection::ShiftByPatch");
  if (!p.empty())
    layered_box->ShiftByPatch(p);
  Timing::Instance()->EndTimer("FlipWithDirection::ShiftByPatch");
  return !p.empty();
}

void NeighborAttack::UpdateResult(const Point& adv_point,
                                  const Point& victim_point,
                                  Result* result) const {
  Timing::Instance()->StartTimer("NeighborAttack::UpdateResult");

  auto p = (adv_point - victim_point);

  for (const int norm_type : kAllowedNormTypes) {
    double new_norm = p.Norm(norm_type);
    if (new_norm < result->best_norms[norm_type]) {
      result->best_norms[norm_type] = new_norm;
      result->best_points[norm_type] = adv_point;
    }
  }

  Timing::Instance()->EndTimer("NeighborAttack::UpdateResult");
}

std::string NeighborAttack::GetNormStringForLogging(
    const Point& adv_point,
    const Point& victim_point) const {
  auto p = (adv_point - victim_point);

  std::string s;
  for (const int norm_type : kAllowedNormTypes) {
    s += " Norm(" + std::to_string(norm_type) +
         ")=" + std::to_string(p.Norm(norm_type)) + " ";
  }
  return s;
}

}  // namespace cz


// Copyright (c) 2018 yshurik
//
// Distributed under the MIT/X11 software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.
//
// The use in another cyptocurrency project the code is licensed under
// Jelurida Public License (JPL). See https://www.jelurida.com/resources/jpl

#include "pegdata.h"

#include <map>
#include <set>
#include <cstdint>
#include <utility>
#include <algorithm>
#include <type_traits>

#include <boost/multiprecision/cpp_int.hpp>

#include "util.h"

#include <zconf.h>
#include <zlib.h>

using namespace std;
using namespace boost;

CFractions::CFractions()
    :nFlags(VALUE)
    ,f(new int64_t[PEG_SIZE])
{
    f[0] = 0; // fast init first item
}
CFractions::CFractions(int64_t value, uint32_t flags)
    :nFlags(flags)
    ,f(new int64_t[PEG_SIZE])
{
    if (flags & VALUE)
        f[0] = value; // fast init first item
    else if (flags & STD) {
        f[0] = value;
        nFlags = VALUE;
        ToStd();
        nFlags = flags;
    }
    else {
        assert(0);
    }
}
CFractions::CFractions(const CFractions & o)
    :nFlags(o.nFlags)
    ,nLockTime(o.nLockTime)
    ,sReturnAddr(o.sReturnAddr)
    ,f(new int64_t[PEG_SIZE])
{
    for(int i=0; i< PEG_SIZE; i++) {
        f[i] = o.f[i];
    }
}

CFractions& CFractions::operator=(const CFractions& o)
{
    nFlags = o.nFlags;
    nLockTime = o.nLockTime;
    sReturnAddr = o.sReturnAddr;
    for(int i=0; i< PEG_SIZE; i++) {
        f[i] = o.f[i];
    }
    return *this;
}

void CFractions::ToDeltas(int64_t* deltas) const
{
    int64_t fp = 0;
    for(int i=0; i<PEG_SIZE; i++) {
        if (i==0) {
            fp = deltas[0] = f[0];
            continue;
        }
        deltas[i] = f[i]-fp*(PEG_RATE-1)/PEG_RATE;
        fp = f[i];
    }
}

void CFractions::FromDeltas(const int64_t* deltas)
{
    int64_t fp = 0;
    for(int i=0; i<PEG_SIZE; i++) {
        if (i==0) {
            fp = f[0] = deltas[0];
            continue;
        }
        f[i] = deltas[i]+fp*(PEG_RATE-1)/PEG_RATE;
        fp = f[i];
    }
}

bool CFractions::Pack(CDataStream& out, unsigned long* report_len, bool compress) const
{
    if (nFlags & VALUE) {
        if (report_len) *report_len = sizeof(int64_t);
        out << nVersion;
        out << uint32_t(nFlags | SER_VALUE);
        out << nLockTime;
        out << sReturnAddr;
        out << f[0];
    } else if (compress) {
        int64_t deltas[PEG_SIZE];
        ToDeltas(deltas);

        int zlevel = 9;
        unsigned char zout[2*PEG_SIZE*sizeof(int64_t)];
        unsigned long n = PEG_SIZE*sizeof(int64_t);
        unsigned long zlen = PEG_SIZE*2*sizeof(int64_t);
        auto src = reinterpret_cast<const unsigned char *>(deltas);
        int res = ::compress2(zout, &zlen, src, n, zlevel);
        if (res == Z_OK) {
            if (report_len) *report_len = zlen;
            out << nVersion;
            out << uint32_t(nFlags | SER_ZDELTA);
            out << nLockTime;
            out << sReturnAddr;
            auto ser = reinterpret_cast<const char *>(zout);
            out << zlen;
            out.write(ser, zlen);
        }
        else {
            if (report_len) *report_len = PEG_SIZE*sizeof(int64_t);
            out << nVersion;
            out << uint32_t(nFlags | SER_RAW);
            out << nLockTime;
            out << sReturnAddr;
            auto ser = reinterpret_cast<const char *>(f.get());
            out.write(ser, PEG_SIZE*sizeof(int64_t));
        }
    } else {
        if (report_len) *report_len = PEG_SIZE*sizeof(int64_t);
        out << nVersion;
        out << uint32_t(nFlags | SER_RAW);
        out << nLockTime;
        out << sReturnAddr;
        auto ser = reinterpret_cast<const char *>(f.get());
        out.write(ser, PEG_SIZE*sizeof(int64_t));
    }
    return true;
}

bool CFractions::Unpack(CDataStream& inp)
{
    uint32_t nSerFlags = 0;
    inp >> nVersion;
    inp >> nSerFlags;
    inp >> nLockTime;
    inp >> sReturnAddr;

    if (nSerFlags & SER_VALUE) {
        nFlags = nSerFlags | VALUE;
        inp >> f[0];
    }
    else if (nSerFlags & SER_ZDELTA) {
        unsigned long zlen = 0;
        inp >> zlen;

        if (zlen>(2*PEG_SIZE*sizeof(int64_t))) {
            // data are broken, no read
            return false;
        }

        unsigned char zinp[2*PEG_SIZE*sizeof(int64_t)];
        unsigned long n = PEG_SIZE*sizeof(int64_t);
        auto ser = reinterpret_cast<char *>(zinp);
        inp.read(ser, zlen);

        int64_t deltas[PEG_SIZE];
        auto src = reinterpret_cast<const unsigned char *>(ser);
        auto dst = reinterpret_cast<unsigned char *>(deltas);
        int res = ::uncompress(dst, &n, src, zlen);
        if (res != Z_OK) {
            // data are broken, can not uncompress
            return false;
        }
        FromDeltas(deltas);
        nFlags = nSerFlags | STD;
    }
    else if (nSerFlags & SER_RAW) {
        auto ser = reinterpret_cast<char *>(f.get());
        inp.read(ser, PEG_SIZE*sizeof(int64_t));
        nFlags = nSerFlags | STD;
    }
    nFlags &= SER_MASK;

    return true;
}

CFractions CFractions::Std() const
{
    if ((nFlags & VALUE) ==0)
        return *this;

    CFractions fstd;
    fstd.sReturnAddr = sReturnAddr;
    fstd.nLockTime = nLockTime;
    fstd.nFlags = nFlags;
    fstd.nFlags &= ~uint32_t(VALUE);
    fstd.nFlags |= STD;

    int64_t v = f[0];
    for(int i=0;i<PEG_SIZE;i++) {
        if (i == PEG_SIZE-1) {
            fstd.f[i] = v;
            break;
        }
        int64_t frac = v/PEG_RATE;
        fstd.f[i] = frac;
        v -= frac;
    }
    return fstd;
}

bool CFractions::IsPositive() const
{
    if (nFlags & VALUE)
        return true;

    for(int i=0;i<PEG_SIZE;i++) {
        if (f[i] <0)
            return false;
    }
    return true;
}

bool CFractions::IsNegative() const
{
    if (nFlags & VALUE)
        return false;

    for(int i=0;i<PEG_SIZE;i++) {
        if (f[i] >0)
            return false;
    }
    return true;
}

int64_t CFractions::Total() const
{
    int64_t nValue =0;
    if (nFlags & VALUE)
        return f[0];

    for(int i=0;i<PEG_SIZE;i++) {
        nValue += f[i];
    }
    return nValue;
}

int64_t CFractions::Low(int supply) const
{
    int64_t nValue =0;
    if (nFlags & VALUE)
        return Std().Low(supply);

    for(int i=0;i<supply;i++) {
        nValue += f[i];
    }
    return nValue;
}

int64_t CFractions::High(int supply) const
{
    int64_t nValue =0;
    if (nFlags & VALUE)
        return Std().High(supply);

    for(int i=supply;i<PEG_SIZE;i++) {
        nValue += f[i];
    }
    return nValue;
}

int64_t CFractions::Low(const CPegLevel & peglevel) const
{
    int64_t nValue =0;
    if (nFlags & VALUE)
        return Std().Low(peglevel);

    int to = peglevel.nSupply + peglevel.nShift;
    if (to <0) return 0;
    if (to >= PEG_SIZE) return 0;

    if (peglevel.nShiftLastPart >0 &&
        peglevel.nShiftLastTotal >0) {
        // partial value to use
        int64_t v = f[to];
        int64_t vpart = ::RatioPart(v,
                                    peglevel.nShiftLastPart,
                                    peglevel.nShiftLastTotal);
        if (vpart < v) vpart++; // better rounding
        nValue += vpart;
    }

    for(int i=0;i<to;i++) {
        nValue += f[i];
    }

    return nValue;
}

int64_t CFractions::High(const CPegLevel & peglevel) const
{
    int64_t nValue =0;
    if (nFlags & VALUE)
        return Std().High(peglevel);

    int from = peglevel.nSupply + peglevel.nShift;
    if (from <0) return 0;
    if (from >= PEG_SIZE) return 0;

    if (peglevel.nShiftLastPart >0 &&
        peglevel.nShiftLastTotal >0) {
        // partial value to use
        int64_t v = f[from];
        int64_t vpart = ::RatioPart(v,
                                    peglevel.nShiftLastPart,
                                    peglevel.nShiftLastTotal);
        if (vpart < v) vpart++; // better rounding
        nValue += (v - vpart);
        from++;
    }

    for(int i=from;i<PEG_SIZE;i++) {
        nValue += f[i];
    }

    return nValue;
}

int64_t CFractions::NChange(const CPegLevel & peglevel) const
{
    CPegLevel peglevel_next = peglevel;
    peglevel_next.nSupply = peglevel_next.nSupplyNext;
    peglevel_next.nSupplyNext = peglevel_next.nSupplyNextNext;

    int64_t nValueSrc = High(peglevel);
    int64_t nValueDst = High(peglevel_next);
    return nValueDst - nValueSrc;
}

int64_t CFractions::NChange(int src_supply, int dst_supply) const
{
    int64_t nValueSrc = High(src_supply);
    int64_t nValueDst = High(dst_supply);
    return nValueDst - nValueSrc;
}

int16_t CFractions::HLI() const
{
    if (nFlags & VALUE)
        return Std().HLI();

    int64_t half = 0;
    int64_t total = Total();
    if (total ==0) {
        return 0;
    }
    for(int16_t i=0;i<PEG_SIZE;i++) {
        half += f[i];
        if (half > total/2) {
            return i;
        }
    }
    return 0;
}

void CFractions::ToStd()
{
    if ((nFlags & VALUE) == 0)
        return;

    nFlags &= ~uint32_t(VALUE);
    nFlags |= STD;

    int64_t v = f[0];
    for(int i=0;i<PEG_SIZE;i++) {
        if (i == PEG_SIZE-1) {
            f[i] = v;
            break;
        }
        int64_t frac = v/PEG_RATE;
        f[i] = frac;
        v -= frac;
    }
}

CFractions CFractions::Positive(int64_t* total) const
{
    if ((nFlags & STD) == 0) {
        return Std().Positive(total);
    }
    CFractions frPositive(0, CFractions::STD);
    for(int i=0; i<PEG_SIZE; i++) {
        if (f[i] <=0) continue;
        frPositive.f[i] = f[i];
        if (total) *total += f[i];
    }
    return frPositive;
}
CFractions CFractions::Negative(int64_t* total) const
{
    if ((nFlags & STD) == 0) {
        return Std().Negative(total);
    }
    CFractions frNegative(0, CFractions::STD);
    for(int i=0; i<PEG_SIZE; i++) {
        if (f[i] >=0) continue;
        frNegative.f[i] = f[i];
        if (total) *total += f[i];
    }
    return frNegative;
}

CFractions CFractions::LowPart(int supply, int64_t* total) const
{
    if ((nFlags & STD) == 0) {
        return Std().LowPart(supply, total);
    }
    CFractions frLowPart(0, CFractions::STD);
    for(int i=0; i<supply; i++) {
        if (total) *total += f[i];
        frLowPart.f[i] += f[i];
    }
    return frLowPart;
}
CFractions CFractions::HighPart(int supply, int64_t* total) const
{
    if ((nFlags & STD) == 0) {
        return Std().HighPart(supply, total);
    }
    CFractions frHighPart(0, CFractions::STD);
    for(int i=supply; i<PEG_SIZE; i++) {
        if (total) *total += f[i];
        frHighPart.f[i] += f[i];
    }
    return frHighPart;
}

CFractions CFractions::LowPart(const CPegLevel & peglevel, int64_t* total) const
{
    if ((nFlags & STD) == 0) {
        return Std().LowPart(peglevel, total);
    }

    CFractions frLowPart(0, CFractions::STD);

    int to = peglevel.nSupply + peglevel.nShift;
    if (to >=0 &&
            to <PEG_SIZE &&
            peglevel.nShiftLastPart >0 &&
            peglevel.nShiftLastTotal >0) {
        // partial value to use
        int64_t v = f[to];
        int64_t vpart = ::RatioPart(v,
                                    peglevel.nShiftLastPart,
                                    peglevel.nShiftLastTotal);
        if (vpart < v) vpart++;
        frLowPart.f[to] = vpart;
        if (total) *total += vpart;
    }

    for(int i=0; i<to; i++) {
        if (total) *total += f[i];
        frLowPart.f[i] = f[i];
    }
    return frLowPart;
}
CFractions CFractions::HighPart(const CPegLevel & peglevel, int64_t* total) const
{
    if ((nFlags & STD) == 0) {
        return Std().HighPart(peglevel, total);
    }

    CFractions frHighPart(0, CFractions::STD);

    int from = peglevel.nSupply + peglevel.nShift;
    if (from >=0 &&
            from <PEG_SIZE &&
            peglevel.nShiftLastPart >0 &&
            peglevel.nShiftLastTotal >0) {
        // partial value to use
        int64_t v = f[from];
        int64_t vpart = ::RatioPart(v,
                                    peglevel.nShiftLastPart,
                                    peglevel.nShiftLastTotal);
        if (vpart < v) vpart++;
        frHighPart.f[from] = (v - vpart);
        if (total) *total += (v - vpart);
        from++;
    }

    for(int i=from; i<PEG_SIZE; i++) {
        if (total) *total += f[i];
        frHighPart.f[i] = f[i];
    }
    return frHighPart;
}

CFractions CFractions::MidPart(const CPegLevel & peglevel_low,
                               const CPegLevel & peglevel_high) const
{
    int from = peglevel_low.nSupply + peglevel_low.nShift;
    int to = peglevel_high.nSupply + peglevel_high.nShift;

    if (from != to) {
        return HighPart(peglevel_low, nullptr).LowPart(peglevel_high, nullptr);
    }

    // same supply, different partial ratio
    CFractions mid(0, STD);
    int one = from;
    int64_t v = f[one];
    int64_t vone = f[one];

    if (peglevel_low.nShiftLastPart >0 &&
        peglevel_low.nShiftLastTotal >0) {
        int64_t vpart = ::RatioPart(vone,
                                    peglevel_low.nShiftLastPart,
                                    peglevel_low.nShiftLastTotal);
        v -= vpart;
    }

    if (peglevel_high.nShiftLastPart >0 &&
        peglevel_high.nShiftLastTotal >0) {
        int64_t vpart = ::RatioPart(vone,
                                    peglevel_high.nShiftLastPart,
                                    peglevel_high.nShiftLastTotal);
        if (vpart < vone) vpart++;
        v -= (vone - vpart);
    }

    mid.f[from] = v;
    return mid;
}

/** Take a part as ration part/total where part is also value (sum fraction)
 *  Returned fractions are also adjusted for part for rounding differences.
 *  #NOTE8
 */
CFractions CFractions::RatioPart(int64_t nPartValue) const {
    if ((nFlags & STD) == 0) {
        return Std().RatioPart(nPartValue);
    }
    int64_t nTotalValue = Total();
    int64_t nPartValueSum = 0;
    CFractions fPart(0, CFractions::STD);

    if (nPartValue == 0 && nTotalValue == 0)
        return fPart;
    if (nPartValue == 0)
        return fPart;
    if (nPartValue > nTotalValue)
        return Std();

    int adjust_from = PEG_SIZE;
    for(int i=0; i<PEG_SIZE; i++) {
        int64_t v = f[i];

        if (v != 0 && i < adjust_from) {
            adjust_from = i;
        }

        bool has_overflow = false;
        if (std::is_same<int64_t,long>()) {
            long m_test;
            has_overflow = __builtin_smull_overflow(v, nPartValue, &m_test);
        } else if (std::is_same<int64_t,long long>()) {
            long long m_test;
            has_overflow = __builtin_smulll_overflow(v, nPartValue, &m_test);
        } else {
            assert(0); // todo: compile error
        }

        if (has_overflow) {
            multiprecision::uint128_t v128(v);
            multiprecision::uint128_t part128(nPartValue);
            multiprecision::uint128_t f128 = (v128*part128)/nTotalValue;
            fPart.f[i] = f128.convert_to<int64_t>();
        }
        else {
            fPart.f[i] = (v*nPartValue)/nTotalValue;
        }

        nPartValueSum += fPart.f[i];
    }

    if (nPartValueSum == nPartValue)
        return fPart;
    if (nPartValueSum > nPartValue)
        return fPart; // todo:peg: validate if possible

    int idx = adjust_from;
    int64_t nAdjustValue = nPartValue - nPartValueSum;
    while(nAdjustValue >0) {
        // todo:peg: review all possible cases if rounding mismatch with adjust_from
        if (fPart.f[idx] < f[idx]) {
            nAdjustValue--;
            fPart.f[idx]++;
        }
        idx++;
        if (idx >= PEG_SIZE) {
            idx = adjust_from;
        }
    }

    return fPart;
}

/** Take a part as ration part/total where part is also value (sum fraction)
 *  Add taken part into destination fractions. Adjusted from adjust_from.
 *  Returns not completed amount (if source("this") has no enough)
 *  #NOTE8
 */
int64_t CFractions::MoveRatioPartTo(int64_t nValueToMove,
                                    CFractions& b)
{
    int64_t nTotalValue = Total();
    int64_t nPartValue = nValueToMove;

    if (nTotalValue == 0)
        return nValueToMove;
    if (nValueToMove == 0)
        return 0;

    if ((nFlags & STD) == 0)
        ToStd();
    if ((b.nFlags & STD) == 0)
        b.ToStd();

    if (nPartValue >= nTotalValue) {
        nPartValue = nTotalValue;
        b += *this; // move all
        for(int i=0; i<PEG_SIZE; i++) f[i] = 0; // taken all
        return nValueToMove - nPartValue;
    }

    int64_t nPartValueSum = 0;
    int adjust_from = PEG_SIZE;
    for(int i=0; i<PEG_SIZE; i++) {
        int64_t v = f[i];

        if (v != 0 && i < adjust_from) {
            adjust_from = i;
        }

        bool has_overflow = false;
        if (std::is_same<int64_t,long>()) {
            long m_test;
            has_overflow = __builtin_smull_overflow(v, nPartValue, &m_test);
        } else if (std::is_same<int64_t,long long>()) {
            long long m_test;
            has_overflow = __builtin_smulll_overflow(v, nPartValue, &m_test);
        } else {
            assert(0); // todo: compile error
        }

        int64_t vp = 0;

        if (has_overflow) {
            multiprecision::uint128_t v128(v);
            multiprecision::uint128_t part128(nPartValue);
            multiprecision::uint128_t f128 = (v128*part128)/nTotalValue;
            vp = f128.convert_to<int64_t>();
        }
        else {
            vp = (v*nPartValue)/nTotalValue;
        }

        nPartValueSum += vp;
        b.f[i] += vp;
        f[i] -= vp;
    }

    if (nPartValueSum == nPartValue)
        return 0;
    if (nPartValueSum > nPartValue)
        return 0; // todo:peg: validate if possible

    int idx = adjust_from;
    int64_t nAdjustValue = nPartValue - nPartValueSum;
    while(nAdjustValue >0) {
        if (f[idx] >0) {
            nAdjustValue--;
            b.f[idx]++;
            f[idx]--;
        }
        idx++;
        if (idx >= PEG_SIZE) {
            idx = adjust_from;
        }
    }

    return 0;
}

CFractions& CFractions::operator+=(const CFractions& b)
{
    if ((b.nFlags & STD) == 0) {
        return operator+=(b.Std());
    }
    if ((nFlags & STD) == 0) {
        ToStd();
    }
    for(int i=0; i<PEG_SIZE; i++) {
        f[i] += b.f[i];
    }
    return *this;
}

CFractions& CFractions::operator-=(const CFractions& b)
{
    if ((b.nFlags & STD) == 0) {
        return operator-=(b.Std());
    }
    if ((nFlags & STD) == 0) {
        ToStd();
    }
    for(int i=0; i<PEG_SIZE; i++) {
        f[i] -= b.f[i];
    }
    return *this;
}

CFractions CFractions::operator&(const CFractions& b) const
{
    CFractions a = *this;
    for(int i=0; i<PEG_SIZE; i++) {
        int64_t va = a.f[i];
        int64_t vb = b.f[i];
        if      (va >=0 && vb >=0) a.f[i] = std::min(va, vb);
        else if (va >=0 && vb < 0) a.f[i] = 0;
        else if (va < 0 && vb >=0) a.f[i] = 0;
        else if (va < 0 && vb < 0) a.f[i] = std::max(va, vb);
    }
    return a;
}

CFractions CFractions::operator-() const
{
    CFractions a = *this;
    for(int i=0; i<PEG_SIZE; i++) {
        int64_t va = a.f[i];
        a.f[i] = -va;
    }
    return a;
}

double CFractions::Distortion(const CFractions& b) const
{
    int64_t nTotalA = Total();
    int64_t nTotalB = b.Total();

    if (nTotalA == nTotalB) {

        if (nTotalA == 0) {
            return 0;
        }

        int64_t nDiff = 0;

        for(int i=0; i<PEG_SIZE; i++) {
            int64_t va = f[i];
            int64_t vb = b.f[i];
            if (va > vb) {
                nDiff += (va - vb);
            }
        }

        return double(nDiff) / double(nTotalA);
    }

    else if (nTotalA < nTotalB) {
        if (nTotalA == 0) {
            return nTotalB; // does not make sense to scale
        }

        CFractions b1 = b.RatioPart(nTotalA);
        return Distortion(b1);
    }

    else if (nTotalA > nTotalB) {
        if (nTotalB == 0) {
            return nTotalA; // does not make sense to scale
        }

        CFractions a1 = RatioPart(nTotalB);
        return a1.Distortion(b);
    }

    return 0;
}


/*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
|  Phycas: Python software for phylogenetic analysis                          |
|  Copyright (C) 2006 Mark T. Holder, Paul O. Lewis and David L. Swofford     |
|                                                                             |
|  This program is free software; you can redistribute it and/or modify       |
|  it under the terms of the GNU General Public License as published by       |
|  the Free Software Foundation; either version 2 of the License, or          |
|  (at your option) any later version.                                        |
|                                                                             |
|  This program is distributed in the hope that it will be useful,            |
|  but WITHOUT ANY WARRANTY; without even the implied warranty of             |
|  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              |
|  GNU General Public License for more details.                               |
|                                                                             |
|  You should have received a copy of the GNU General Public License along    |
|  with this program; if not, write to the Free Software Foundation, Inc.,    |
|  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.                |
\~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~*/

#if ! defined(GTR_MODEL_HPP)
#define GTR_MODEL_HPP

#include "states_patterns.hpp"
#include <boost/noncopyable.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/lambda/lambda.hpp>
#include <vector>
#include <numeric>
#include <algorithm>
#include "xlikelihood.hpp"
#include "mcmc_param.hpp"
#include "basic_cdf.hpp"
#include "q_matrix.hpp"
#include "model.hpp"
#include "multivariate_probability_distribution.hpp"

namespace phycas{

/*----------------------------------------------------------------------------------------------------------------------
|	Specialization of the base class Model that represents the General Time Reversible (GTR) model.
*/
class GTR: public Model
	{
	public:
                                        GTR();
                                        ~GTR()
                                            {
                                            //std::cerr << "GTR dying..." << std::endl;
                                            }

		virtual std::string             getModelName() const;
		virtual void                    createParameters(JointPriorManagerShPtr jpm, TreeShPtr t, MCMCUpdaterVect & edgelens, MCMCUpdaterVect & edgelen_hyperparams, MCMCUpdaterVect & parameters, int subset_pos);
		double                          calcUniformizationLambda() const;
        double                          calcLMat(double * * lMat) const;
        double                          calcUMat(double * * uMat) const;
		void                            calcPMat(double * * pMat, double edgeLength) const;

        void                            fixRelRates();
		void                            freeRelRates();

        std::vector<double>             getRelRates();
		void                            setRelRates(const std::vector<double> & rates);

        void                            calcRelRatesFromRatios();
		void                            setRelRateRatio(unsigned which, double value);

        void                            setStateFreqRatio(unsigned ratio_index, double value);

		void                            setRelRateUnnorm(unsigned param_index, double value);
		double                          getRelRateUnnorm(unsigned param_index);

		void                            setRelRateParamPrior(ProbDistShPtr d);
		ProbDistShPtr                   getRelRateParamPrior();

		void                            setRelRatePrior(MultivarProbDistShPtr d);
		MultivarProbDistShPtr           getRelRatePrior();

        void                            setNucleotideFreqs(double freqA, double freqC, double freqG, double freqT);
        void                            setAllFreqsEqual();
		virtual void                    setStateFreqUnnorm(unsigned param_index, double value);
        virtual void                    setStateFreqsUnnorm(const std::vector<double> & values);

        void                            setStateFreqPrior(MultivarProbDistShPtr d);
		MultivarProbDistShPtr           getStateFreqPrior();

        void                            setStateFreqParamPrior(ProbDistShPtr d);
		ProbDistShPtr                   getStateFreqParamPrior();

        virtual std::string             paramHeader() const;
		virtual std::string             paramReport(unsigned ndecimals, bool include_edgelen_hyperparams) const;
		double                          calcTRatio();

        virtual unsigned                getNumFreeParameters() const;
        virtual void                    appendPWKParamNames(std::vector<std::string> & names, std::string prefix = "") const;
        virtual void                    appendFreeParamNames(std::vector<std::string> & names, std::string prefix = "") const;
        virtual void                    appendParamNames(std::vector<std::string> & names, std::string prefix = "") const;
        virtual void                    appendUntransformedParamValues(std::vector<double> & values) const;
        virtual void                    appendTransformedParamValues(std::vector<double> & values) const;
        virtual bool                    setParamValueFromTransformed(std::string parameter_name, double transformed_value, TreeShPtr tree);
        virtual double                  calcLogDetJacobian() const;

	protected:

		MultivarProbDistShPtr           rel_rate_prior;		    /**< The prior distribution governing each relative rate (usually a gamma distribution with scale 1 and shape equal to the desired Dirichlet parameter) */
		ProbDistShPtr                   rel_rate_param_prior;	/**< The joint prior distribution governing all six relative rates */
		ProbDistShPtr                   freq_param_prior;	    /**< The prior distribution governing each frequency parameter (usually a gamma distribution with scale 1 and shape equal to the desired Dirichlet parameter; used if frequencies are updated separately by slice sampling) */
    	MultivarProbDistShPtr           freq_prior;	            /**< The prior distribution governing the vector of frequencies (used if frequencies are updated jointly by StateFreqMove) */
		bool                            rel_rates_fixed;	    /**< If true, the relative rate values will not change during MCMC updates */
		mutable MCMCUpdaterVect         rel_rate_params;	    /**< A vector containing copies of all six relative rate parameters (saved so that fixed/free status can be changed) */
		mutable QMatrix                 q_matrix;			    /**< A QMatrix object used to compute transition probabilities */

		string_vect_t                   relrate_name;			/**< Holds names of the 6 relative rates (e.g. rAC, rAG, rAT, rCG, rCT, rGT) used as headers in the param file */
		string_vect_t                   freq_name;				/**< Holds names of the 4 state frequencies (e.g. freqA, freqC, freqG, freqT) used as headers in the param file */

		// Below here are quantities that directly affect likelihood calculations and which should increment time_stamp when modified
		std::vector<double>             rel_rates;			    /**< A vector containing the six relative rates */
		std::vector<double>             rel_rate_ratios;	    /**< A vector for temporary storage of the five relative rate ratios: rAG/rAC, rAT/rAC, rCG/rAC, rCT/rAC, and rGT/rAC */
	};

typedef boost::shared_ptr<GTR> GTRShPtr;

} // namespace phycas

//#include "phycas/src/gtr_model.inl"

#endif


#include <koinos/system/system_calls.hpp>
#include <koinos/contracts/token/token.h>

#include <koinos/buffer.hpp>
#include <koinos/common.h>

#include <boost/multiprecision/cpp_int.hpp>

#include <string>

using namespace koinos;
using namespace koinos::contracts;

using int128_t = boost::multiprecision::int128_t;

namespace constants {

static const std::string token_name    = "Token";
static const std::string token_symbol  = "TOKEN";
constexpr uint32_t token_decimals      = 8;
constexpr std::size_t max_address_size = 25;
constexpr std::size_t max_name_size    = 32;
constexpr std::size_t max_symbol_size  = 8;
constexpr std::size_t max_buffer_size  = 2048;
std::string supply_key                 = "";

} // constants

namespace state {

system::object_space contract_space()
{
   system::object_space obj_space;
   auto contract_id = system::get_contract_id();
   obj_space.mutable_zone().set( reinterpret_cast< const uint8_t* >( contract_id.data() ), contract_id.size() );
   obj_space.set_id( 0 );
   return obj_space;
}

}

enum class multicodec : std::uint64_t
{
   identity   = 0x00,
   sha1       = 0x11,
   sha2_256   = 0x12,
   sha2_512   = 0x13,
   ripemd_160 = 0x1053
};

enum entries : uint32_t
{
   name_entry               = 0x76ea4297,
   symbol_entry             = 0x7e794b24,
   decimals_entry           = 0x59dc15ce,
   total_supply_entry       = 0xcf2e8212,
   balance_of_entry         = 0x15619248,
   transfer_entry           = 0x62efa292,
   mint_entry               = 0xc2f82bdc
};

token::name_result< constants::max_name_size > name()
{
   token::name_result< constants::max_name_size > res;
   res.mutable_value() = constants::token_name.c_str();
   return res;
}

token::symbol_result< constants::max_symbol_size > symbol()
{
   token::symbol_result< constants::max_symbol_size > res;
   res.mutable_value() = constants::token_symbol.c_str();
   return res;
}

token::decimals_result decimals()
{
   token::decimals_result res;
   res.mutable_value() = constants::token_decimals;
   return res;
}

token::total_supply_result total_supply()
{
   token::total_supply_result res;

   token::balance_object bal_obj;
   system::get_object( state::contract_space(), constants::supply_key, bal_obj );

   res.mutable_value() = bal_obj.get_value();
   return res;
}

token::balance_of_result balance_of( const token::balance_of_arguments< constants::max_address_size >& args )
{
   token::balance_of_result res;

   std::string owner( reinterpret_cast< const char* >( args.get_owner().get_const() ), args.get_owner().get_length() );

   token::mana_balance_object bal_obj;
   system::get_object( state::contract_space(), owner, bal_obj );

   res.set_value( bal_obj.get_balance() );
   return res;
}

token::transfer_result transfer( const token::transfer_arguments< constants::max_address_size, constants::max_address_size >& args )
{
   token::transfer_result res;
   res.set_value( false );

   std::string from( reinterpret_cast< const char* >( args.get_from().get_const() ), args.get_from().get_length() );
   std::string to( reinterpret_cast< const char* >( args.get_to().get_const() ), args.get_to().get_length() );
   uint64_t value = args.get_value();

   if ( from == to )
   {
      system::print( "cannot transfer to self\n" );
      return res;
   }

   system::require_authority( from );

   token::balance_object from_bal_obj;
   system::get_object( state::contract_space(), from, from_bal_obj );

   if ( from_bal_obj.value() < value )
   {
      system::print( "'from' has insufficient balance\n" );
      return res;
   }

   token::balance_object to_bal_obj;
   system::get_object( state::contract_space(), to, to_bal_obj );

   from_bal_obj.set_value( from_bal_obj.value() - value );
   to_bal_obj.set_value( to_bal_obj.value() + value );

   system::put_object( state::contract_space(), from, from_bal_obj );
   system::put_object( state::contract_space(), to, to_bal_obj );

   res.set_value( true );
   return res;
}

token::mint_result mint( const token::mint_arguments< constants::max_address_size >& args )
{
   token::mint_result res;
   res.set_value( false );

   std::string to( reinterpret_cast< const char* >( args.get_to().get_const() ), args.get_to().get_length() );
   uint64_t amount = args.get_value();

   system::require_authority( system::get_contract_id() );

   auto supply = total_supply().get_value();
   auto new_supply = supply + amount;

   // Check overflow
   if ( new_supply < supply )
   {
      system::print( "mint would overflow supply\n" );
      return res;
   }

   token::balance_object to_bal_obj;
   system::get_object( state::contract_space(), to, to_bal_obj );

   to_bal_obj.set_value( to_bal_obj.value() + amount );

   token::balance_object supply_obj;
   supply_obj.set_value( new_supply );

   system::put_object( state::contract_space(), constants::supply_key, supply_obj );
   system::put_object( state::contract_space(), to, to_bal_obj );

   res.set_value( true );
   return res;
}

int main()
{
   auto entry_point = system::get_entry_point();
   auto args = system::get_contract_arguments();

   std::array< uint8_t, constants::max_buffer_size > retbuf;

   koinos::read_buffer rdbuf( (uint8_t*)args.c_str(), args.size() );
   koinos::write_buffer buffer( retbuf.data(), retbuf.size() );

   switch( std::underlying_type_t< entries >( entry_point ) )
   {
      case entries::name_entry:
      {
         auto res = name();
         res.serialize( buffer );
         break;
      }
      case entries::symbol_entry:
      {
         auto res = symbol();
         res.serialize( buffer );
         break;
      }
      case entries::decimals_entry:
      {
         auto res = decimals();
         res.serialize( buffer );
         break;
      }
      case entries::total_supply_entry:
      {
         auto res = total_supply();
         res.serialize( buffer );
         break;
      }
      case entries::balance_of_entry:
      {
         token::balance_of_arguments< constants::max_address_size > arg;
         arg.deserialize( rdbuf );

         auto res = balance_of( arg );
         res.serialize( buffer );
         break;
      }
      case entries::transfer_entry:
      {
         token::transfer_arguments< constants::max_address_size, constants::max_address_size > arg;
         arg.deserialize( rdbuf );

         auto res = transfer( arg );
         res.serialize( buffer );
         break;
      }
      case entries::mint_entry:
      {
         token::mint_arguments< constants::max_address_size > arg;
         arg.deserialize( rdbuf );

         auto res = mint( arg );
         res.serialize( buffer );
         break;
      }
      default:
         system::exit_contract( 1 );
   }

   std::string retval( reinterpret_cast< const char* >( buffer.data() ), buffer.get_size() );
   system::set_contract_result_bytes( retval );

   system::exit_contract( 0 );
   return 0;
}