#include <petscsys.h>
#include <petscmat.h>
#include <petscvec.h>
#include <petscts.h>
#include <iostream>
#include <cstdlib>
#include <cmath>

extern "C" {

PetscErrorCode monitor(TS ts, PetscInt steps, PetscReal time, Vec u, void *mctx);

PetscErrorCode tri_diag_mat(
	PetscInt n,
	PetscScalar a_diag,
	PetscScalar a_L,
	PetscScalar a_U,
	Mat* A_ptr);

PetscErrorCode initial_solution_vector(
	PetscInt n,
	Vec* u_ptr);

PetscErrorCode F_eval(TS ts, PetscReal t, Vec u, Vec u_t, Vec F, void* ctx);
PetscErrorCode Jac_F_eval(TS ts, PetscReal t, Vec u, Vec u_t, PetscReal a, Mat *A, Mat *B, MatStructure *flag, void *ctx);

} // extern "C"

struct ts_ctx
{
	Mat A_lhs, A_rhs;
};

int main(int argc, char* args[])
{
	PetscErrorCode ierr = PetscInitialize(&argc, &args, 0, 0); CHKERRQ(ierr);

	int comm_size;
	ierr = MPI_Comm_size(PETSC_COMM_WORLD, &comm_size); CHKERRQ(ierr);

	if(comm_size > 1)
	{
		std::cerr << "ts_test is for single processor only.\n";
		exit(1);
	}

	PetscInt n_elems = 100;
	PetscReal dt = 0.001;
	PetscReal ftime = -1;

	TS ts;
	ts_ctx ctx;
	Vec u, residual;
	Mat A_combine;
	ierr = tri_diag_mat(n_elems - 1, 4.0/(3*n_elems), 1.0/(3*n_elems), 1.0/(3*n_elems), &ctx.A_lhs); CHKERRQ(ierr);
	ierr = tri_diag_mat(n_elems - 1, n_elems, -n_elems/2.0, -n_elems/2.0, &ctx.A_rhs); CHKERRQ(ierr);

	ierr = MatDuplicate(ctx.A_rhs, MAT_DO_NOT_COPY_VALUES, &A_combine); CHKERRQ(ierr);

	ierr = initial_solution_vector(n_elems - 1, &u); CHKERRQ(ierr);

	//ierr = VecView(u, PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);

	ierr = VecDuplicate(u, &residual); CHKERRQ(ierr);

	ierr = TSCreate(PETSC_COMM_WORLD, &ts); CHKERRQ(ierr);
    ierr = TSSetProblemType(ts, TS_NONLINEAR);CHKERRQ(ierr);

    const TSType type = TSTHETA;
    ierr = TSSetType(ts, type); CHKERRQ(ierr);

    ierr = TSSetIFunction(ts, residual, F_eval, &ctx); CHKERRQ(ierr);
    ierr = TSSetIJacobian(ts, A_combine, A_combine, Jac_F_eval, &ctx); CHKERRQ(ierr);

	ierr = TSSetDuration(ts,10000,10); CHKERRQ(ierr);
	ierr = TSSetInitialTimeStep(ts, 0.0, dt); CHKERRQ(ierr);
	ierr = TSSetSolution(ts, u); CHKERRQ(ierr);
	ierr = TSSetFromOptions(ts); CHKERRQ(ierr);
	ierr = TSMonitorSet(ts, monitor, PETSC_NULL, PETSC_NULL); CHKERRQ(ierr);
	ierr = TSSolve(ts, u, &ftime); CHKERRQ(ierr);

	//ierr = VecView(u, PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);

	ierr = VecDestroy(&u); CHKERRQ(ierr);
	ierr = VecDestroy(&residual); CHKERRQ(ierr);

	ierr = MatDestroy(&ctx.A_lhs); CHKERRQ(ierr);
	ierr = MatDestroy(&ctx.A_rhs); CHKERRQ(ierr);
	ierr = MatDestroy(&A_combine); CHKERRQ(ierr);
	ierr = TSDestroy(&ts); CHKERRQ(ierr);

	ierr = PetscFinalize(); CHKERRQ(ierr);
	return 0;
}

PetscErrorCode tri_diag_mat(
	PetscInt n,         // size of matrix to create
	PetscScalar a_diag, // value to place on diagonal
	PetscScalar a_L,    // value to place on lower diagonal
	PetscScalar a_U,    // value to place on upper diagonal
	Mat* A_ptr)         // output matrix
{
	// compute tridiagonal matrix with each diagonal having
	// the given identical values.

	PetscErrorCode ierr;

	Mat& A = *A_ptr;

	ierr = MatCreate(PETSC_COMM_WORLD, &A); CHKERRQ(ierr);
	ierr = MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, n, n); CHKERRQ(ierr);
	ierr = MatSetFromOptions(A); CHKERRQ(ierr);
	ierr = MatSetUp(A); CHKERRQ(ierr);

    if(n == 1)
    {
		// special case
		ierr = MatSetValue(A, 0, 0, a_diag, INSERT_VALUES); CHKERRQ(ierr);
    }
	else if(n > 1)
	{
		ierr = MatSetValue(A, 0, 0, a_diag, INSERT_VALUES); CHKERRQ(ierr);
		ierr = MatSetValue(A, 0, 1, a_U,    INSERT_VALUES); CHKERRQ(ierr);
    
		PetscScalar row_vals[3] = {a_L, a_diag, a_U};
    
		for(PetscInt row = 1; row < n - 1; ++row)
		{
			PetscInt col_ids[3] = {row-1, row, row+1};
			ierr = MatSetValues(A, 1, &row, 3, col_ids, row_vals, INSERT_VALUES); CHKERRQ(ierr);
		}
    
		ierr = MatSetValue(A, n-1, n-2, a_L,    INSERT_VALUES); CHKERRQ(ierr);
		ierr = MatSetValue(A, n-1, n-1, a_diag, INSERT_VALUES); CHKERRQ(ierr);
	}

	ierr = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
	ierr = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);

	return ierr;
}

PetscErrorCode initial_solution_vector(
	PetscInt n,   // size of vector
	Vec* u_ptr)   // output vector
{
	// evaluate sin(pi x) in [-1, 1] at n uniformly distributed interior
	// points (that is exclude boundary points)

	PetscErrorCode ierr;

	Vec& u = *u_ptr;
	ierr = VecCreate(PETSC_COMM_WORLD, &u); CHKERRQ(ierr);
	ierr = VecSetSizes(u, PETSC_DECIDE, n); CHKERRQ(ierr);
	ierr = VecSetFromOptions(u); CHKERRQ(ierr);

	const double pi = 4*std::atan(1.0);

	const double h = 2.0/(n+1);
	double x = -1 + h;
	for(PetscInt i = 0; i < n; ++i, x += h)
	{
		const double val = std::sin(pi*x);
		ierr = VecSetValue(u, i, val, INSERT_VALUES);
	}

	ierr = VecAssemblyBegin(u); CHKERRQ(ierr);
	ierr = VecAssemblyEnd(u); CHKERRQ(ierr);

	return ierr;
}

PetscErrorCode F_eval(TS ts, PetscReal t, Vec u, Vec u_t, Vec F, void* ctx)
{
	ts_ctx& ctx_obj = *reinterpret_cast<ts_ctx*>(ctx);
	Mat A_lhs = ctx_obj.A_lhs, A_rhs = ctx_obj.A_rhs;
	PetscErrorCode ierr;
	ierr = MatMult(A_lhs, u_t, F); CHKERRQ(ierr);
	ierr = MatMultAdd(A_rhs, u, F, F); CHKERRQ(ierr);
	return ierr;
}

PetscErrorCode Jac_F_eval(TS ts, PetscReal t, Vec u, Vec u_t, PetscReal a, Mat *A, Mat *B, MatStructure *flag, void *ctx)
{
	ts_ctx& ctx_obj = *reinterpret_cast<ts_ctx*>(ctx);
	Mat A_lhs = ctx_obj.A_lhs, A_rhs = ctx_obj.A_rhs;
	PetscErrorCode ierr;
	*flag = SAME_PRECONDITIONER;
	
	ierr = MatCopy(A_rhs, *A, SAME_NONZERO_PATTERN); CHKERRQ(ierr);
	ierr = MatAXPY(*A, a, A_lhs, SAME_NONZERO_PATTERN); CHKERRQ(ierr);
	*B = *A;
	return ierr;
}

PetscErrorCode monitor(TS ts, PetscInt steps, PetscReal time, Vec u, void *mctx)
{
	PetscErrorCode ierr;
    PetscReal u_inf = -1;
	ierr = VecNorm(u, NORM_INFINITY, &u_inf); CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "log10(u_inf) = %25.15g\n", log10(u_inf));
	return ierr;
}