#include "petscmat.h"

static char help[]="A demo program";

typedef struct{
  PetscInt dim;     // Dimensions 
  PetscInt m;       // Number of data points 
  PetscInt maxnnz;  // Maximum number of nonzeros in any row
  Mat      data;    // Store data
} AppCtx;

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc,char **args)
{
  AppCtx         user; 
  PetscErrorCode info;
  PetscViewer    view;
  char           bdata_path[PETSC_MAX_PATH_LEN];
  PetscTruth     bdata_flg; 
  PetscScalar    *vals;
  PetscInt       *idxs;

  PetscInt       zero=0;
  PetscInt       one=1;
  
  user.m=20000000;

  // svnvish: BUGBUG
  // do not change these values
  user.dim=800;
  user.maxnnz=200;
  
  PetscInitialize(&argc,&args,(char *)0,help);

  info=PetscOptionsGetString(PETSC_NULL,"-data",bdata_path,PETSC_MAX_PATH_LEN-1,&bdata_flg);CHKERRQ(info);
  
  // Create data matrix 
  info=MatCreate(PETSC_COMM_WORLD,&user.data);CHKERRQ(info);
  info=MatSetSizes(user.data,PETSC_DECIDE,PETSC_DECIDE,user.m,user.dim);CHKERRQ(info);
  info=MatSetFromOptions(user.data);CHKERRQ(info);
  info=MatSetType(user.data,MATMPIAIJ);
  info=PetscObjectSetName((PetscObject) user.data,"Data");CHKERRQ(info);

  info=PetscPrintf(PETSC_COMM_WORLD,"\n Create finished \n");CHKERRQ(info);

  // Allocate space 
  PetscInt begin,end;
  info=MatGetOwnershipRange(user.data,&begin,&end);CHKERRQ(info);
  PetscInt locm=end-begin;

  PetscInt cbegin,cend;
  MatGetOwnershipRangeColumn(user.data,&cbegin,&cend);CHKERRQ(info);

  PetscInt *diag_nnz;
  PetscInt *offdiag_nnz;
  info=PetscMalloc(locm*sizeof(PetscInt),&diag_nnz); CHKERRQ(info);
  info=PetscMalloc(locm*sizeof(PetscInt),&offdiag_nnz); CHKERRQ(info);
  
  for (PetscInt ii=0; ii<locm; ii++){
    diag_nnz[ii]=0;
    for (PetscInt idx=0; idx<user.dim; idx=idx+4){
      if (idx>=cbegin && idx<cend)
        diag_nnz[ii]++;
    }
    offdiag_nnz[ii]=user.maxnnz-diag_nnz[ii];
  }
  
  info=MatMPIAIJSetPreallocation(user.data,zero,diag_nnz,zero,offdiag_nnz);CHKERRQ(info);

  info=PetscPrintf(PETSC_COMM_WORLD,"\n Preallocate finished \n");CHKERRQ(info);

  info=PetscMalloc(user.maxnnz*sizeof(PetscScalar),&vals);CHKERRQ(info);
  info=PetscMalloc(user.maxnnz*sizeof(PetscInt),&idxs);CHKERRQ(info);
  PetscInt j=0;
  for (PetscInt idx=0; idx<user.dim; idx=idx+4){
    idxs[j]=idx;
    vals[j]=200.0;
    j++;
  }
  
  for (PetscInt ii=begin; ii<end; ii++)
    info=MatSetValues(user.data,one,&ii,user.maxnnz,idxs,vals,INSERT_VALUES);CHKERRQ(info);

  info=PetscPrintf(PETSC_COMM_WORLD,"\n setting values finished \n");CHKERRQ(info);
  
  info=MatAssemblyBegin(user.data,MAT_FINAL_ASSEMBLY);CHKERRQ(info);
  info=MatAssemblyEnd(user.data,MAT_FINAL_ASSEMBLY);CHKERRQ(info);
    
  // Write the data and labels in Petsc binary file
  info=PetscPrintf(PETSC_COMM_WORLD,"\n Writing data in binary format to %s \n",bdata_path);CHKERRQ(info);
  info=PetscViewerBinaryOpen(PETSC_COMM_WORLD,bdata_path,FILE_MODE_WRITE,&view);CHKERRQ(info);
  info=MatView(user.data,view);CHKERRQ(info);
  
  info=PetscPrintf(PETSC_COMM_WORLD,"\n writing finished \n");CHKERRQ(info);
  
  info=PetscViewerDestroy(view);CHKERRQ(info);
  info=MatDestroy(user.data);CHKERRQ(info);
  info=PetscFree(diag_nnz);CHKERRQ(info);
  info=PetscFree(offdiag_nnz);CHKERRQ(info);
  info=PetscFree(vals);CHKERRQ(info);
  info=PetscFree(idxs);CHKERRQ(info);
  
  info=PetscFinalize();CHKERRQ(info);
  return 0;
}