#include <iostream>
#include <petsc.h>
#include <mpi.h>

typedef struct _p_Mesh* Mesh;
struct _p_Mesh{
	Mat  A, B, C, BT, CT;
	Mat  system;
	Vec  f1, f2, p;
	Vec  rhs, solution;
};


int main(int argc, char* argv[]){
	PetscInitialize(&argc, &argv, NULL, NULL);

	Mesh testMesh;
	PetscNew(&testMesh);

	//create matrix A, B, C and BT, CT
	MatCreate(PETSC_COMM_WORLD, &testMesh->A);
	MatSetSizes(testMesh->A, PETSC_DECIDE, PETSC_DECIDE, 5, 5);
	MatSetFromOptions(testMesh->A);
	MatSetUp(testMesh->A);
	for (int i = 0; i < 5; ++i){
		MatSetValue(testMesh->A, i, i, 1, INSERT_VALUES);
	}
	MatAssemblyBegin(testMesh->A,MAT_FINAL_ASSEMBLY);
	MatAssemblyEnd(testMesh->A,MAT_FINAL_ASSEMBLY);

	MatCreate(PETSC_COMM_WORLD, &testMesh->B);
	MatSetSizes(testMesh->B, PETSC_DECIDE, PETSC_DECIDE, 5, 5);
	MatSetFromOptions(testMesh->B);
	MatSetUp(testMesh->B);
	for (int i = 0; i < 5; ++i){
		MatSetValue(testMesh->B, i, i, 1, INSERT_VALUES);
	}
	MatAssemblyBegin(testMesh->B,MAT_FINAL_ASSEMBLY);
	MatAssemblyEnd(testMesh->B,MAT_FINAL_ASSEMBLY);


	MatCreate(PETSC_COMM_WORLD, &testMesh->C);
	MatSetSizes(testMesh->C, PETSC_DECIDE, PETSC_DECIDE, 5, 5);
	MatSetFromOptions(testMesh->C);
	MatSetUp(testMesh->C);
	for (int i = 0; i < 5; ++i){
		MatSetValue(testMesh->C, i, i, 1, INSERT_VALUES);
	}
	MatAssemblyBegin(testMesh->C,MAT_FINAL_ASSEMBLY);
	MatAssemblyEnd(testMesh->C,MAT_FINAL_ASSEMBLY);

	MatTranspose(testMesh->B, MAT_INITIAL_MATRIX, &testMesh->BT);
	MatTranspose(testMesh->C, MAT_INITIAL_MATRIX, &testMesh->CT);

	//Create nested Matrix
	Mat mblock[9] = {testMesh->A, NULL, testMesh->BT, NULL, testMesh->A, testMesh->CT, testMesh->B, testMesh->C, NULL};
	MatCreateNest(PETSC_COMM_WORLD, 3, NULL, 3, NULL, mblock, &testMesh->system);
	MatView(testMesh->system, PETSC_VIEWER_STDOUT_WORLD);

	//Create vectors
	VecCreate(PETSC_COMM_WORLD, &testMesh->f1);
	VecSetSizes(testMesh->f1, PETSC_DECIDE, 5);
	VecSetFromOptions(testMesh->f1);
	VecSetUp(testMesh->f1);
	for (int i = 0; i < 5; ++i){
		VecSetValue(testMesh->f1, i, 1, INSERT_VALUES);
	}
	VecAssemblyBegin(testMesh->f1);
	VecAssemblyEnd(testMesh->f1);

	VecCreate(PETSC_COMM_WORLD, &testMesh->f2);
	VecSetSizes(testMesh->f2, PETSC_DECIDE, 5);
	VecSetFromOptions(testMesh->f2);
	VecSetUp(testMesh->f2);
	for (int i = 0; i < 5; ++i){
		VecSetValue(testMesh->f2, i, 1, INSERT_VALUES);
	}
	VecAssemblyBegin(testMesh->f2);
	VecAssemblyEnd(testMesh->f2);

	VecCreate(PETSC_COMM_WORLD, &testMesh->p);
	VecSetSizes(testMesh->p, PETSC_DECIDE, 5);
	VecSetFromOptions(testMesh->p);
	VecSetUp(testMesh->p);
	for (int i = 0; i < 5; ++i){
		VecSetValue(testMesh->p, i, 1, INSERT_VALUES);
	}
	VecAssemblyBegin(testMesh->p);
	VecAssemblyEnd(testMesh->p);

	//create nested vector
	Vec vblock[3] = {testMesh->f1, testMesh->f2, testMesh->p};
	VecCreateNest(PETSC_COMM_WORLD, 3, NULL, vblock, &testMesh->rhs);
	VecView(testMesh->rhs, PETSC_VIEWER_STDOUT_WORLD);

	VecDuplicate(testMesh->rhs, &testMesh->solution);

	//try to solve and encounter error
	KSP ksp;
	PC  pc;
	KSPCreate(PETSC_COMM_WORLD, &ksp);
	KSPSetType(ksp, KSPGMRES);
	KSPSetOperators(ksp, testMesh->system, testMesh->system);
	KSPGetPC(ksp, &pc);
	PCSetType(pc, PCJACOBI);
	KSPSetUp(ksp);
	KSPSolve(ksp, testMesh->rhs, testMesh->solution);

	PetscFinalize();
	return 0;
}