#include <petsc.h>
#include <algorithm>

int main(int argc, char** argv) {
    Mat A, B, C, seq;
    int rank;
    PetscErrorCode ierr;
    PetscInitialize(&argc, &argv, nullptr, nullptr);
    MPI_Comm_size(PETSC_COMM_WORLD, &rank);
    if(rank != 2)
        exit(0);
    MPI_Comm_rank(PETSC_COMM_WORLD, &rank);
    ierr = MatCreate(PETSC_COMM_WORLD, &A);CHKERRQ(ierr);
    ierr = MatSetType(A, MATMPIAIJ);CHKERRQ(ierr);
    ierr = MatSetSizes(A, 1, 1, 2, 2);CHKERRQ(ierr);
    PetscInt ia[2] = { 0, 2 };
    PetscInt ja[2] = { 0, 1 };
    PetscScalar a[2] = { 1.0, 1.0 };
    if(rank == 0) {
        ierr = MatMPIAIJSetPreallocationCSR(A, ia, ja, a);CHKERRQ(ierr);
    }
    else {
        ierr = MatMPIAIJSetPreallocationCSR(A, ia, ja, a);CHKERRQ(ierr);
    }
#if 0
    int lda = 1;
#else
    int lda = 4;
#endif
    PetscScalar* data = new PetscScalar[4 * lda]();
    for (int i = 0; i < 4; ++i)
        std::fill_n(data + lda * i, 1, i * 1.0);
    ierr = MatCreateDense(PETSC_COMM_WORLD, 1, PETSC_DECIDE, 2, 4, data, &B);CHKERRQ(ierr);
    ierr = MatDenseGetLocalMatrix(B, &seq);CHKERRQ(ierr);
    ierr = MatSeqDenseSetLDA(seq, lda);CHKERRQ(ierr);
    ierr = MatView(A, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
    ierr = MatView(B, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
    ierr = MatMatMult(A, B, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &C);CHKERRQ(ierr);
    ierr = MatView(C, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
    ierr = MatDestroy(&B);CHKERRQ(ierr);
    ierr = MatDestroy(&A);CHKERRQ(ierr);
    PetscFinalize();
    return 0;
}