#include "HybridSolveHandler.h"

#include "CallAndTimeBlas.h"
#include "DataCollector.h"
#include "FactorHiGHSSettings.h"
#include "FormatHandler.h"
#include "Swaps.h"
#include "ipm/hipo/auxiliary/Auxiliary.h"

namespace hipo {

HybridSolveHandler::HybridSolveHandler(
    const Symbolic& S, const std::vector<std::vector<double>>& sn_columns,
    const std::vector<std::vector<Int>>& swaps,
    const std::vector<std::vector<double>>& pivot_2x2, DataCollector& data)
    : SolveHandler(S, sn_columns, data), swaps_{swaps}, pivot_2x2_{pivot_2x2} {}

void HybridSolveHandler::forwardSolve(std::vector<double>& x) const {
  // Forward solve.
  // Blas calls: dtrsv, dgemv

  // supernode columns in format FH

  HIPO_CLOCK_CREATE;

  const Int nb = S_.blockSize();

  for (Int sn = 0; sn < S_.sn(); ++sn) {
    // leading size of supernode
    const Int ldSn = S_.ptr(sn + 1) - S_.ptr(sn);

    // number of columns in the supernode
    const Int sn_size = S_.snStart(sn + 1) - S_.snStart(sn);

    // first colums of the supernode
    const Int sn_start = S_.snStart(sn);

    // index to access S->rows for this supernode
    const Int64 start_row = S_.ptr(sn);

    // number of blocks of columns
    const Int n_blocks = (sn_size - 1) / nb + 1;

    // index to access snColumns[sn]
    Int64 SnCol_ind{};

    // go through blocks of columns for this supernode
    for (Int j = 0; j < n_blocks; ++j) {
      // number of columns in the block
      const Int jb = std::min(nb, sn_size - nb * j);

      // number of entries in diagonal part
      const Int diag_entries = jb * jb;

      // index to access vector x
      const Int x_start = sn_start + nb * j;

#ifdef HIPO_PIVOTING
      HIPO_CLOCK_START(2);
      // apply swaps to portion of rhs that is affected
      const Int* current_swaps = &swaps_[sn][nb * j];
      permuteWithSwaps(&x[x_start], current_swaps, jb);
      HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_swap);
#endif

      HIPO_CLOCK_START(2);
      callAndTime_dtrsv('U', 'T', 'U', jb, &sn_columns_[sn][SnCol_ind], jb,
                        &x[x_start], 1, data_);

      SnCol_ind += diag_entries;

      // temporary space for gemv
      const Int gemv_space = ldSn - nb * j - jb;
      std::vector<double> y(gemv_space);
      if (gemv_space > 0) {
        callAndTime_dgemv('T', jb, gemv_space, 1.0, &sn_columns_[sn][SnCol_ind],
                          jb, &x[x_start], 1, 0.0, y.data(), 1, data_);
        SnCol_ind += jb * gemv_space;
        HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_dense);

        HIPO_CLOCK_START(2);
        // scatter solution of gemv
        for (Int i = 0; i < gemv_space; ++i) {
          const Int row = S_.rows(start_row + nb * j + jb + i);
          x[row] -= y[i];
        }
        HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_sparse);
      }

#ifdef HIPO_PIVOTING
      HIPO_CLOCK_START(2);
      // apply inverse swaps
      permuteWithSwaps(&x[x_start], current_swaps, jb, true);
      HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_swap);
#endif
    }
  }
}

void HybridSolveHandler::backwardSolve(std::vector<double>& x) const {
  // Backward solve.
  // Blas calls: dtrsv, dgemv

  // supernode columns in format FH

  HIPO_CLOCK_CREATE;

  const Int nb = S_.blockSize();

  // go through the sn in reverse order
  for (Int sn = S_.sn() - 1; sn >= 0; --sn) {
    // leading size of supernode
    const Int ldSn = S_.ptr(sn + 1) - S_.ptr(sn);

    // number of columns in the supernode
    const Int sn_size = S_.snStart(sn + 1) - S_.snStart(sn);

    // first colums of the supernode
    const Int sn_start = S_.snStart(sn);

    // index to access S->rows for this supernode
    const Int64 start_row = S_.ptr(sn);

    // number of blocks of columns
    const Int n_blocks = (sn_size - 1) / nb + 1;

    // index to access snColumns[sn]
    // initialised with the total number of entries of snColumns[sn]
    Int64 SnCol_ind = sn_columns_[sn].size() - extra_space_frontal;

    // go through blocks of columns for this supernode in reverse order
    for (Int j = n_blocks - 1; j >= 0; --j) {
      // number of columns in the block
      const Int jb = std::min(nb, sn_size - nb * j);

      // number of entries in diagonal part
      const Int diag_entries = jb * jb;

      // index to access vector x
      const Int x_start = sn_start + nb * j;

#ifdef HIPO_PIVOTING
      HIPO_CLOCK_START(2);
      // apply swaps to portion of rhs that is affected
      const Int* current_swaps = &swaps_[sn][nb * j];
      permuteWithSwaps(&x[x_start], current_swaps, jb);
      HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_swap);
#endif

      // temporary space for gemv
      const Int gemv_space = ldSn - nb * j - jb;
      std::vector<double> y(gemv_space);
      if (gemv_space > 0) {
        HIPO_CLOCK_START(2);
        // scatter entries into y
        for (Int i = 0; i < gemv_space; ++i) {
          const Int row = S_.rows(start_row + nb * j + jb + i);
          y[i] = x[row];
        }
        HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_sparse);

        HIPO_CLOCK_START(2);
        SnCol_ind -= jb * gemv_space;
        callAndTime_dgemv('N', jb, gemv_space, -1.0,
                          &sn_columns_[sn][SnCol_ind], jb, y.data(), 1, 1.0,
                          &x[x_start], 1, data_);
        HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_dense);
      }

      HIPO_CLOCK_START(2);
      SnCol_ind -= diag_entries;
      callAndTime_dtrsv('U', 'N', 'U', jb, &sn_columns_[sn][SnCol_ind], jb,
                        &x[x_start], 1, data_);
      HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_dense);

#ifdef HIPO_PIVOTING
      HIPO_CLOCK_START(2);
      // apply inverse swaps
      permuteWithSwaps(&x[x_start], current_swaps, jb, true);
      HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_swap);
#endif
    }
  }
}

void HybridSolveHandler::diagSolve(std::vector<double>& x) const {
  // Diagonal solve

  // supernode columns in format FH

  HIPO_CLOCK_CREATE;

  const Int nb = S_.blockSize();

  for (Int sn = 0; sn < S_.sn(); ++sn) {
    // leading size of supernode
    const Int ldSn = S_.ptr(sn + 1) - S_.ptr(sn);

    // number of columns in the supernode
    const Int sn_size = S_.snStart(sn + 1) - S_.snStart(sn);

    // first colums of the supernode
    const Int sn_start = S_.snStart(sn);

    // number of blocks of columns
    const Int n_blocks = (sn_size - 1) / nb + 1;

    // index to access diagonal part of block
    Int diag_start{};

    // go through blocks of columns for this supernode
    for (Int j = 0; j < n_blocks; ++j) {
      // number of columns in the block
      const Int jb = std::min(nb, sn_size - nb * j);

#ifdef HIPO_PIVOTING
      HIPO_CLOCK_START(2);
      // apply swaps to portion of rhs that is affected
      const Int* current_swaps = &swaps_[sn][nb * j];
      permuteWithSwaps(&x[sn_start + nb * j], current_swaps, jb);
      HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_swap);
#endif

      HIPO_CLOCK_START(2);

      const double* current_2x2 = &pivot_2x2_[sn][nb * j];
      Int step = 1;

      // go through columns of block
      for (Int col = 0; col < jb; col += step) {
        if (current_2x2[col] == 0.0) {
          // 1x1 pivots
          step = 1;
          const double inv_d = sn_columns_[sn][diag_start + col + jb * col];
          x[sn_start + nb * j + col] *= inv_d;
        } else {
          // 2x2 pivots
          step = 2;

          // inverse of 2x2 pivot
          const double i_d1 = sn_columns_[sn][diag_start + col + jb * col];
          const double i_d2 =
              sn_columns_[sn][diag_start + col + 1 + jb * (col + 1)];
          const double i_off = current_2x2[col];

          double x1 = x[sn_start + nb * j + col];
          double x2 = x[sn_start + nb * j + col + 1];

          x[sn_start + nb * j + col] = i_d1 * x1 + i_off * x2;
          x[sn_start + nb * j + col + 1] = i_d2 * x2 + i_off * x1;
        }
      }

      HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_dense);

#ifdef HIPO_PIVOTING
      HIPO_CLOCK_START(2);
      // apply inverse swaps
      permuteWithSwaps(&x[sn_start + nb * j], current_swaps, jb, true);
      HIPO_CLOCK_STOP(2, data_, kTimeSolveSolve_swap);
#endif

      // move diag_start forward by number of diagonal entries in block
      diag_start += jb * jb;

      // move diag_start forward by number of sub-diagonal entries in block
      diag_start += (ldSn - nb * j - jb) * jb;
    }
  }
}

}  // namespace hipo
