/**
 * @file kmedoids_ucb.cpp
 * @date 2020-06-10
 *
 * This file contains the primary C++ implementation of the BanditPAM code.
 *
 */
#include "kmedoids_ucb.hpp"

#include <carma.h>
#include <armadillo>
#include <unordered_map>
#include <regex>
//#include <sstream>

/**
 *  \brief Class implementation for running KMedoids methods.
 *
 *  KMedoids class. Creates a KMedoids object that can be used to find the medoids
 *  for a particular set of input data.
 *
 *  @param n_medoids Number of medoids/clusters to create
 *  @param algorithm Algorithm used to find medoids; options are "BanditPAM" for
 *  the "Bandit-PAM" algorithm, or "naive" to use the naive method
 *  @param verbosity Verbosity of the algorithm, 0 will have no log file
 *  emitted, 1 will emit a log file
 *  @param max_iter The maximum number of iterations the algorithm runs for
 *  @param logFilename The name of the output log file
 */
KMedoids::KMedoids(int n_medoids, std::string algorithm, int verbosity,
                                          int max_iter, std::string logFilename
    ): n_medoids(n_medoids),
       algorithm(algorithm),
       max_iter(max_iter),
       verbosity(verbosity),
       logFilename(logFilename) {
  KMedoids::checkAlgorithm(algorithm);
}

/**
 *  \brief Destroys KMedoids object.
 *
 *  Destructor for the KMedoids class.
 */
KMedoids::~KMedoids() {;}

/**
 *  \brief Checks whether algorithm input is valid
 *
 *  Checks whether the user's selected algorithm is a valid option.
 *
 *  @param algorithm Name of the algorithm input by the user.
 */
void KMedoids::checkAlgorithm(std::string algorithm) {
  if (algorithm == "BanditPAM") {
    fitFn = &KMedoids::fit_bpam;
  } else if (algorithm == "naive") {
    fitFn = &KMedoids::fit_naive;
  } else {
    throw "unrecognized algorithm";
  }
}

/**
 *  \brief Returns the final medoids
 *
 *  Returns the final medoids at the end of the SWAP step after KMedoids::fit
 *  has been called.
 */
arma::rowvec KMedoids::getMedoidsFinal() {
  return medoid_indices_final;
}

/**
 *  \brief Returns the build medoids
 *
 *  Returns the build medoids at the end of the BUILD step after KMedoids::fit
 *  has been called.
 */
arma::rowvec KMedoids::getMedoidsBuild() {
  return medoid_indices_build;
}

/**
 *  \brief Returns the medoid assignments for each datapoint
 *
 *  Returns the medoid each input datapoint is assigned to after KMedoids::fit
 *  has been called and the final medoids have been identified
 */
arma::rowvec KMedoids::getLabels() {
  return labels;
}

/**
 *  \brief Returns the number of swap steps
 *
 *  Returns the number of SWAP steps completed during the last call to
 *  KMedoids::fit
 */
int KMedoids::getSteps() {
  return steps;
}

/**
 *  \brief Sets the loss function
 *
 *  Sets the loss function used during KMedoids::fit
 *
 *  @param loss Loss function to be used e.g. L2
 */
void KMedoids::setLossFn(std::string loss) {
  if (std::regex_match(loss, std::regex("L\\d*"))) {
      loss = loss.substr(1);
  }
  try {
    if (loss == "manhattan") {
        lossFn = &KMedoids::manhattan;
    } else if (loss == "cos") {
        lossFn = &KMedoids::cos;
    } else if (loss == "inf") {
        lossFn = &KMedoids::LINF;
    } else if (std::isdigit(loss.at(0))) {
        lossFn = &KMedoids::LP;
        lp     = atoi(loss.c_str());
    } else {
        throw std::invalid_argument("error: unrecognized loss function");
    }
  } catch (std::invalid_argument& e) {
      std::cout << e.what() << std::endl;
    }
}

/**
 *  \brief Returns the number of medoids
 *
 *  Returns the number of medoids to be identified during KMedoids::fit
 */
int KMedoids::getNMedoids() {
  return n_medoids;
}

/**
 *  \brief Sets the number of medoids
 *
 *  Sets the number of medoids to be identified during KMedoids::fit
 */
void KMedoids::setNMedoids(int new_num) {
  n_medoids = new_num;
}

/**
 *  \brief Returns the algorithm for KMedoids
 *
 *  Returns the algorithm used for identifying the medoids during KMedoids::fit
 */
std::string KMedoids::getAlgorithm() {
  return algorithm;
}

/**
 *  \brief Sets the algorithm for KMedoids
 *
 *  Sets the algorithm used for identifying the medoids during KMedoids::fit
 *
 *  @param new_alg New algorithm to use
 */
void KMedoids::setAlgorithm(std::string new_alg) {
  algorithm = new_alg;
}

/**
 *  \brief Returns the verbosity for KMedoids
 *
 *  Returns the verbosity used during KMedoids::fit, with 0 not creating a
 *  logfile, and >0 creating a detailed logfile.
 */
int KMedoids::getVerbosity() {
  return verbosity;
}

/**
 *  \brief Sets the verbosity for KMedoids
 *
 *  Sets the verbosity used during KMedoids::fit, with 0 not creating a
 *  logfile, and >0 creating a detailed logfile.
 *
 *  @param new_ver New verbosity to use
 */
void KMedoids::setVerbosity(int new_ver) {
  verbosity = new_ver;
}

/**
 *  \brief Returns the maximum number of iterations for KMedoids
 *
 *  Returns the maximum number of iterations that can be run during
 *  KMedoids::fit
 */
int KMedoids::getMaxIter() {
  return max_iter;
}

/**
 *  \brief Sets the maximum number of iterations for KMedoids
 *
 *  Sets the maximum number of iterations that can be run during KMedoids::fit
 *
 *  @param new_max New maximum number of iterations to use
 */
void KMedoids::setMaxIter(int new_max) {
  max_iter = new_max;
}

/**
 *  \brief Returns the log filename for KMedoids
 *
 *  Returns the name of the logfile that will be output at the end of
 *  KMedoids::fit if verbosity is >0
 */
std::string KMedoids::getLogfileName() {
  return logFilename;
}

/**
 *  \brief Sets the log filename for KMedoids
 *
 *  Sets the name of the logfile that will be output at the end of
 *  KMedoids::fit if verbosity is >0
 *
 *  @param new_lname New logfile name
 */
void KMedoids::setLogFilename(std::string new_lname) {
  logFilename = new_lname;
}

/**
 * \brief Finds medoids for the input data under identified loss function
 *
 * Primary function of the KMedoids class. Identifies medoids for input dataset
 * after both the SWAP and BUILD steps, and outputs logs if verbosity is >0
 *
 * @param input_data Input data to find the medoids of
 * @param loss The loss function used during medoid computation
 */
void KMedoids::fit(arma::mat input_data, std::string loss) {
  KMedoids::setLossFn(loss);
  (this->*fitFn)(input_data);
  if (verbosity > 0) {
      logHelper.init(logFilename);
      logHelper.writeProfile(medoid_indices_build, medoid_indices_final, steps,
                                                        logHelper.loss_swap.back());
      logHelper.close();
  }
}

/**
 * \brief Runs naive PAM algorithm.
 *
 * Run the naive PAM algorithm to identify a dataset's medoids.
 *
 * @param input_data Input data to find the medoids of
 */
void KMedoids::fit_naive(arma::mat input_data) {
  data = input_data;
  data = arma::trans(data);
  arma::rowvec medoid_indices(n_medoids);
  // runs build step
  KMedoids::build_naive(data, medoid_indices);
  steps = 0;

  medoid_indices_build = medoid_indices;
  arma::rowvec assignments(data.n_cols);
  size_t i = 0;
  bool medoidChange = true;
  while (i < max_iter && medoidChange) {
    auto previous(medoid_indices);
    // runs swa step as necessary
    KMedoids::swap_naive(data, medoid_indices, assignments);
    medoidChange = arma::any(medoid_indices != previous);
    i++;
  }
  medoid_indices_final = medoid_indices;
  labels = assignments;
  steps = i;
}

/**
 * \brief Build step for the naive algorithm
 *
 * Runs build step for the naive PAM algorithm. Loops over all datapoint and
 * checks its distance from every other datapoint in the dataset, then checks if
 * the total cost is less than that of the medoid (if a medoid exists yet).
 *
 * @param data Transposed input data to find the medoids of
 * @param medoid_indices Uninitialized array of medoids that is modified in place
 * as medoids are identified
 */
void KMedoids::build_naive(
  arma::mat& data,
  arma::rowvec& medoid_indices)
{
  size_t N = data.n_cols;
  int p = (buildConfidence * N); // reciprocal
  bool use_absolute = true;
  arma::rowvec estimates(N, arma::fill::zeros);
  arma::rowvec best_distances(N);
  best_distances.fill(std::numeric_limits<double>::infinity());
  arma::rowvec sigma(N); // standard deviation of induced losses on reference points
  for (size_t k = 0; k < n_medoids; k++) {
    double minDistance = std::numeric_limits<double>::infinity();
    int best = 0;
    KMedoids::build_sigma(
           data, best_distances, sigma, batchSize, use_absolute); // computes std dev amongst batch of reference points
    // fixes a base datapoint
    for (int i = 0; i < data.n_cols; i++) {
      double total = 0;
      for (size_t j = 0; j < data.n_cols; j++) {
        // computes distance between base and all other points
        double cost = (this->*lossFn)(data, i, j);
        for (size_t medoid = 0; medoid < k; medoid++) {
          double current = (this->*lossFn)(data, medoid_indices(medoid), j);
          // compares this for cost of the medoid
          if (current < cost) {
            cost = current;
          }
        }
        total += cost;
      }
      if (total < minDistance) {
        minDistance = total;
        best = i;
      }
    }
    // updates the medoid index for that of lowest cost.
    medoid_indices(k) = best;

    // don't need to do this on final iteration
    for (size_t l = 0; l < N; l++) {
        double cost = (this->*lossFn)(data, l, medoid_indices(k));
        if (cost < best_distances(l)) {
            best_distances(l) = cost;
        }
    }
    use_absolute = false; // use difference of loss for sigma and sampling,
                          // not absolute
    logHelper.loss_build.push_back(minDistance/N);
    logHelper.p_build.push_back((float)1/(float)p);
    logHelper.comp_exact_build.push_back(N);
  }
}

/**
 * \brief Swap step for the naive algorithm
 *
 * Runs build step for the naive PAM algorithm. Loops over all datapoint and
 * checks its distance from every other datapoint in the dataset, then checks if
 * the total cost is less than that of the medoid.
 *
 * @param data Transposed input data to find the medoids of
 * @param medoid_indices Array of medoid indices created from the build step
 * that is modified in place as better medoids are identified
 * @param assignments Uninitialized array of indices corresponding to each
 * datapoint assigned the index of the medoid it is closest to
 */
void KMedoids::swap_naive(
  arma::mat& data,
  arma::rowvec& medoid_indices,
  arma::rowvec& assignments)
{
  double minDistance = std::numeric_limits<double>::infinity();
  size_t best = 0;
  size_t medoid_to_swap = 0;
  size_t N = data.n_cols;
  int p = (N * n_medoids * swapConfidence); // reciprocal
  arma::mat sigma(n_medoids, N, arma::fill::zeros);
  arma::rowvec best_distances(N);
  arma::rowvec second_distances(N);

  // calculate quantities needed for swap, best_distances and sigma
  calc_best_distances_swap(
    data, medoid_indices, best_distances, second_distances, assignments);

  swap_sigma(data,
              sigma,
              batchSize,
              best_distances,
              second_distances,
              assignments);

  // write the sigma distribution to logfile
  sigma_log(sigma);

  // iterate across the current medoids
  for (size_t k = 0; k < n_medoids; k++) {
    // for every point in our dataset, let it serve as a "base" point
    for (size_t i = 0; i < data.n_cols; i++) {
      double total = 0;
      for (size_t j = 0; j < data.n_cols; j++) {
        // compute distance between base point and every other datapoint
        double cost = (this->*lossFn)(data, i, j);
        for (size_t medoid = 0; medoid < n_medoids; medoid++) {
          if (medoid == k) {
            continue;
          }
          double current = (this->*lossFn)(data, medoid_indices(medoid), j);
          if (current < cost) {
            cost = current;
          }
        }
        total += cost;
      }
      // if total distance for new base point is better than that of the medoid,
      // update the best index identified so far
      if (total < minDistance) {
        minDistance = total;
        best = i;
        medoid_to_swap = k;
      }
    }
  }
  medoid_indices(medoid_to_swap) = best;
  logHelper.loss_swap.push_back(minDistance/N);
  logHelper.p_swap.push_back((float)1/(float)p);
  logHelper.comp_exact_swap.push_back(N*n_medoids);
}

/**
 * \brief Runs BanditPAM algorithm.
 *
 * Run the BanditPAM algorithm to identify a dataset's medoids.
 *
 * @param input_data Input data to find the medoids of
 */
void KMedoids::fit_bpam(arma::mat input_data) {
  data = input_data;
  data = arma::trans(data);
  arma::mat medoids_mat(data.n_rows, n_medoids);
  arma::rowvec medoid_indices(n_medoids);
  // runs build step
  KMedoids::build(data, medoid_indices, medoids_mat);
  steps = 0;

  medoid_indices_build = medoid_indices;
  arma::rowvec assignments(data.n_cols);
  // runs swap step
  KMedoids::swap(data, medoid_indices, medoids_mat, assignments);
  medoid_indices_final = medoid_indices;
  labels = assignments;
}

/**
 * \brief Build step for BanditPAM
 *
 * Runs build step for the BanditPAM algorithm. Draws batch sizes with replacement
 * from reference set, and uses the estimated reward of the potential medoid
 * solutions on the reference set to update the reward confidence intervals and
 * accordingly narrow the solution set.
 *
 * @param data Transposed input data to find the medoids of
 * @param medoid_indices Uninitialized array of medoids that is modified in place
 * as medoids are identified
 * @param medoids Matrix of possible medoids that is updated as the bandit
 * learns which datapoints will be unlikely to be good candidates
 */
void KMedoids::build(
  arma::mat& data,
  arma::rowvec& medoid_indices,
  arma::mat& medoids)
{
    // Parameters
    size_t N = data.n_cols;
    arma::rowvec N_mat(N);
    N_mat.fill(N);
    int p = (buildConfidence * N); // reciprocal of
    bool use_absolute = true;
    arma::rowvec estimates(N, arma::fill::zeros);
    arma::rowvec best_distances(N);
    best_distances.fill(std::numeric_limits<double>::infinity());
    arma::rowvec sigma(N); // standard deviation of induced losses on reference points
    arma::urowvec candidates(
      N,
      arma::fill::ones); // one hot encoding of candidates -- points not filtered out yet
    arma::rowvec lcbs(N);
    arma::rowvec ucbs(N);
    arma::rowvec T_samples(N, arma::fill::zeros); // number of times calculating induced loss for reference point
    arma::rowvec exact_mask(N, arma::fill::zeros); // computed the loss exactly for this datapoint

    for (size_t k = 0; k < n_medoids; k++) {
        // instantiate medoids one-by-online
        size_t step_count = 0;
        candidates.fill(1);
        T_samples.fill(0);
        exact_mask.fill(0);
        estimates.fill(0);
        KMedoids::build_sigma(
           data, best_distances, sigma, batchSize, use_absolute); // computes std dev amongst batch of reference points

        while (arma::sum(candidates) > precision) { // while some candidates exist
            arma::umat compute_exactly =
              ((T_samples + batchSize) >= N_mat) != exact_mask;
            if (arma::accu(compute_exactly) > 0) {
                arma::uvec targets = find(compute_exactly);
                logHelper.comp_exact_build.push_back(targets.n_rows);
                arma::rowvec result =
                  build_target(data, targets, N, best_distances, use_absolute); // induced loss for these targets over all reference points
                estimates.cols(targets) = result;
                ucbs.cols(targets) = result;
                lcbs.cols(targets) = result;
                exact_mask.cols(targets).fill(1);
                T_samples.cols(targets) += N;
                candidates.cols(targets).fill(0);
            }
            if (arma::sum(candidates) < precision) {
                break;
            }
            arma::uvec targets = arma::find(candidates);
            arma::rowvec result = build_target(
              data, targets, batchSize, best_distances, use_absolute); // induced loss for the targets (sample)
            estimates.cols(targets) =
              ((T_samples.cols(targets) % estimates.cols(targets)) +
               (result * batchSize)) /
              (batchSize + T_samples.cols(targets)); // update the running average
            T_samples.cols(targets) += batchSize;
            arma::rowvec adjust(targets.n_rows);
            adjust.fill(p);
            adjust = arma::log(adjust);
            arma::rowvec cb_delta =
              sigma.cols(targets) %
              arma::sqrt(adjust / T_samples.cols(targets));
            ucbs.cols(targets) = estimates.cols(targets) + cb_delta;
            lcbs.cols(targets) = estimates.cols(targets) - cb_delta;
            candidates = (lcbs < ucbs.min()) && (exact_mask == 0);
            step_count++;
        }

        medoid_indices.at(k) = lcbs.index_min();
        medoids.unsafe_col(k) = data.unsafe_col(medoid_indices(k));

        // don't need to do this on final iteration
        for (size_t i = 0; i < N; i++) {
            double cost = (this->*lossFn)(data, i, medoid_indices(k));
            if (cost < best_distances(i)) {
                best_distances(i) = cost;
            }
        }
        use_absolute = false; // use difference of loss for sigma and sampling,
                              // not absolute
        logHelper.loss_build.push_back(arma::mean(arma::mean(best_distances)));
        logHelper.p_build.push_back((float)1/(float)p);
    }
}

/**
 * \brief Calculates confidence intervals in build step
 *
 * Calculates the confidence intervals about the reward for each arm
 *
 * @param data Transposed input data to find the medoids of
 * @param sigma Dispersion paramater for each datapoint
 * @param batch_size Number of datapoints sampled for updating confidence
 * intervals
 * @param best_distances Array of best distances from each point to previous set
 * of medoids
 * @param use_aboslute Determines whether the absolute cost is added to the total
 */
void KMedoids::build_sigma(
  arma::mat& data,
  arma::rowvec& best_distances,
  arma::rowvec& sigma,
  arma::uword batch_size,
  bool use_absolute)
{
    size_t N = data.n_cols;
    // without replacement, requires updated version of armadillo
    arma::uvec tmp_refs = arma::randperm(N, batch_size);
    arma::vec sample(batch_size);
// for each possible swap
#pragma omp parallel for
    for (size_t i = 0; i < N; i++) {
        // gather a sample of points
        for (size_t j = 0; j < batch_size; j++) {
            double cost = (this->*lossFn)(data, i, tmp_refs(j));
            if (use_absolute) {
                sample(j) = cost;
            } else {
                sample(j) = cost < best_distances(tmp_refs(j))
                              ? cost
                              : best_distances(tmp_refs(j));
                sample(j) -= best_distances(tmp_refs(j));
            }
        }
        sigma(i) = arma::stddev(sample);
    }
    arma::rowvec P = {0.25, 0.5, 0.75};
    arma::rowvec Q = arma::quantile(sigma, P);
    std::ostringstream sigma_out;
    sigma_out << "min: " << arma::min(sigma)
              << ", 25th: " << Q(0)
              << ", median: " << Q(1)
              << ", 75th: " << Q(2)
              << ", max: " << arma::max(sigma)
              << ", mean: " << arma::mean(sigma);
    logHelper.sigma_build.push_back(sigma_out.str());
}

/**
 * \brief Estimates the mean reward for each arm in build step
 *
 * Estimates the mean reward (or loss) for each arm in the identified targets
 * in the build step and returns a list of the estimated reward.
 *
 * @param data Transposed input data to find the medoids of
 * @param target Set of target datapoints to be estimated
 * @param batch_size Number of datapoints sampled for updating confidence
 * intervals
 * @param best_distances Array of best distances from each point to previous set
 * of medoids
 * @param use_absolute Determines whether the absolute cost is added to the total
 */
arma::rowvec KMedoids::build_target(
  arma::mat& data,
  arma::uvec& target,
  size_t batch_size,
  arma::rowvec& best_distances,
  bool use_absolute)
{
    size_t N = data.n_cols;
    arma::rowvec estimates(target.n_rows, arma::fill::zeros);
    arma::uvec tmp_refs = arma::randperm(N,
                                   batch_size); // without replacement, requires
                                                // updated version of armadillo
#pragma omp parallel for
    for (size_t i = 0; i < target.n_rows; i++) {
        double total = 0;
        for (size_t j = 0; j < tmp_refs.n_rows; j++) {
            double cost =
              (this->*lossFn)(data, tmp_refs(j), target(i));
            if (use_absolute) {
                total += cost;
            } else {
                total += cost < best_distances(tmp_refs(j))
                           ? cost
                           : best_distances(tmp_refs(j));
                total -= best_distances(tmp_refs(j));
            }
        }
        estimates(i) = total / batch_size;
    }
    return estimates;
}

/**
 * \brief Swap step for BanditPAM
 *
 * Runs Swap step for the BanditPAM algorithm. Draws batch sizes with replacement
 * from reference set, and uses the estimated reward of the potential medoid
 * solutions on the reference set to update the reward confidence intervals and
 * accordingly narrow the solution set.
 *
 * @param data Transposed input data to find the medoids of
 * @param medoid_indices Array of medoid indices created from the build step
 * that is modified in place as better medoids are identified
 * @param medoids Matrix of possible medoids that is updated as the bandit
 * learns which datapoints will be unlikely to be good candidates
 * @param assignments Uninitialized array of indices corresponding to each
 * datapoint assigned the index of the medoid it is closest to
 */
void KMedoids::swap(
  arma::mat& data,
  arma::rowvec& medoid_indices,
  arma::mat& medoids,
  arma::rowvec& assignments)
{
    size_t N = data.n_cols;
    int p = (N * n_medoids * swapConfidence); // reciprocal

    arma::mat sigma(n_medoids, N, arma::fill::zeros);

    arma::rowvec best_distances(N);
    arma::rowvec second_distances(N);
    size_t iter = 0;
    bool swap_performed = true;
    arma::umat candidates(n_medoids, N, arma::fill::ones);
    arma::umat exact_mask(n_medoids, N, arma::fill::zeros);
    arma::mat estimates(n_medoids, N, arma::fill::zeros);
    arma::mat lcbs(n_medoids, N);
    arma::mat ucbs(n_medoids, N);
    arma::umat T_samples(n_medoids, N, arma::fill::zeros);

    // continue making swaps while loss is decreasing
    while (swap_performed && iter < max_iter) {
        iter++;

        // calculate quantities needed for swap, best_distances and sigma
        calc_best_distances_swap(
          data, medoid_indices, best_distances, second_distances, assignments);

        swap_sigma(data,
                   sigma,
                   batchSize,
                   best_distances,
                   second_distances,
                   assignments);

        candidates.fill(1);
        exact_mask.fill(0);
        estimates.fill(0);
        T_samples.fill(0);

        // while there is at least one candidate (double comparison issues)
        while (arma::accu(candidates) > 0.5) {
            calc_best_distances_swap(
              data, medoid_indices, best_distances, second_distances, assignments);

            // compute exactly if it's been samples more than N times and hasn't
            // been computed exactly already
            arma::umat compute_exactly =
              ((T_samples + batchSize) >= N) != (exact_mask);
            arma::uvec targets = arma::find(compute_exactly);

            if (targets.size() > 0) {
                logHelper.comp_exact_swap.push_back(targets.size());
                arma::vec result = swap_target(data,
                                               medoid_indices,
                                               targets,
                                               N,
                                               best_distances,
                                               second_distances,
                                               assignments);
                estimates.elem(targets) = result;
                ucbs.elem(targets) = result;
                lcbs.elem(targets) = result;
                exact_mask.elem(targets).fill(1);
                T_samples.elem(targets) += N;

                candidates = (lcbs < ucbs.min()) && (exact_mask == 0);
            }
            if (arma::accu(candidates) < precision) {
                break;
            }
            targets = arma::find(candidates);
            arma::vec result = swap_target(data,
                                           medoid_indices,
                                           targets,
                                           batchSize,
                                           best_distances,
                                           second_distances,
                                           assignments);
            estimates.elem(targets) =
              ((T_samples.elem(targets) % estimates.elem(targets)) +
               (result * batchSize)) /
              (batchSize + T_samples.elem(targets));
            T_samples.elem(targets) += batchSize;
            arma::vec adjust(targets.n_rows);
            adjust.fill(p);
            adjust = arma::log(adjust);
            arma::vec cb_delta = sigma.elem(targets) %
                                 arma::sqrt(adjust / T_samples.elem(targets));

            ucbs.elem(targets) = estimates.elem(targets) + cb_delta;
            lcbs.elem(targets) = estimates.elem(targets) - cb_delta;
            candidates = (lcbs < ucbs.min()) && (exact_mask == 0);
            targets = arma::find(candidates);
        }
        // now switch medoids
        arma::uword new_medoid = lcbs.index_min();
        // extract medoid of swap
        size_t k = new_medoid % medoids.n_cols;

        // extract data point of swap
        size_t n = new_medoid / medoids.n_cols;
        swap_performed = medoid_indices(k) != n;
        steps++;

        medoid_indices(k) = n;
        medoids.col(k) = data.col(medoid_indices(k));
        calc_best_distances_swap(
          data, medoid_indices, best_distances, second_distances, assignments);
        sigma_log(sigma);
        logHelper.loss_swap.push_back(arma::mean(arma::mean(best_distances)));
        logHelper.p_swap.push_back((float)1/(float)p);
    }
}

/**
 * \brief Calculates distances in swap step
 *
 * Calculates the best and second best distances for each datapoint to one of
 * the medoids in the current medoid set.
 *
 * @param data Transposed input data to find the medoids of
 * @param medoid_indices Array of medoid indices corresponding to dataset entries
 * @param best_distances Array of best distances from each point to previous set
 * of medoids
 * @param second_best_distances Array of second smallest distances from each
 * point to previous set of medoids
 * @param assignments Assignments of datapoints to their closest medoid
 */
void KMedoids::calc_best_distances_swap(
  arma::mat& data,
  arma::rowvec& medoid_indices,
  arma::rowvec& best_distances,
  arma::rowvec& second_distances,
  arma::rowvec& assignments)
{
#pragma omp parallel for
    for (size_t i = 0; i < data.n_cols; i++) {
        double best = std::numeric_limits<double>::infinity();
        double second = std::numeric_limits<double>::infinity();
        for (size_t k = 0; k < medoid_indices.n_cols; k++) {
            double cost = (this->*lossFn)(data, medoid_indices(k), i);
            if (cost < best) {
                assignments(i) = k;
                second = best;
                best = cost;
            } else if (cost < second) {
                second = cost;
            }
        }
        best_distances(i) = best;
        second_distances(i) = second;
    }
}

/**
 * \brief Estimates the mean reward for each arm in swap step
 *
 * Estimates the mean reward (or loss) for each arm in the identified targets
 * in the swap step and returns a list of the estimated reward.
 *
 * @param data Transposed input data to find the medoids of
 * @param sigma Dispersion paramater for each datapoint
 * @param targets Set of target datapoints to be estimated
 * @param batch_size Number of datapoints sampled for updating confidence
 * intervals
 * @param best_distances Array of best distances from each point to previous set
 * of medoids
 * @param second_best_distances Array of second smallest distances from each
 * point to previous set of medoids
 * @param assignments Assignments of datapoints to their closest medoid
 */
arma::vec KMedoids::swap_target(
  arma::mat& data,
  arma::rowvec& medoid_indices,
  arma::uvec& targets,
  size_t batch_size,
  arma::rowvec& best_distances,
  arma::rowvec& second_best_distances,
  arma::rowvec& assignments)
{
    size_t N = data.n_cols;
    arma::vec estimates(targets.n_rows, arma::fill::zeros);
    arma::uvec tmp_refs = arma::randperm(N,
                                   batch_size); // without replacement, requires
                                                // updated version of armadillo

// for each considered swap
#pragma omp parallel for
    for (size_t i = 0; i < targets.n_rows; i++) {
        double total = 0;
        // extract data point of swap
        size_t n = targets(i) / medoid_indices.n_cols;
        size_t k = targets(i) % medoid_indices.n_cols;
        // calculate total loss for some subset of the data
        for (size_t j = 0; j < batch_size; j++) {
            double cost = (this->*lossFn)(data, n, tmp_refs(j));
            if (k == assignments(tmp_refs(j))) {
                if (cost < second_best_distances(tmp_refs(j))) {
                    total += cost;
                } else {
                    total += second_best_distances(tmp_refs(j));
                }
            } else {
                if (cost < best_distances(tmp_refs(j))) {
                    total += cost;
                } else {
                    total += best_distances(tmp_refs(j));
                }
            }
            total -= best_distances(tmp_refs(j));
        }
        estimates(i) = total / tmp_refs.n_rows;
    }
    return estimates;
}

/**
 * \brief Calculates confidence intervals in swap step
 *
 * Calculates the confidence intervals about the reward for each arm
 *
 * @param data Transposed input data to find the medoids of
 * @param sigma Dispersion paramater for each datapoint
 * @param batch_size Number of datapoints sampled for updating confidence
 * intervals
 * @param best_distances Array of best distances from each point to previous set
 * of medoids
 * @param second_best_distances Array of second smallest distances from each
 * point to previous set of medoids
 * @param assignments Assignments of datapoints to their closest medoid
 */
void KMedoids::swap_sigma(
  arma::mat& data,
  arma::mat& sigma,
  size_t batch_size,
  arma::rowvec& best_distances,
  arma::rowvec& second_best_distances,
  arma::rowvec& assignments)
{
    size_t N = data.n_cols;
    size_t K = sigma.n_rows;
    arma::uvec tmp_refs = arma::randperm(N,
                                   batch_size); // without replacement, requires
                                                // updated version of armadillo

    arma::vec sample(batch_size);
// for each considered swap
#pragma omp parallel for
    for (size_t i = 0; i < K * N; i++) {
        // extract data point of swap
        size_t n = i / K;
        size_t k = i % K;

        // calculate change in loss for some subset of the data
        for (size_t j = 0; j < batch_size; j++) {
            double cost = (this->*lossFn)(data, n,tmp_refs(j));

            if (k == assignments(tmp_refs(j))) {
                if (cost < second_best_distances(tmp_refs(j))) {
                    sample(j) = cost;
                } else {
                    sample(j) = second_best_distances(tmp_refs(j));
                }
            } else {
                if (cost < best_distances(tmp_refs(j))) {
                    sample(j) = cost;
                } else {
                    sample(j) = best_distances(tmp_refs(j));
                }
            }
            sample(j) -= best_distances(tmp_refs(j));
        }
        sigma(k, n) = arma::stddev(sample);
    }
}

/**
* \brief Write the sigma distribution into logfile
*
* Calculates the statistical measures of the sigma distribution
* and writes the results to the log file.
*
* @param sigma Dispersion paramater for each datapoint
*/
void KMedoids::sigma_log(arma::mat& sigma) {
  arma::rowvec flat_sigma = sigma.as_row();
  arma::rowvec P = {0.25, 0.5, 0.75};
  arma::rowvec Q = arma::quantile(flat_sigma, P);
  std::ostringstream sigma_out;
  sigma_out << "min: " << arma::min(flat_sigma)
            << ", 25th: " << Q(0)
            << ", median: " << Q(1)
            << ", 75th: " << Q(2)
            << ", max: " << arma::max(flat_sigma)
            << ", mean: " << arma::mean(flat_sigma);
  logHelper.sigma_swap.push_back(sigma_out.str());
};

/**
 * \brief Calculate loss for medoids
 *
 * Calculates the loss under the previously identified loss function of the
 * medoid indices.
 *
 * @param data Transposed input data to find the medoids of
 * @param medoid_indices Indices of the medoids in the dataset.
 */
double KMedoids::calc_loss(
  arma::mat& data,
  arma::rowvec& medoid_indices)
{
    double total = 0;

    for (size_t i = 0; i < data.n_cols; i++) {
        double cost = std::numeric_limits<double>::infinity();
        for (size_t k = 0; k < n_medoids; k++) {
            double currCost = (this->*lossFn)(data, medoid_indices(k), i);
            if (currCost < cost) {
                cost = currCost;
            }
        }
        total += cost;
    }
    return total;
}

// Loss and miscellaneous functions

/**
 * \brief LP loss
 *
 * Calculates the LP loss between the datapoints at index i and j of the dataset
 *
 * @param data Transposed input data to find the medoids of
 * @param i Index of first datapoint
 * @param j Index of second datapoint
 */
double KMedoids::LP(arma::mat& data, int i, int j) const {
    return arma::norm(data.col(i) - data.col(j), lp);
}

/**
 * \brief L2 loss
 *
 * Calculates the L2 loss between the datapoints at index i and j of the dataset
 *
 * @param i Index of first datapoint
 * @param j Index of second datapoint
 */
//double KMedoids::L2(int i, int j) const {
//    return arma::norm(data.col(i) - data.col(j), 2);
//}

/**
 * \brief cos loss
 *
 * Calculates the cosine loss between the datapoints at index i and j of the
 * dataset
 *
 * @param data Transposed input data to find the medoids of
 * @param i Index of first datapoint
 * @param j Index of second datapoint
 */
double KMedoids::cos(arma::mat& data, int i, int j) const {
    return arma::dot(data.col(i), data.col(j)) / (arma::norm(data.col(i))
                                                    * arma::norm(data.col(j)));
}

/**
 * \brief Manhattan loss
 *
 * Calculates the Manhattan loss between the datapoints at index i and j of the
 * dataset
 *
 * @param data Transposed input data to find the medoids of
 * @param i Index of first datapoint
 * @param j Index of second datapoint
 */
double KMedoids::manhattan(arma::mat& data, int i, int j) const {
    return arma::accu(arma::abs(data.col(i) - data.col(j)));
}

/**
 * \brief L_INFINITY loss
 *
 * Calculates the Manhattan loss between the datapoints at index i and j of the
 * dataset
 *
 * @param data Transposed input data to find the medoids of
 * @param i Index of first datapoint
 * @param j Index of second datapoint
 */
double KMedoids::LINF(arma::mat& data, int i, int j) const {
    return arma::max(arma::abs(data.col(i) - data.col(j)));
}


/*
  Copyright 2008 Larry Gritz and the other authors and contributors.
  All Rights Reserved.

  Redistribution and use in source and binary forms, with or without
  modification, are permitted provided that the following conditions are
  met:
  * Redistributions of source code must retain the above copyright
    notice, this list of conditions and the following disclaimer.
  * Redistributions in binary form must reproduce the above copyright
    notice, this list of conditions and the following disclaimer in the
    documentation and/or other materials provided with the distribution.
  * Neither the name of the software's owners nor the names of its
    contributors may be used to endorse or promote products derived from
    this software without specific prior written permission.
  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

  (This is the Modified BSD License)
*/

/// \file
/// ImageBufAlgo functions for filtered transformations

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

#include <OpenEXR/half.h>
#include <OpenEXR/ImathMatrix.h>
#include <OpenEXR/ImathBox.h>

#include <cmath>

#include "OpenImageIO/imagebuf.h"
#include "OpenImageIO/imagebufalgo.h"
#include "OpenImageIO/imagebufalgo_util.h"
#include "OpenImageIO/dassert.h"
#include "OpenImageIO/filter.h"
#include "OpenImageIO/thread.h"

OIIO_NAMESPACE_ENTER {

namespace {

// Poor man's Dual2<float> makes it easy to compute with differentials. For
// a rich man's implementation and full documentation, see
// OpenShadingLanguage (dual2.h).
class Dual2 {
public:
    float val() const { return m_val; }
    float dx() const { return m_dx; }
    float dy() const { return m_dy; }
    Dual2 (float val) : m_val(val), m_dx(0.0f), m_dy(0.0f) {}
    Dual2 (float val, float dx, float dy) : m_val(val), m_dx(dx), m_dy(dy) {}
    Dual2& operator= (float f) { m_val = f; m_dx = 0.0f; m_dy = 0.0f; return *this; }
    friend Dual2 operator+ (const Dual2 &a, const Dual2 &b) {
        return Dual2 (a.m_val+b.m_val, a.m_dx+b.m_dx, a.m_dy+b.m_dy);
    }
    friend Dual2 operator+ (const Dual2 &a, float b) {
        return Dual2 (a.m_val+b, a.m_dx, a.m_dy);
    }
    friend Dual2 operator* (const Dual2 &a, float b) {
        return Dual2 (a.m_val*b, a.m_dx*b, a.m_dy*b);
    }
    friend Dual2 operator* (const Dual2 &a, const Dual2 &b) {
        // Use the chain rule
        return Dual2 (a.m_val*b.m_val,
                      a.m_val*b.m_dx + a.m_dx*b.m_val,
                      a.m_val*b.m_dy + a.m_dy*b.m_val);
    }
    friend Dual2 operator/ (const Dual2 &a, const Dual2 &b) {
        float bvalinv = 1.0f / b.m_val;
        float aval_bval = a.m_val * bvalinv;
        return Dual2 (aval_bval,
                      bvalinv * (a.m_dx - aval_bval * b.m_dx),
                      bvalinv * (a.m_dy - aval_bval * b.m_dy));
    }
private:
    float m_val, m_dx, m_dy;
};

/// Transform a 2D point (x,y) with derivatives by a 3x3 affine matrix to
/// obtain a transformed point with derivatives.
inline void
robust_multVecMatrix (const Imath::M33f &M,
                      const Dual2 &x, const Dual2 &y,
                      Dual2 &outx, Dual2 &outy)
{
    Dual2 a = x * M[0][0] + y * M[1][0] + M[2][0];
    Dual2 b = x * M[0][1] + y * M[1][1] + M[2][1];
    Dual2 w = x * M[0][2] + y * M[1][2] + M[2][2];

    if (w.val() != 0.0f) {
       outx = a / w;
       outy = b / w;
    } else {
       outx = 0.0f;
       outy = 0.0f;
    }
}

// Transform an ROI by an affine matrix.
ROI
transform (const Imath::M33f &M, ROI roi)
{
    Imath::V2f ul (roi.xbegin+0.5f, roi.ybegin+0.5f);
    Imath::V2f ur (roi.xend-0.5f, roi.ybegin+0.5f);
    Imath::V2f ll (roi.xbegin+0.5f, roi.yend-0.5f);
    Imath::V2f lr (roi.xend-0.5f, roi.yend-0.5f);
    M.multVecMatrix (ul, ul);
    M.multVecMatrix (ur, ur);
    M.multVecMatrix (ll, ll);
    M.multVecMatrix (lr, lr);
    Imath::Box2f box (ul);
    box.extendBy (ll);
    box.extendBy (ur);
    box.extendBy (lr);
    int xmin = int (floorf(box.min.x));
    int ymin = int (floorf(box.min.y));
    int xmax = int (floorf(box.max.x)) + 1;
    int ymax = int (floorf(box.max.y)) + 1;
    return ROI (xmin, xmax, ymin, ymax, roi.zbegin, roi.zend, roi.chbegin, roi.chend);
}

// Given s,t image space coordinates and their derivatives, compute a
// filtered sample using the derivatives to guide the size of the filter
// footprint.
template<typename SRCTYPE>
inline void
filtered_sample (const ImageBuf &src, float s, float t,
                 float dsdx, float dtdx, float dsdy, float dtdy,
                 const Filter2D *filter, ImageBuf::WrapMode wrap,
                 float *result)
{
    DASSERT (filter);
    // Just use isotropic filtering
    float ds = std::max (1.0f, std::max (fabsf(dsdx), fabsf(dsdy)));
    float dt = std::max (1.0f, std::max (fabsf(dtdx), fabsf(dtdy)));
    float ds_inv = 1.0f / ds;
    float dt_inv = 1.0f / dt;
    float filterrad_s = 0.5f * ds * filter->width();
    float filterrad_t = 0.5f * dt * filter->width();
    ImageBuf::ConstIterator<SRCTYPE> samp (src,
                      (int)floorf(s-filterrad_s), (int)ceilf(s+filterrad_s),
                      (int)floorf(t-filterrad_t), (int)ceilf(t+filterrad_t),
                      0, 1, wrap);
    int nc = src.nchannels();
    float *sum = ALLOCA (float, nc);
    memset (sum, 0, nc*sizeof(float));
    float total_w = 0.0f;
    for ( ; ! samp.done(); ++samp) {
        float w = (*filter) (ds_inv*(samp.x()+0.5f-s), dt_inv*(samp.y()+0.5f-t));
        for (int c = 0; c < nc; ++c)
            sum[c] += w * samp[c];
        total_w += w;
    }
    if (total_w != 0.0f)
        for (int c = 0; c < nc; ++c)
            result[c] = sum[c] / total_w;
    else
        for (int c = 0; c < nc; ++c)
            result[c] = 0.0f;
}

} // end anon namespace

template<typename DSTTYPE, typename SRCTYPE>
static bool
resize_ (ImageBuf &dst, const ImageBuf &src,
         Filter2D *filter, ROI roi, int nthreads)
{
    if (nthreads != 1 && roi.npixels() >= 1000) {
        // Lots of pixels and request for multi threads? Parallelize.
        ImageBufAlgo::parallel_image (
            boost::bind(resize_<DSTTYPE,SRCTYPE>, boost::ref(dst),
                        boost::cref(src), filter,
                        _1 /*roi*/, 1 /*nthreads*/),
            roi, nthreads);
        return true;
    }

    // Serial case

    const ImageSpec &srcspec (src.spec());
    const ImageSpec &dstspec (dst.spec());
    int nchannels = dstspec.nchannels;

    // Local copies of the source image window, converted to float
    float srcfx = srcspec.full_x;
    float srcfy = srcspec.full_y;
    float srcfw = srcspec.full_width;
    float srcfh = srcspec.full_height;

    // Ratios of dst/src size.  Values larger than 1 indicate that we
    // are maximizing (enlarging the image), and thus want to smoothly
    // interpolate.  Values less than 1 indicate that we are minimizing
    // (shrinking the image), and thus want to properly filter out the
    // high frequencies.
    float xratio = float(dstspec.full_width) / srcfw; // 2 upsize, 0.5 downsize
    float yratio = float(dstspec.full_height) / srcfh;

    float dstfx = dstspec.full_x;
    float dstfy = dstspec.full_y;
    float dstfw = dstspec.full_width;
    float dstfh = dstspec.full_height;
    float dstpixelwidth = 1.0f / dstfw;
    float dstpixelheight = 1.0f / dstfh;
    float *pel = ALLOCA (float, nchannels);
    float filterrad = filter->width() / 2.0f;

    // radi,radj is the filter radius, as an integer, in source pixels.  We
    // will filter the source over [x-radi, x+radi] X [y-radj,y+radj].
    int radi = (int) ceilf (filterrad/xratio);
    int radj = (int) ceilf (filterrad/yratio);
    int xtaps = 2*radi + 1;
    int ytaps = 2*radj + 1;
    bool separable = filter->separable();
    float *xfiltval = NULL, *yfiltval = NULL;
    if (separable) {
        // Allocate temp space to cache the filter weights
        xfiltval = ALLOCA (float, xtaps);
        yfiltval = ALLOCA (float, ytaps);
    }
#if 0
    std::cerr << "Resizing " << srcspec.full_width << "x" << srcspec.full_height
              << " to " << dstspec.full_width << "x" << dstspec.full_height << "\n";
    std::cerr << "ratios = " << xratio << ", " << yratio << "\n";
    std::cerr << "examining src filter support radius of " << radi << " x " << radj << " pixels\n";
    std::cerr << "dst range " << roi << "\n";
    std::cerr << "separable filter\n";
#endif

    // We're going to loop over all output pixels we're interested in.
    //
    // (s,t) = NDC space coordinates of the output sample we are computing.
    //     This is the "sample point".
    // (src_xf, src_xf) = source pixel space float coordinates of the
    //     sample we're computing. We want to compute the weighted sum
    //     of all the source image pixels that fall under the filter when
    //     centered at that location.
    // (src_x, src_y) = image space integer coordinates of the floor,
    //     i.e., the closest pixel in the source image.
    // src_xf_frac and src_yf_frac are the position within that pixel
    //     of our sample.
    ImageBuf::Iterator<DSTTYPE> out (dst, roi);
    for (int y = roi.ybegin;  y < roi.yend;  ++y) {
        float t = (y-dstfy+0.5f)*dstpixelheight;
        float src_yf = srcfy + t * srcfh;
        int src_y;
        float src_yf_frac = floorfrac (src_yf, &src_y);

        // If using separable filters, our vertical set of filter tap
        // weights will be the same for the whole scanline we're on.  Just
        // compute and normalize them once.
        float totalweight_y = 0.0f;
        if (separable) {
            for (int j = 0;  j < ytaps;  ++j) {
                float w = filter->yfilt (yratio * (j-radj-(src_yf_frac-0.5f)));
                yfiltval[j] = w;
                totalweight_y += w;
            }
            for (int i = 0;  i < ytaps;  ++i)
                yfiltval[i] /= totalweight_y;
        }

        for (int x = roi.xbegin;  x < roi.xend;  ++x) {
            float s = (x-dstfx+0.5f)*dstpixelwidth;
            float src_xf = srcfx + s * srcfw;
            int src_x;
            float src_xf_frac = floorfrac (src_xf, &src_x);
            for (int c = 0;  c < nchannels;  ++c)
                pel[c] = 0.0f;
            if (separable) {
                // Cache and normalize the horizontal filter tap weights
                // just once for this (x,y) position, reuse for all vertical
                // taps.
                float totalweight_x = 0.0f;
                for (int i = 0;  i < xtaps;  ++i) {
                    float w = filter->xfilt (xratio * (i-radi-(src_xf_frac-0.5f)));
                    xfiltval[i] = w;
                    totalweight_x += w;
                }

                if (totalweight_x != 0.0f) {
                    for (int i = 0;  i < xtaps;  ++i)  // normalize x filter
                        xfiltval[i] /= totalweight_x;  // weights
                    ImageBuf::ConstIterator<SRCTYPE> srcpel (src, src_x-radi, src_x+radi+1,
                                                             src_y-radj, src_y+radj+1,
                                                             0, 1, ImageBuf::WrapClamp);
                    for (int j = -radj;  j <= radj;  ++j) {
                        float wy = yfiltval[j+radj];
                        if (wy == 0.0f) {
                            // 0 weight for this y tap -- move to next line
                            srcpel.pos (srcpel.x(), srcpel.y()+1, srcpel.z());
                            continue;
                        }
                        for (int i = 0;  i < xtaps; ++i, ++srcpel) {
                            float w = wy * xfiltval[i];
                            for (int c = 0;  c < nchannels;  ++c)
                                pel[c] += w * srcpel[c];
                        }
                    }
                }
                // Copy the pixel value (already normalized) to the output.
                DASSERT (out.x() == x && out.y() == y);
                if (totalweight_y == 0.0f) {
                    // zero it out
                    for (int c = 0;  c < nchannels;  ++c)
                        out[c] = 0.0f;
                } else {
                    for (int c = 0;  c < nchannels;  ++c)
                        out[c] = pel[c];
                }
            } else {
                // Non-separable
                float totalweight = 0.0f;
                ImageBuf::ConstIterator<SRCTYPE> srcpel (src, src_x-radi, src_x+radi+1,
                                                       src_y-radi, src_y+radi+1,
                                                       0, 1, ImageBuf::WrapClamp);
                for (int j = -radj;  j <= radj;  ++j) {
                    for (int i = -radi;  i <= radi;  ++i, ++srcpel) {
                        float w = (*filter)(xratio * (i-(src_xf_frac-0.5f)),
                                            yratio * (j-(src_yf_frac-0.5f)));
                        totalweight += w;
                        if (w == 0.0f)
                            continue;
                        DASSERT (! srcpel.done());
                        for (int c = 0;  c < nchannels;  ++c)
                            pel[c] += w * srcpel[c];
                    }
                }
                DASSERT (srcpel.done());
                // Rescale pel to normalize the filter and write it to the
                // output image.
                DASSERT (out.x() == x && out.y() == y);
                if (totalweight == 0.0f) {
                    // zero it out
                    for (int c = 0;  c < nchannels;  ++c)
                        out[c] = 0.0f;
                } else {
                    for (int c = 0;  c < nchannels;  ++c)
                        out[c] = pel[c] / totalweight;
                }
            }

            ++out;
        }
    }

    return true;
}

bool
ImageBufAlgo::resize (ImageBuf &dst, const ImageBuf &src,
                      Filter2D *filter, ROI roi, int nthreads)
{
    if (! IBAprep (roi, &dst, &src,
            IBAprep_REQUIRE_SAME_NCHANNELS | IBAprep_NO_SUPPORT_VOLUME |
            IBAprep_NO_COPY_ROI_FULL))
        return false;

    // Set up a shared pointer with custom deleter to make sure any
    // filter we allocate here is properly destroyed.
    boost::shared_ptr<Filter2D> filterptr ((Filter2D*)NULL, Filter2D::destroy);
    bool allocfilter = (filter == NULL);
    if (allocfilter) {
        // If no filter was provided, punt and just linearly interpolate.
        const ImageSpec &srcspec (src.spec());
        const ImageSpec &dstspec (dst.spec());
        float wratio = float(dstspec.full_width) / float(srcspec.full_width);
        float hratio = float(dstspec.full_height) / float(srcspec.full_height);
        float w = 2.0f * std::max (1.0f, wratio);
        float h = 2.0f * std::max (1.0f, hratio);
        filter = Filter2D::create ("triangle", w, h);
        filterptr.reset (filter);
    }

    bool ok;
    OIIO_DISPATCH_TYPES2 (ok, "resize", resize_,
                          dst.spec().format, src.spec().format,
                          dst, src, filter, roi, nthreads);
    return ok;
}

bool
ImageBufAlgo::resize (ImageBuf &dst, const ImageBuf &src,
                      string_view filtername_, float fwidth,
                      ROI roi, int nthreads)
{
    if (! IBAprep (roi, &dst, &src,
            IBAprep_REQUIRE_SAME_NCHANNELS | IBAprep_NO_SUPPORT_VOLUME |
            IBAprep_NO_COPY_ROI_FULL))
        return false;
    const ImageSpec &srcspec (src.spec());
    const ImageSpec &dstspec (dst.spec());

    // Resize ratios
    float wratio = float(dstspec.full_width) / float(srcspec.full_width);
    float hratio = float(dstspec.full_height) / float(srcspec.full_height);

    // Set up a shared pointer with custom deleter to make sure any
    // filter we allocate here is properly destroyed.
    boost::shared_ptr<Filter2D> filter ((Filter2D*)NULL, Filter2D::destroy);
    std::string filtername = filtername_;
    if (filtername.empty()) {
        // No filter name supplied -- pick a good default
        if (wratio > 1.0f || hratio > 1.0f)
            filtername = "blackman-harris";
        else
            filtername = "lanczos3";
    }
    for (int i = 0, e = Filter2D::num_filters();  i < e;  ++i) {
        FilterDesc fd;
        Filter2D::get_filterdesc (i, &fd);
        if (fd.name == filtername) {
            float w = fwidth > 0.0f ? fwidth : fd.width * std::max (1.0f, wratio);
            float h = fwidth > 0.0f ? fwidth : fd.width * std::max (1.0f, hratio);
            filter.reset (Filter2D::create (filtername, w, h));
            break;
        }
    }
    if (! filter) {
        dst.error ("Filter \"%s\" not recognized", filtername);
        return false;
    }

    bool ok;
    OIIO_DISPATCH_TYPES2 (ok, "resize", resize_,
                          dstspec.format, srcspec.format,
                          dst, src, filter.get(), roi, nthreads);
    return ok;
}

template<typename DSTTYPE, typename SRCTYPE>
static bool
resample_ (ImageBuf &dst, const ImageBuf &src, bool interpolate,
           ROI roi, int nthreads)
{
    if (nthreads != 1 && roi.npixels() >= 1000) {
        // Lots of pixels and request for multi threads? Parallelize.
        ImageBufAlgo::parallel_image (
            boost::bind(resample_<DSTTYPE,SRCTYPE>, boost::ref(dst),
                        boost::cref(src), interpolate,
                        _1 /*roi*/, 1 /*nthreads*/),
            roi, nthreads);
        return true;
    }

    // Serial case

    const ImageSpec &srcspec (src.spec());
    const ImageSpec &dstspec (dst.spec());
    int nchannels = src.nchannels();

    // Local copies of the source image window, converted to float
    float srcfx = srcspec.full_x;
    float srcfy = srcspec.full_y;
    float srcfw = srcspec.full_width;
    float srcfh = srcspec.full_height;

    float dstfx = dstspec.full_x;
    float dstfy = dstspec.full_y;
    float dstfw = dstspec.full_width;
    float dstfh = dstspec.full_height;
    float dstpixelwidth = 1.0f / dstfw;
    float dstpixelheight = 1.0f / dstfh;
    float *pel = ALLOCA (float, nchannels);

    ImageBuf::Iterator<DSTTYPE> out (dst, roi);
    ImageBuf::ConstIterator<SRCTYPE> srcpel (src);
    for (int y = roi.ybegin;  y < roi.yend;  ++y) {
        // s,t are NDC space
        float t = (y-dstfy+0.5f)*dstpixelheight;
        // src_xf, src_xf are image space float coordinates
        float src_yf = srcfy + t * srcfh - 0.5f;
        // src_x, src_y are image space integer coordinates of the floor
        int src_y;
        (void) floorfrac (src_yf, &src_y);
        for (int x = roi.xbegin;  x < roi.xend;  ++x) {
            float s = (x-dstfx+0.5f)*dstpixelwidth;
            float src_xf = srcfx + s * srcfw - 0.5f;
            int src_x;
            (void) floorfrac (src_xf, &src_x);

            if (interpolate) {
                src.interppixel (src_xf, src_yf, pel);
                for (int c = roi.chbegin; c < roi.chend; ++c)
                    out[c] = pel[c];
            } else {
                srcpel.pos (src_x, src_y, 0);
                for (int c = roi.chbegin; c < roi.chend; ++c)
                    out[c] = srcpel[c];
            }
            ++out;
        }
    }

    return true;
}

bool
ImageBufAlgo::resample (ImageBuf &dst, const ImageBuf &src,
                        bool interpolate, ROI roi, int nthreads)
{
    if (! IBAprep (roi, &dst, &src,
            IBAprep_REQUIRE_SAME_NCHANNELS | IBAprep_NO_SUPPORT_VOLUME |
            IBAprep_NO_COPY_ROI_FULL))
        return false;
    bool ok;
    OIIO_DISPATCH_TYPES2 (ok, "resample", resample_,
                          dst.spec().format, src.spec().format,
                          dst, src, interpolate, roi, nthreads);
    return ok;
}

#if 0
template<typename DSTTYPE, typename SRCTYPE>
static bool
affine_resample_ (ImageBuf &dst, const ImageBuf &src, const Imath::M33f &Minv,
                  ROI roi, int nthreads)
{
    if (nthreads != 1 && roi.npixels() >= 1000) {
        // Possible multiple thread case -- recurse via parallel_image
        ImageBufAlgo::parallel_image (
            boost::bind(affine_resample_<DSTTYPE,SRCTYPE>,
                        boost::ref(dst), boost::cref(src), Minv,
                        _1 /*roi*/, 1 /*nthreads*/),
            roi, nthreads);
        return true;
    }

    // Serial case
    ImageBuf::Iterator<DSTTYPE,DSTTYPE> d (dst, roi);
    ImageBuf::ConstIterator<SRCTYPE,DSTTYPE> s (src);
    for (  ;  ! d.done();  ++d) {
        Imath::V2f P (d.x()+0.5f, d.y()+0.5f);
        Minv.multVecMatrix (P, P);
        s.pos (int(floorf(P.x)), int(floorf(P.y)), d.z());
        for (int c = roi.chbegin;  c < roi.chend;  ++c)
            d[c] = s[c];
    }
    return true;
}
#endif

template<typename DSTTYPE, typename SRCTYPE>
static bool
warp_ (ImageBuf &dst, const ImageBuf &src, const Imath::M33f &M,
       const Filter2D *filter, ImageBuf::WrapMode wrap,
       ROI roi, int nthreads)
{
    if (nthreads != 1 && roi.npixels() >= 1000) {
        // Possible multiple thread case -- recurse via parallel_image
        ImageBufAlgo::parallel_image (
            boost::bind(warp_<DSTTYPE,SRCTYPE>,
                        boost::ref(dst), boost::cref(src), M,
                        filter, wrap, _1 /*roi*/, 1 /*nthreads*/),
            roi, nthreads);
        return true;
    }

    // Serial case
    int nc = dst.nchannels();
    float *pel = ALLOCA (float, nc);
    memset (pel, 0, nc*sizeof(float));
    Imath::M33f Minv = M.inverse();
    ImageBuf::Iterator<DSTTYPE> out (dst, roi);
    for (  ;  ! out.done();  ++out) {
        Dual2 x (out.x()+0.5f, 1.0f, 0.0f);
        Dual2 y (out.y()+0.5f, 0.0f, 1.0f);
        robust_multVecMatrix (Minv, x, y, x, y);
        filtered_sample<SRCTYPE> (src, x.val(), y.val(),
                                  x.dx(), y.dx(), x.dy(), y.dy(),
                                  filter, wrap, pel);
        for (int c = roi.chbegin;  c < roi.chend;  ++c)
            out[c] = pel[c];

    }
    return true;
}

bool
ImageBufAlgo::warp (ImageBuf &dst, const ImageBuf &src,
                    const Imath::M33f &M,
                    const Filter2D *filter,
                    bool recompute_roi, ImageBuf::WrapMode wrap,
                    ROI roi, int nthreads)
{
    ROI src_roi_full = src.roi_full();
    ROI dst_roi, dst_roi_full;
    if (dst.initialized()) {
        dst_roi = roi.defined() ? roi : dst.roi();
        dst_roi_full = dst.roi_full();
    } else {
        dst_roi = roi.defined() ? roi : (recompute_roi ? transform (M, src.roi()) : src.roi());
        dst_roi_full = src_roi_full;
    }
    dst_roi.chend = std::min (dst_roi.chend, src.nchannels());
    dst_roi_full.chend = std::min (dst_roi_full.chend, src.nchannels());

    if (! IBAprep (dst_roi, &dst, &src, IBAprep_NO_SUPPORT_VOLUME))
        return false;

    // Set up a shared pointer with custom deleter to make sure any
    // filter we allocate here is properly destroyed.
    boost::shared_ptr<Filter2D> filterptr ((Filter2D*)NULL, Filter2D::destroy);
    if (filter == NULL) {
        // If no filter was provided, punt and just linearly interpolate.
        filterptr.reset (Filter2D::create ("lanczos3", 6.0f, 6.0f));
        filter = filterptr.get();
    }

    bool ok;
    OIIO_DISPATCH_TYPES2 (ok, "warp", warp_,
                          dst.spec().format, src.spec().format,
                          dst, src, M, filter, wrap, dst_roi, nthreads);
    return ok;
}

bool
ImageBufAlgo::warp (ImageBuf &dst, const ImageBuf &src,
                    const Imath::M33f &M,
                    string_view filtername_, float filterwidth,
                    bool recompute_roi, ImageBuf::WrapMode wrap,
                    ROI roi, int nthreads)
{
    // Set up a shared pointer with custom deleter to make sure any
    // filter we allocate here is properly destroyed.
    boost::shared_ptr<Filter2D> filter ((Filter2D*)NULL, Filter2D::destroy);
    std::string filtername = filtername_.size() ? filtername_ : "lanczos3";
    for (int i = 0, e = Filter2D::num_filters();  i < e;  ++i) {
        FilterDesc fd;
        Filter2D::get_filterdesc (i, &fd);
        if (fd.name == filtername) {
            float w = filterwidth > 0.0f ? filterwidth : fd.width;
            float h = filterwidth > 0.0f ? filterwidth : fd.width;
            filter.reset (Filter2D::create (filtername, w, h));
            break;
        }
    }
    if (! filter) {
        dst.error ("Filter \"%s\" not recognized", filtername);
        return false;
    }

    return warp (dst, src, M, filter.get(), recompute_roi, wrap, roi, nthreads);
}

bool
ImageBufAlgo::rotate (ImageBuf &dst, const ImageBuf &src,
                      float angle, float center_x, float center_y,
                      Filter2D *filter, bool recompute_roi,
                      ROI roi, int nthreads)
{
    // Calculate the rotation matrix
    Imath::M33f M;
    M.translate(Imath::V2f(-center_x, -center_y));
    M.rotate(angle);
    M *= Imath::M33f().translate(Imath::V2f(center_x, center_y));
    return ImageBufAlgo::warp (dst, src, M, filter,
                               recompute_roi, ImageBuf::WrapBlack,
                               roi, nthreads);
}

bool
ImageBufAlgo::rotate (ImageBuf &dst, const ImageBuf &src,
                      float angle, float center_x, float center_y,
                      string_view filtername, float filterwidth,
                      bool recompute_roi, ROI roi, int nthreads)
{
    // Calculate the rotation matrix
    Imath::M33f M;
    M.translate(Imath::V2f(-center_x, -center_y));
    M.rotate(angle);
    M *= Imath::M33f().translate(Imath::V2f(center_x, center_y));
    return ImageBufAlgo::warp (dst, src, M, filtername, filterwidth,
                               recompute_roi, ImageBuf::WrapBlack,
                               roi, nthreads);
}

bool
ImageBufAlgo::rotate (ImageBuf &dst, const ImageBuf &src, float angle,
                      Filter2D *filter,
                      bool recompute_roi, ROI roi, int nthreads)
{
    ROI src_roi_full = src.roi_full();
    float center_x = 0.5f * (src_roi_full.xbegin + src_roi_full.xend);
    float center_y = 0.5f * (src_roi_full.ybegin + src_roi_full.yend);
    return ImageBufAlgo::rotate (dst, src, angle, center_x, center_y,
                                 filter, recompute_roi, roi, nthreads);
}

bool
ImageBufAlgo::rotate (ImageBuf &dst, const ImageBuf &src, float angle,
                      string_view filtername, float filterwidth,
                      bool recompute_roi,
                      ROI roi, int nthreads)
{
    ROI src_roi_full = src.roi_full();
    float center_x = 0.5f * (src_roi_full.xbegin + src_roi_full.xend);
    float center_y = 0.5f * (src_roi_full.ybegin + src_roi_full.yend);
    return ImageBufAlgo::rotate (dst, src, angle, center_x, center_y,
                                 filtername, filterwidth, recompute_roi,
                                 roi, nthreads);
}

} OIIO_NAMESPACE_EXIT


//---------------------------------------------------------------------------//
// Copyright (c) 2018-2020 Mikhail Komarov <nemo@nil.foundation>
//
// Distributed under the Boost Software License, Version 1.0
// See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt
//---------------------------------------------------------------------------//

#ifndef CRYPTO3_RIJNDAEL_IMPL_HPP
#define CRYPTO3_RIJNDAEL_IMPL_HPP

#include <boost/crypto3/block/algorithm/move.hpp>

#include <boost/crypto3/detail/stream_endian.hpp>
#include <boost/crypto3/detail/pack.hpp>

namespace boost {
    namespace crypto3 {
        namespace block {
            /*!
             * @cond DETAIL_IMPL
             */
            namespace detail {
                template<std::size_t KeyBitsImpl, std::size_t BlockBitsImpl, typename PolicyType>
                class rijndael_impl {
                    typedef PolicyType policy_type;

                    typedef typename policy_type::key_type key_type;
                    typedef typename policy_type::key_schedule_type key_schedule_type;
                    typedef typename policy_type::key_schedule_word_type key_schedule_word_type;

                    typedef typename policy_type::block_type block_type;

                    typedef typename policy_type::constants_type constants_type;
                    typedef typename policy_type::shift_offsets_type shift_offsets_type;
                    typedef typename policy_type::mm_type mm_type;

                    BOOST_STATIC_ASSERT(KeyBitsImpl == PolicyType::key_bits);
                    BOOST_STATIC_ASSERT(BlockBitsImpl == PolicyType::block_bits);

                    static inline key_schedule_word_type sub_word(const key_schedule_word_type &x,
                                                                  const constants_type &constants) {
                        key_schedule_word_type result = {0};
#pragma clang loop unroll(full)
                        for (std::size_t i = 0; i < policy_type::word_bytes; ++i) {
                            result =
                                result << CHAR_BIT | constants[::boost::crypto3::detail::extract_uint_t<CHAR_BIT>(x, i)];
                        }

                        return result;
                    }

                    static inline void sub_bytes(block_type &state, const constants_type &sbox) {
#pragma clang loop unroll(full)
                        for (std::size_t i = 0; i < policy_type::word_bytes * policy_type::block_words; ++i) {
                            state[i] = sbox[state[i]];
                        }
                    }

                    static inline void shift_rows(block_type &state, const shift_offsets_type &offset) {
                        std::array<typename policy_type::byte_type, policy_type::block_words> tmp = {0};

                        // row 0 never gets shifted
#pragma clang loop unroll(full)
                        for (std::size_t row = 1; row < policy_type::word_bytes; ++row) {
                            const std::size_t off = offset[row - 1];
#pragma clang loop unroll(full)
                            for (std::size_t i = 0; i < off; ++i) {
                                tmp[i] = state[i * policy_type::word_bytes + row];
                            }
#pragma clang loop unroll(full)
                            for (std::size_t i = 0; i < policy_type::block_words - 1; ++i) {
                                state[i * policy_type::word_bytes + row] =
                                    state[(i + off) * policy_type::word_bytes + row];
                            }
#pragma clang loop unroll(full)
                            for (std::size_t i = 0; i < off; ++i) {
                                state[(policy_type::block_words - 1 - i) * policy_type::word_bytes + row] =
                                    tmp[off - 1 - i];
                            }
                        }
                    }

                    template<typename StateType>
                    static inline block_type mix_columns(const StateType &state, const mm_type &mm) {
                        block_type tmp = {0};

#pragma clang loop unroll(full)
                        for (std::size_t col = 0; col < policy_type::block_words; ++col) {
#pragma clang loop unroll(full)
                            for (std::size_t row = 0; row < policy_type::word_bytes; ++row) {
#pragma clang loop unroll(full)
                                for (std::size_t k = 0; k < policy_type::word_bytes; ++k) {
                                    tmp[col * policy_type::word_bytes + row] ^=
                                        policy_type::mul(mm[row * policy_type::word_bytes + k],
                                                         state[col * policy_type::word_bytes + k]);
                                }
                            }
                        }

                        return tmp;
                    }

                    template<typename InputIterator>
                    static inline void add_round_key(block_type &state, InputIterator first, InputIterator last) {
                        BOOST_ASSERT(std::distance(first, last) == policy_type::block_words &&
                                     state.size() == policy_type::block_bytes);

#pragma clang loop unroll(full)
                        for (std::size_t i = 0; i < policy_type::block_words && first != last; ++i && ++first) {
#pragma clang loop unroll(full)
                            for (std::size_t j = 0; j < policy_type::word_bytes; ++j) {
                                state[i * policy_type::word_bytes + j] ^=
                                    ::boost::crypto3::detail::extract_uint_t<CHAR_BIT>(*first,
                                                                                policy_type::word_bytes - (j + 1));
                            }
                        }
                    }

                    static inline void apply_round(std::uint8_t round, block_type &state, const key_schedule_type &w,
                                                   const constants_type &sbox, const shift_offsets_type &offsets,
                                                   const mm_type &mm) {
                        sub_bytes(state, sbox);
                        shift_rows(state, offsets);
                        state = mix_columns(state, mm);
                        add_round_key(state, w.begin() + round * policy_type::block_words,
                                      w.begin() + (round + 1) * policy_type::block_words);
                    }

                public:
                    static typename policy_type::block_type encrypt_block(const block_type &plaintext,
                                                                          const key_schedule_type &encryption_key) {
                        block_type state = plaintext;

                        add_round_key(state, encryption_key.begin(), encryption_key.begin() + policy_type::block_words);

#pragma clang loop unroll(full)
                        for (std::size_t round = 1; round < policy_type::rounds; ++round) {
                            apply_round(round, state, encryption_key, policy_type::constants,
                                        policy_type::shift_offsets, policy_type::mm);
                        }

                        sub_bytes(state, policy_type::constants);
                        shift_rows(state, policy_type::shift_offsets);
                        add_round_key(state, encryption_key.begin() + policy_type::rounds * policy_type::block_words,
                                      encryption_key.begin() + (policy_type::rounds + 1) * policy_type::block_words);

                        return state;
                    }

                    static typename policy_type::block_type decrypt_block(const block_type &plaintext,
                                                                          const key_schedule_type &decryption_key) {
                        block_type state = plaintext;

                        add_round_key(state, decryption_key.begin() + policy_type::rounds * policy_type::block_words,
                                      decryption_key.begin() + (policy_type::rounds + 1) * policy_type::block_words);

#pragma clang loop unroll(full)
                        for (std::size_t round = policy_type::rounds - 1; round > 0; --round) {
                            apply_round(round, state, decryption_key, policy_type::inverted_constants,
                                        policy_type::inverted_shift_offsets, policy_type::inverted_mm);
                        }

                        sub_bytes(state, policy_type::inverted_constants);
                        shift_rows(state, policy_type::inverted_shift_offsets);
                        add_round_key(state, decryption_key.begin(), decryption_key.begin() + policy_type::block_words);

                        return state;
                    }

                    static void schedule_key(const key_type &key, key_schedule_type &encryption_key,
                                             key_schedule_type &decryption_key) {
                        // the first key_words words are the original key
                        ::boost::crypto3::detail::pack<stream_endian::big_octet_big_bit,
                                                     stream_endian::little_octet_big_bit, CHAR_BIT,
                                                     policy_type::word_bits>(
                            key.begin(), key.begin() + policy_type::key_words * policy_type::word_bytes,
                            encryption_key.begin());

#pragma clang loop unroll(full)
                        for (std::size_t i = policy_type::key_words; i < policy_type::key_schedule_words; ++i) {
                            typename policy_type::key_schedule_word_type tmp = encryption_key[i - 1];
                            if (i % policy_type::key_words == 0) {
                                tmp = sub_word(policy_type::rotate_left(tmp), policy_type::constants) ^
                                      policy_type::round_constants[i / policy_type::key_words - 1];
                            } else if (policy_type::key_words > 6 && i % policy_type::key_words == 4) {
                                tmp = sub_word(tmp, policy_type::constants);
                            }
                            encryption_key[i] = encryption_key[i - policy_type::key_words] ^ tmp;
                        }

                        std::array<typename policy_type::byte_type, policy_type::key_schedule_bytes> bekey = {0};
                        ::boost::crypto3::detail::pack<stream_endian::little_octet_big_bit,
                                                     stream_endian::big_octet_big_bit, policy_type::word_bits,
                                                     CHAR_BIT>(encryption_key.begin(), encryption_key.end(),
                                                               bekey.begin());

#pragma clang loop unroll(full)
                        for (std::uint8_t round = 1; round < policy_type::rounds; ++round) {
                            move(mix_columns(boost::adaptors::slice(bekey, round * policy_type::block_bytes,
                                                                    (round + 1) * policy_type::block_bytes),
                                             policy_type::inverted_mm),
                                 bekey.begin() + round * policy_type::block_bytes);
                        }

                        ::boost::crypto3::detail::pack<stream_endian::big_octet_big_bit,
                                                     stream_endian::little_octet_big_bit, CHAR_BIT,
                                                     policy_type::word_bits>(bekey.begin(), bekey.end(),
                                                                             decryption_key.begin());
                    }
                };
            }    // namespace detail
            /*!
             * @endcond
             */
        }    // namespace block
    }        // namespace crypto3
}    // namespace boost

#endif    // CRYPTO3_RIJNDAEL_IMPL_HPP


// Copyright (c) 2016 Jack Grigg
// Copyright (c) 2016 The Zcash developers
// Copyright (c) 2019 BTC.COM
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#pragma once

#include "sodium.h"

#include <cstring>
#include <exception>
#include <functional>
#include <memory>
#include <set>
#include <vector>

#include <boost/static_assert.hpp>

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wsign-compare"
#pragma GCC diagnostic ignored "-Wunused-variable"

namespace equihash_zcash {

typedef crypto_generichash_blake2b_state eh_HashState;
typedef uint32_t eh_index;
typedef uint8_t eh_trunc;

void ExpandArray(const unsigned char* in, size_t in_len,
                 unsigned char* out, size_t out_len,
                 size_t bit_len, size_t byte_pad=0);
void CompressArray(const unsigned char* in, size_t in_len,
                   unsigned char* out, size_t out_len,
                   size_t bit_len, size_t byte_pad=0);

eh_index ArrayToEhIndex(const unsigned char* array);
eh_trunc TruncateIndex(const eh_index i, const unsigned int ilen);

std::vector<eh_index> GetIndicesFromMinimal(std::vector<unsigned char> minimal,
                                            size_t cBitLen);
std::vector<unsigned char> GetMinimalFromIndices(std::vector<eh_index> indices,
                                                 size_t cBitLen);

template<size_t WIDTH>
class StepRow
{
    template<size_t W>
    friend class StepRow;
    friend class CompareSR;

protected:
    unsigned char hash[WIDTH];

public:
    StepRow(const unsigned char* hashIn, size_t hInLen,
            size_t hLen, size_t cBitLen);
    ~StepRow() { }

    template<size_t W>
    StepRow(const StepRow<W>& a);

    bool IsZero(size_t len);
    std::string GetHex(size_t len) { return HexStr(hash, hash+len); }

    template<size_t W>
    friend bool HasCollision(StepRow<W>& a, StepRow<W>& b, int l);
};

class CompareSR
{
private:
    size_t len;

public:
    CompareSR(size_t l) : len {l} { }

    template<size_t W>
    inline bool operator()(const StepRow<W>& a, const StepRow<W>& b) { return memcmp(a.hash, b.hash, len) < 0; }
};

template<size_t WIDTH>
bool HasCollision(StepRow<WIDTH>& a, StepRow<WIDTH>& b, int l);

template<size_t WIDTH>
class FullStepRow : public StepRow<WIDTH>
{
    template<size_t W>
    friend class FullStepRow;

    using StepRow<WIDTH>::hash;

public:
    FullStepRow(const unsigned char* hashIn, size_t hInLen,
                size_t hLen, size_t cBitLen, eh_index i);
    ~FullStepRow() { }

    FullStepRow(const FullStepRow<WIDTH>& a) : StepRow<WIDTH> {a} { }
    template<size_t W>
    FullStepRow(const FullStepRow<W>& a, const FullStepRow<W>& b, size_t len, size_t lenIndices, int trim);
    FullStepRow& operator=(const FullStepRow<WIDTH>& a);

    inline bool IndicesBefore(const FullStepRow<WIDTH>& a, size_t len, size_t lenIndices) const { return memcmp(hash+len, a.hash+len, lenIndices) < 0; }
    std::vector<unsigned char> GetIndices(size_t len, size_t lenIndices,
                                          size_t cBitLen) const;

    template<size_t W>
    friend bool DistinctIndices(const FullStepRow<W>& a, const FullStepRow<W>& b,
                                size_t len, size_t lenIndices);
    template<size_t W>
    friend bool IsValidBranch(const FullStepRow<W>& a, const size_t len, const unsigned int ilen, const eh_trunc t);
};

template<size_t WIDTH>
class TruncatedStepRow : public StepRow<WIDTH>
{
    template<size_t W>
    friend class TruncatedStepRow;

    using StepRow<WIDTH>::hash;

public:
    TruncatedStepRow(const unsigned char* hashIn, size_t hInLen,
                     size_t hLen, size_t cBitLen,
                     eh_index i, unsigned int ilen);
    ~TruncatedStepRow() { }

    TruncatedStepRow(const TruncatedStepRow<WIDTH>& a) : StepRow<WIDTH> {a} { }
    template<size_t W>
    TruncatedStepRow(const TruncatedStepRow<W>& a, const TruncatedStepRow<W>& b, size_t len, size_t lenIndices, int trim);
    TruncatedStepRow& operator=(const TruncatedStepRow<WIDTH>& a);

    inline bool IndicesBefore(const TruncatedStepRow<WIDTH>& a, size_t len, size_t lenIndices) const { return memcmp(hash+len, a.hash+len, lenIndices) < 0; }
    std::shared_ptr<eh_trunc> GetTruncatedIndices(size_t len, size_t lenIndices) const;
};

enum EhSolverCancelCheck
{
    ListGeneration,
    ListSorting,
    ListColliding,
    RoundEnd,
    FinalSorting,
    FinalColliding,
    PartialGeneration,
    PartialSorting,
    PartialSubtreeEnd,
    PartialIndexEnd,
    PartialEnd
};

class EhSolverCancelledException : public std::exception
{
    virtual const char* what() const throw() {
        return "Equihash solver was cancelled";
    }
};

inline constexpr const size_t max(const size_t A, const size_t B) { return A > B ? A : B; }

inline constexpr size_t equihash_solution_size(unsigned int N, unsigned int K) {
    return (1 << K)*(N/(K+1)+1)/8;
}

template<unsigned int N, unsigned int K>
class Equihash
{
private:
    BOOST_STATIC_ASSERT(K < N);
    BOOST_STATIC_ASSERT(N % 8 == 0);
    BOOST_STATIC_ASSERT((N/(K+1)) + 1 < 8*sizeof(eh_index));

public:
    enum : size_t { IndicesPerHashOutput=512/N };
    enum : size_t { HashOutput=IndicesPerHashOutput*N/8 };
    enum : size_t { CollisionBitLength=N/(K+1) };
    enum : size_t { CollisionByteLength=(CollisionBitLength+7)/8 };
    enum : size_t { HashLength=(K+1)*CollisionByteLength };
    enum : size_t { FullWidth=2*CollisionByteLength+sizeof(eh_index)*(1 << (K-1)) };
    enum : size_t { FinalFullWidth=2*CollisionByteLength+sizeof(eh_index)*(1 << (K)) };
    enum : size_t { TruncatedWidth=max(HashLength+sizeof(eh_trunc), 2*CollisionByteLength+sizeof(eh_trunc)*(1 << (K-1))) };
    enum : size_t { FinalTruncatedWidth=max(HashLength+sizeof(eh_trunc), 2*CollisionByteLength+sizeof(eh_trunc)*(1 << (K))) };
    enum : size_t { SolutionWidth=(1 << K)*(CollisionBitLength+1)/8 };

    Equihash() { }

    int InitialiseState(eh_HashState& base_state);
#ifdef ENABLE_MINING
    bool BasicSolve(const eh_HashState& base_state,
                    const std::function<bool(std::vector<unsigned char>)> validBlock,
                    const std::function<bool(EhSolverCancelCheck)> cancelled);
    bool OptimisedSolve(const eh_HashState& base_state,
                        const std::function<bool(std::vector<unsigned char>)> validBlock,
                        const std::function<bool(EhSolverCancelCheck)> cancelled);
#endif
    bool IsValidSolution(const eh_HashState& base_state, std::vector<unsigned char> soln);
};

} // namespace

#include "equihash_zcash.tcc"
#include "equihash_zcash.inl"

namespace equihash_zcash {

static Equihash<96,3> Eh96_3;
static Equihash<200,9> Eh200_9;
static Equihash<96,5> Eh96_5;
static Equihash<48,5> Eh48_5;

#define EhZecInitialiseState(n, k, base_state)  \
    if (n == 96 && k == 3) {                 \
        equihash_zcash::Eh96_3.InitialiseState(base_state);  \
    } else if (n == 200 && k == 9) {         \
        equihash_zcash::Eh200_9.InitialiseState(base_state); \
    } else if (n == 96 && k == 5) {          \
        equihash_zcash::Eh96_5.InitialiseState(base_state);  \
    } else if (n == 48 && k == 5) {          \
        equihash_zcash::Eh48_5.InitialiseState(base_state);  \
    } else {                                 \
        throw std::invalid_argument("Unsupported Equihash parameters"); \
    }

#ifdef ENABLE_MINING
inline bool EhBasicSolve(unsigned int n, unsigned int k, const eh_HashState& base_state,
                    const std::function<bool(std::vector<unsigned char>)> validBlock,
                    const std::function<bool(EhSolverCancelCheck)> cancelled)
{
    if (n == 96 && k == 3) {
        return Eh96_3.BasicSolve(base_state, validBlock, cancelled);
    } else if (n == 200 && k == 9) {
        return Eh200_9.BasicSolve(base_state, validBlock, cancelled);
    } else if (n == 96 && k == 5) {
        return Eh96_5.BasicSolve(base_state, validBlock, cancelled);
    } else if (n == 48 && k == 5) {
        return Eh48_5.BasicSolve(base_state, validBlock, cancelled);
    } else {
        throw std::invalid_argument("Unsupported Equihash parameters");
    }
}

inline bool EhBasicSolveUncancellable(unsigned int n, unsigned int k, const eh_HashState& base_state,
                    const std::function<bool(std::vector<unsigned char>)> validBlock)
{
    return EhBasicSolve(n, k, base_state, validBlock,
                        [](EhSolverCancelCheck pos) { return false; });
}

inline bool EhOptimisedSolve(unsigned int n, unsigned int k, const eh_HashState& base_state,
                    const std::function<bool(std::vector<unsigned char>)> validBlock,
                    const std::function<bool(EhSolverCancelCheck)> cancelled)
{
    if (n == 96 && k == 3) {
        return Eh96_3.OptimisedSolve(base_state, validBlock, cancelled);
    } else if (n == 200 && k == 9) {
        return Eh200_9.OptimisedSolve(base_state, validBlock, cancelled);
    } else if (n == 96 && k == 5) {
        return Eh96_5.OptimisedSolve(base_state, validBlock, cancelled);
    } else if (n == 48 && k == 5) {
        return Eh48_5.OptimisedSolve(base_state, validBlock, cancelled);
    } else {
        throw std::invalid_argument("Unsupported Equihash parameters");
    }
}

inline bool EhOptimisedSolveUncancellable(unsigned int n, unsigned int k, const eh_HashState& base_state,
                    const std::function<bool(std::vector<unsigned char>)> validBlock)
{
    return EhOptimisedSolve(n, k, base_state, validBlock,
                            [](EhSolverCancelCheck pos) { return false; });
}
#endif // ENABLE_MINING

#define EhZecIsValidSolution(n, k, base_state, soln, ret)   \
    if (n == 96 && k == 3) {                             \
        ret = equihash_zcash::Eh96_3.IsValidSolution(base_state, soln);  \
    } else if (n == 200 && k == 9) {                     \
        ret = equihash_zcash::Eh200_9.IsValidSolution(base_state, soln); \
    } else if (n == 96 && k == 5) {                      \
        ret = equihash_zcash::Eh96_5.IsValidSolution(base_state, soln);  \
    } else if (n == 48 && k == 5) {                      \
        ret = equihash_zcash::Eh48_5.IsValidSolution(base_state, soln);  \
    } else {                                             \
        throw std::invalid_argument("Unsupported Equihash parameters"); \
    }

} // namespace

#pragma GCC diagnostic pop