//
//  main.cpp
//

#include "Matrix.hpp"
#include <mpi.h>

#include <cstdlib>
#include <ctime>
#include <iostream>
#include <stdexcept>

using namespace std;

size_t ROOT_PROCESS_RANK = 0;

int ROW_TAG = 1;
int COL_TAG = 2;

// Inverser la matrice par la méthode de Gauss-Jordan; implantation séquentielle.
void invertSequential(Matrix& iA) {

    // vérifier que la matrice est carrée
    assert(iA.rows() == iA.cols());
    // construire la matrice [A I]
    MatrixConcatCols lAI(iA, MatrixIdentity(iA.rows()));

    // traiter chaque rangée
    for (size_t k=0; k<iA.rows(); ++k) {
        // trouver l'index p du plus grand pivot de la colonne k en valeur absolue
        // (pour une meilleure stabilité numérique).
        size_t p = k;
        double lMax = fabs(lAI(k,k));
        for(size_t i = k; i < lAI.rows(); ++i) {
            if(fabs(lAI(i,k)) > lMax) {
                lMax = fabs(lAI(i,k));
                p = i;
            }
        }
        // vérifier que la matrice n'est pas singulière
        if (lAI(p, k) == 0) throw runtime_error("Matrix not invertible");

        // échanger la ligne courante avec celle du pivot

        cout << "P = " << p << ", K = " << k << endl;
        if (p != k) lAI.swapRows(p, k);

        double lValue = lAI(k, k);
        for (size_t j=0; j<lAI.cols(); ++j) {
            // On divise les éléments de la rangée k
            // par la valeur du pivot.
            // Ainsi, lAI(k,k) deviendra égal à 1.
            lAI(k, j) /= lValue;
        }

        // Pour chaque rangée...
        for (size_t i=0; i<lAI.rows(); ++i) {
            if (i != k) { // ...différente de k
                // On soustrait la rangée k
                // multipliée par l'élément k de la rangée courante
                double lValue = lAI(i, k);
                lAI.getRowSlice(i) -= lAI.getRowCopy(k)*lValue;
            }
        }
        cout << lAI.str() <<endl;
    }

    // On copie la partie droite de la matrice AI ainsi transformée
    // dans la matrice courante (this).
    for (unsigned int i=0; i<iA.rows(); ++i) {
        iA.getRowSlice(i) = lAI.getDataArray()[slice(i*lAI.cols()+iA.cols(), iA.cols(), 1)];
    }
}

// Inverser la matrice par la méthode de Gauss-Jordan; implantation MPI parallèle.
void invertParallel(Matrix& iA, size_t rank, size_t size) {

  // vérifier que la matrice est carrée
  assert(iA.rows() == iA.cols());
  // construire la matrice [A I]
  MatrixConcatCols lAI(iA, MatrixIdentity(iA.rows()));

  size_t startRow = rank * lAI.rows() / size;
  size_t endRow = (rank + 1) * lAI.rows() / size;
  for (size_t currentPivotingRow = 0; currentPivotingRow < lAI.rows(); ++currentPivotingRow) {
    double maxAbsValue = 0;
    if (startRow <= currentPivotingRow && currentPivotingRow < endRow){
      // Find maxValue
      // broadcast it
    } else {
      // get the broadcast from the pivoting row
    }

    if (maxAbsValue == 0) {
      // finalize
      if (rank == ROOT_PROCESS_RANK) {
        //throw runtime_error
      } else {
        // exit program
      }
    }

    if (startRow <= currentPivotingRow && currentPivotingRow < endRow){
      // Find maxValue
      // broadcast it
    } else {
      for (size_t currentRow = startRow; currentRow < endRow; ++currentRow) {
        for (size_t currentCol = 0; currentCol < lAi.cols(); ++currentCol) {
          //lAi(currentRow, currentCol) =
        }
      }
    }

  }
}

// Multiplier deux matrices.
Matrix multiplyMatrix(const Matrix& iMat1, const Matrix& iMat2) {

    // vérifier la compatibilité des matrices
    assert(iMat1.cols() == iMat2.rows());
    // effectuer le produit matriciel
    Matrix lRes(iMat1.rows(), iMat2.cols());
    // traiter chaque rangée
    for(size_t i=0; i < lRes.rows(); ++i) {
        // traiter chaque colonne
        for(size_t j=0; j < lRes.cols(); ++j) {
            lRes(i,j) = (iMat1.getRowCopy(i)*iMat2.getColumnCopy(j)).sum();
        }
    }
    return lRes;
}

void createSizeAndOffsetBuffers(int* &sizeBuffer, int* &offsetBuffer, size_t rows, size_t cols, size_t size) {
  sizeBuffer = new int[size];
  offsetBuffer = new int [size];

  for (int i = 0; i < size; ++i) {
    sizeBuffer[i] = cols * ((i + 1) * rows / size - i * rows / size);
    offsetBuffer[i] = cols * (i * rows / size);
  }
}

Matrix broadcastMatrix(size_t rank, size_t size) {
  size_t rows, cols;
  MPI_Bcast(
    &rows,
    sizeof(size_t),
    MPI_CHAR,
    ROOT_PROCESS_RANK,
    MPI_COMM_WORLD);
  MPI_Bcast(
    &cols,
    sizeof(size_t),
    MPI_CHAR,
    ROOT_PROCESS_RANK,
    MPI_COMM_WORLD);

  Matrix broadcastedMatrix(rows, cols);
  MPI_Bcast(
    &broadcastedMatrix(0, 0),
    rows * cols,
    MPI_DOUBLE,
    ROOT_PROCESS_RANK,
    MPI_COMM_WORLD);

  return broadcastedMatrix;
}

void gatherMatrix(Matrix& m, size_t rank, size_t size) {
  MPI_Gatherv(&m(rank * m.rows() / size, 0),
    m.cols() * ((rank + 1) * m.rows() / size - rank * m.rows() / size),
    MPI_DOUBLE,
    NULL,
    NULL,
    NULL,
    MPI_DOUBLE,
    ROOT_PROCESS_RANK,
    MPI_COMM_WORLD);
}

Matrix broadcastRootMatrix(Matrix& m, size_t rank, size_t size) {
  size_t rows = m.rows();
  size_t cols = m.cols();
  MPI_Bcast(
    &rows,
    sizeof(size_t),
    MPI_CHAR,
    rank,
    MPI_COMM_WORLD);
  MPI_Bcast(
    &rows,
    sizeof(size_t),
    MPI_CHAR,
    rank,
    MPI_COMM_WORLD);

  MPI_Bcast(
    &m(0, 0),
    rows * cols,
    MPI_DOUBLE,
    rank,
    MPI_COMM_WORLD);

  return m;
}

Matrix gatherRootMatrix(Matrix& m, size_t rank, size_t size) {
  int* sizeBuffer;
  int* offsetBuffer;

  createSizeAndOffsetBuffers(sizeBuffer, offsetBuffer, m.rows(), m.cols(), size);

  Matrix gatheredMatrix(m.rows(), m.cols());
  MPI_Gatherv(&m(rank * m.rows() / size, 0),
    sizeBuffer[rank],
    MPI_DOUBLE,
    &gatheredMatrix(0, 0),
    sizeBuffer,
    offsetBuffer,
    MPI_DOUBLE,
    ROOT_PROCESS_RANK,
    MPI_COMM_WORLD);

  delete[] sizeBuffer;
  delete[] offsetBuffer;

  return gatheredMatrix;
}

void invertParallelChild(size_t rank, size_t size) {
  Matrix lB(broadcastMatrix(rank, size));
  invertParallel(lB, rank, size);
  gatherMatrix(lB, rank, size);
}

void invertParallelParent(size_t rank, size_t size, int argc, char** argv) {

  unsigned int lS = 5;
  if (argc == 2) {
      lS = atoi(argv[1]);
  }

  MatrixRandom lA(lS, lS);
  cout << "Matrice random:\n" << lA.str() << endl;

  broadcastRootMatrix(lA, rank, size);
  invertParallel(lA, rank, size);
  Matrix lB(gatherRootMatrix(lA, rank, size));
  cout << "Matrice inverse:\n" << lB.str() << endl;

  Matrix lRes = multiplyMatrix(lA, lB);
  cout << "Produit des deux matrices:\n" << lRes.str() << endl;

  cout << "Erreur: " << lRes.getDataArray().sum() - lS << endl;
}

int main(int argc, char** argv) {
    MPI_Init(NULL, NULL);

    srand((unsigned)time(NULL));

    int worldSize;
    MPI_Comm_size(MPI_COMM_WORLD, &worldSize);

    int worldRank;
    MPI_Comm_rank(MPI_COMM_WORLD, &worldRank);

    if (worldRank == ROOT_PROCESS_RANK) {
      invertParallelParent(worldRank, worldSize, argc, argv);
    } else {
      invertParallelChild(worldRank, worldSize);
    }

    MPI_Finalize();
    return 0;
}