#include "../src/mat/impls/aij/seq/aij.h"
#include "petscblaslapack.h"

#undef __FUNCT__
#define __FUNCT__ "MatAXPY_SeqAIJ"
PetscErrorCode
MatAXPY_SeqAIJ(
    Mat *W,
    PetscScalar a,
    Mat X,
    Mat Y)
{
    PetscErrorCode ierr;
    Mat_SeqAIJ *X_aij = (Mat_SeqAIJ*)X->data;
    Mat_SeqAIJ *Y_aij = (Mat_SeqAIJ*)Y->data;
    PetscInt m, m_X, m_Y, n, n_X, n_Y;
    PetscInt *nnz, row;
    PetscInt *X_ai = X_aij->i, *X_aj = X_aij->j, *X_ailen = X_aij->ilen;
    PetscInt *Y_ai = Y_aij->i, *Y_aj = Y_aij->j, *Y_ailen = Y_aij->ilen;
    PetscInt *j_X, *j_Y, *j_X_next, *j_Y_next;
    PetscInt X_len, X_len_max;
    MPI_Comm comm, comm_X, comm_Y;
    PetscScalar *v;
    PetscBLASInt one = 1, bn;

    PetscFunctionBegin;

    /*
     * Ensure matrices are row-oriented.
     */
    if (!Y_aij->roworiented || !X_aij->roworiented)
    {
        SETERRQ(PETSC_ERR_ARG_WRONG, "X and Y both must be row-oriented matrices");
    }

    /*
     * Ensure matrices have conforming sizes.
     */
    ierr = MatGetSize(X, &m_X, &n_X);CHKERRQ(ierr);
    ierr = MatGetSize(Y, &m_Y, &n_Y);CHKERRQ(ierr);
    if (m_X != m_Y)
    {
        SETERRQ(PETSC_ERR_ARG_SIZ, "X and Y must have the same number of rows");
    }
    if (n_X != n_Y)
    {
        SETERRQ(PETSC_ERR_ARG_SIZ, "X and Y must have the same number of cols");
    }
    m = m_X;
    n = n_X;

    /*
     * Ensure matrices share the same MPI communicator.
     */
    ierr = PetscObjectGetComm((PetscObject)X,&comm_X);CHKERRQ(ierr);
    ierr = PetscObjectGetComm((PetscObject)Y,&comm_Y);CHKERRQ(ierr);
    if (comm_X != comm_Y)
    {
        SETERRQ(PETSC_ERR_ARG_WRONG, "X and Y both must share the same MPI communicator");
    }
    comm = comm_X;

    /*
     * Determine non-zero structure of a*X + Y.
     */
    ierr = PetscMalloc(m*sizeof(PetscInt), &nnz);CHKERRQ(ierr);
    X_len_max = 0;
    for (row = 0; row < m; ++row)
    {
        nnz[row] = 0;

        j_X = X_aj + X_ai[row];
        j_Y = Y_aj + Y_ai[row];

        j_X_next = X_aj + X_ai[row+1];
        j_Y_next = Y_aj + Y_ai[row+1];

        X_len = X_ai[row+1] - X_ai[row];
        if (X_len > X_len_max) X_len_max = X_len;

        while (j_X != j_X_next || j_Y != j_Y_next)
        {
            nnz[row] += 1;
            if (j_X != j_X_next && j_Y != j_Y_next)
            {
                if (*j_X < *j_Y)
                {
                    ++j_X;
                }
                else if (*j_Y < *j_X)
                {
                    ++j_X;
                }
                else
                {
                    ++j_X;
                    ++j_Y;
                }
            }
            else if (j_X != j_X_next)
            {
                ++j_X;
            }
            else if (j_Y != j_Y_next)
            {
                ++j_Y;
            }
        }
    }

    /*
     * Create the new matrix.
     */
    ierr = MatCreateSeqAIJ(comm,m,n,PETSC_DEFAULT,nnz,W);CHKERRQ(ierr);
    ierr = PetscFree(nnz);CHKERRQ(ierr);

    /*
     * Accumulate values in the new matrix and assemble it.
     */
    if (a != 0.0 && a != 1.0) ierr = PetscMalloc(X_len_max*sizeof(PetscScalar), &v);CHKERRQ(ierr);
    for (row = 0; row < m; ++row)
    {
        if (a != 0.0)
        {
            if (a == 1.0)
            {
                ierr = MatSetValues(*W,
                                    1, &row,
                                    X_ailen[row], X_aj + X_ai[row],
                                    X_aij->a + X_ai[row],
                                    ADD_VALUES);CHKERRQ(ierr);
            }
            else
            {
                ierr = PetscMemcpy(X_aij->a + X_ai[row], v, X_ailen[row]*sizeof(PetscScalar));CHKERRQ(ierr);
                bn = PetscBLASIntCast(X_ailen[row]);
                BLASscal_(&bn,&a,v,&one);
                ierr = MatSetValues(*W,
                                    1, &row,
                                    X_ailen[row], X_aj + X_ai[row],
                                    v,
                                    ADD_VALUES);CHKERRQ(ierr);
            }
        }
        ierr = MatSetValues(*W,
                            1, &row,
                            Y_ailen[row], Y_aj + Y_ai[row],
                            Y_aij->a + Y_ai[row],
                            ADD_VALUES);CHKERRQ(ierr);
    }
    ierr = MatAssemblyBegin(*W,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
    if (a != 0.0 && a != 1.0) ierr = PetscFree(v);CHKERRQ(ierr);
    ierr = MatAssemblyEnd(*W,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
    PetscFunctionReturn(0);
}