/**
 * @file fastpam1.cpp
 * @date 2021-08-03
 *
 * Contains the primary C++ implementation of the FastPAM1 code
 * from the paper: Erich Schubert and Peter J. Rousseeuw: Faster k-Medoids
 * Clustering: Improving the PAM, CLARA, and CLARANS Algorithms.
 * (https://arxiv.org/pdf/1810.05691.pdf). The original PAM papers are: 1)
 * Leonard Kaufman and Peter J. Rousseeuw: Clustering by means of medoids. 2)
 * Leonard Kaufman and Peter J. Rousseeuw: Partitioning around medoids (program
 * pam).
 */

#include "fastpam1.hpp"

#include <armadillo>
#include <unordered_map>

namespace km {
void FastPAM1::fitFastPAM1(
  const arma::fmat &inputData,
  std::optional<std::reference_wrapper<const arma::fmat>> distMat) {
  data = inputData;
  data = arma::trans(data);
  arma::urowvec medoidIndices(nMedoids);
  FastPAM1::buildFastPAM1(data, distMat, &medoidIndices);
  steps = 0;
  medoidIndicesBuild = medoidIndices;
  arma::urowvec assignments(data.n_cols);
  size_t iter = 0;
  bool medoidChange = true;
  while (iter < maxIter && medoidChange) {
    auto previous{medoidIndices};
    if (nMedoids > 1) {
      FastPAM1::swapFastPAM1(data, distMat, &medoidIndices, &assignments);
    }

    medoidChange = arma::any(medoidIndices != previous);
    iter++;
  }
  medoidIndicesFinal = medoidIndices;
  labels = assignments;
  steps = iter;
}

void FastPAM1::buildFastPAM1(
  const arma::fmat &data,
  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
  arma::urowvec *medoidIndices) {
  size_t N = data.n_cols;
  arma::frowvec estimates(N, arma::fill::zeros);
  arma::frowvec bestDistances(N);
  bestDistances.fill(std::numeric_limits<float>::infinity());
  arma::frowvec sigma(N);
  float minDistance = std::numeric_limits<float>::infinity();
  int best = 0;
  float total = 0;
  float cost = 0;

  // TODO(@motiwari): pragma omp parallel for?
  for (size_t k = 0; k < nMedoids; k++) {
    minDistance = std::numeric_limits<float>::infinity();
    best = 0;
    // fixes a base datapoint
    // TODO(@motiwari): pragma omp parallel for?
    for (size_t i = 0; i < data.n_cols; i++) {
      total = 0;
      // TODO(@motiwari): pragma omp parallel for?
      for (size_t j = 0; j < data.n_cols; j++) {
        // computes distance between base and all other points
        cost = (this->*lossFn)(data, i, j);
        // compares this with the cached best distance
        if (bestDistances(j) < cost) {
          cost = bestDistances(j);
        }
        total += cost;
      }
      if (total < minDistance) {
        minDistance = total;
        best = i;
      }
    }
    (*medoidIndices)(k) = best;

    // update the medoid assignment and best_distance for this datapoint
    // TODO(@motiwari): pragma omp parallel for?
    for (size_t l = 0; l < N; l++) {
      cost = (this->*lossFn)(data, l, (*medoidIndices)(k));
      if (cost < bestDistances(l)) {
        bestDistances(l) = cost;
      }
    }
  }
}

void FastPAM1::swapFastPAM1(
  const arma::fmat &data,
  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
  arma::urowvec *medoidIndices, arma::urowvec *assignments) {
  float bestChange = 0;
  float prevBestChange = 0;
  size_t swapIn = 0;
  size_t medoidToSwap = 0;
  size_t N = data.n_cols;
  size_t iter = 0;
  bool swapPerformed = true;
  arma::fmat sigma(nMedoids, N, arma::fill::zeros);
  arma::frowvec bestDistances(N);
  arma::frowvec secondBestDistances(N);
  arma::frowvec deltaTD(nMedoids, arma::fill::zeros);

  // calculate quantities needed for swap, bestDistances and sigma
  KMedoids::calcBestDistancesSwap(data, distMat, medoidIndices, &bestDistances,
                                  &secondBestDistances, assignments,
                                  swapPerformed);

  float di = 0;
  float dij = 0;

  while (swapPerformed && iter < maxIter) {
    bestChange = 0;
    deltaTD.zeros();
    iter++;
    // TODO(@motiwari): pragma omp parallel for?
    for (size_t i = 0; i < data.n_cols; i++) {
      di = bestDistances(i);

      // Consider making point i a medoid.
      // The total loss then contains at least one term, -di,
      // because the loss contribution for point i is
      // reduced from di to 0
      deltaTD.fill(-di);
      // TODO(@motiwari): pragma omp parallel for?
      for (size_t j = 0; j < data.n_cols; j++) {
        if (j != i) {
          dij = (this->*lossFn)(data, i, j);
          if (dij < bestDistances(j)) {
            // Case 1: point i becomes the closest
            // medoid for point j,
            // regardless of which medoid j was
            // previously assigned to. deltaTD
            // will be negative across ALL
            // possible medoid indices m
            deltaTD += (dij - bestDistances(j));
          } else if (dij < secondBestDistances(j)) {
            // Case 2: i. If point i is closer
            // than the second best
            // medoid but further than the best
            // medoid (enforced by failing
            // the condition for the above if
            // condition), point i will
            // become the closest medoid only
            // when we remove its associated
            // medoid and add point i
            deltaTD.at((*assignments)(j)) += (dij - bestDistances(j));
          } else {
            // Case 3: dij > secondBestDistances(j).
            // Then the loss for point j
            // will not change for any medoid swapped
            // out except for its
            // assignment, in which case it moves to
            // its second nearest medoid
            deltaTD.at((*assignments)(j)) +=
              (secondBestDistances(j) - bestDistances(j));
          }
        }
      }

      // Determine the best medoid to swap out
      arma::uword swapOut = deltaTD.index_min();
      // If the loss change is better than
      // the best loss change so far,
      // Update our running best statistics
      if (deltaTD.min() < prevBestChange) {
        bestChange = deltaTD.min();
        prevBestChange = bestChange;
        swapIn = i;
        medoidToSwap = swapOut;
      }
    }

    // Update the loss and perform the swap if
    // the loss would be improved
    if (bestChange < 0) {
      (*medoidIndices)(medoidToSwap) = swapIn;
      calcBestDistancesSwap(data, distMat, medoidIndices, &bestDistances,
                            &secondBestDistances, assignments, swapPerformed);
    } else {
      swapPerformed = false;
    }
  }
}
}  // namespace km
