/*
bicgstab.cu (CUDA+MPI)
*/

#include "oth.h"
#include "oth_mpi.h"
#include "cuda_complex.h"
#include "oth_prototype.h"

void bicgstab(int ifreq, int64_t n, cuDoubleComplex *a, cuDoubleComplex *x, cuDoubleComplex *b, cuDoubleComplex *r, cuDoubleComplex *p, cuDoubleComplex *q, cuDoubleComplex *t, cuDoubleComplex *u, cuDoubleComplex *r0, int *niter, double *residual, double *eaverage, FILE *fp, int iout)
{
	const char msg1[] = "  --- converged ---";
	const char msg2[] = "  --- diverged ---";
	const char msg3[] = "  --- not converged ---";
	const char msg4[] = "  --- devided by zero ---";
	const double inf = 1e14;
	char str[BUFSIZ];

	// n0 <= (net index) < n1
	int64_t n0 = NEY(i0Ey, j0Ey, k0Ey);
	int64_t n1 = ((commRank < commSize - 1) ? NEX(i1Ex, j1Ex, k1Ex) : NEZ(i1Ez, j1Ez, k1Ez)) + 1;
	//printf("%d %d n0=%zd n1=%zd %zd %zd %zd %zd %zd %zd\n", commSize, commRank, n0, n1, NEX(i0Ex, j0Ex, k0Ex), NEY(i0Ey, j0Ey, k0Ey), NEZ(i0Ez, j0Ez, k0Ez), NEX(i1Ex, j1Ex, k1Ex), NEY(i1Ey, j1Ey, k1Ey), NEZ(i1Ez, j1Ez, k1Ez)); fflush(stdout);

	// RHS norm
	double bnrm = Dznrm2(n1 - n0, &b[n0]);
	bnrm = sqrt(comm_dsum(bnrm * bnrm));
	//printf("%d %d %f\n", commSize, commRank, b2); fflush(stdout);

	// == initial value ==

	// q = Ax
	copyV(x);
	prodmv(0, ifreq, a, x, q, NULL);

	// r = b - Ax = b - q
	Zcopy(n, b, r);
	Zaxpy(n, make_cuDoubleComplex(-1, 0), q, r);

	// r0 <- b
	r0 = b;

	// r0*
	if (!Iseed) {
		// r0* = r
		Zcopy(n, r, r0);
	}
	else {
		// r0* = p (work array)
		Zcopy(n, p, r0);
	}

	// p = r
	Zcopy(n, r, p);

	// dot1 = (r0*, r)
	cuDoubleComplex l_dot1 = Zdotc(n1 - n0, &r0[n0], &r[n0]);
	cuDoubleComplex dot1 = make_cuDoubleComplex(comm_dsum(l_dot1.x), comm_dsum(l_dot1.y));

	// == iteration ==

	int iter;
	for (iter = 0; iter <= Solver.maxiter; iter++) {
		// update

		// q = Ap
		copyV(p);
		prodmv(0, ifreq, a, p, q, NULL);

		// a = (r0*, r) / (r0*, Ap) = dot1 / (r0*, q)
		cuDoubleComplex l_dot2 = Zdotc(n1 - n0, &r0[n0], &q[n0]);
		cuDoubleComplex dot2 = make_cuDoubleComplex(comm_dsum(l_dot2.x), comm_dsum(l_dot2.y));
		if (cuCabs(dot2) == 0) {
			if (iout) {
				sprintf(str, "%s (1) (%d %e %e)", msg4, iter, cuCabs(dot1), cuCabs(dot2));
				monitor1(fp, str);
			}
			break;
		}
		cuDoubleComplex acg = cuCdiv(dot1, dot2);

		// t = r - aAp = r - aq
		Zaxpy(n, make_cuDoubleComplex(-acg.x, -acg.y), q, r);
		Zcopy(n, r, t);

		// u = At
		copyV(t);
		prodmv(0, ifreq, a, t, u, NULL);

		// z = (At, t) / (At, At) = (u, t) / (u, u)
		cuDoubleComplex l_dot3 = Zdotc(n1 - n0, &u[n0], &t[n0]);
		cuDoubleComplex l_dot4 = Zdotc(n1 - n0, &u[n0], &u[n0]);
		cuDoubleComplex dot3 = make_cuDoubleComplex(comm_dsum(l_dot3.x), comm_dsum(l_dot3.y));
		cuDoubleComplex dot4 = make_cuDoubleComplex(comm_dsum(l_dot4.x), 0);
		if (cuCabs(dot4) == 0) {
			if (iout) {
				sprintf(str, "%s (2) (%d %e %e)", msg4, iter, cuCabs(dot3), cuCabs(dot4));
				monitor1(fp, str);
			}
			break;
		}
		cuDoubleComplex zcg = cuCdiv(dot3, dot4);

		// x = x + ap + zt
/*
		Zaxpy(n, acg, p, x);
		Zaxpy(n, zcg, t, x);
*/
		Zaxbycz(n, make_cuDoubleComplex(1, 0), acg, zcg, x, p, t);

		// r = t - zAt = t - zu
		Zaxpy(n, make_cuDoubleComplex(-zcg.x, -zcg.y), u, t);
		Zcopy(n, t, r);

		// residual

		double xnrm = Dznrm2(n1 - n0, &x[n0]);
		xnrm = sqrt(comm_dsum(xnrm * xnrm));
		eaverage[iter] = xnrm;

		double rnrm = Dznrm2(n1 - n0, &r[n0]);
		rnrm = sqrt(comm_dsum(rnrm * rnrm));
		residual[iter] = rnrm / bnrm;

		// monitor

		if ((iter % Solver.nout) == 0 || (residual[iter] < Solver.converg)) {
			if (iout) {
				monitor6(fp, iter, residual[iter], eaverage[iter]);
			}
		}

		// converged / diverged / not converged

		if ((iter > 0) && (residual[iter] < Solver.converg)) {
			if (iout) {
				monitor1(fp, msg1);
			}
			break;
		}
		if (residual[iter] > inf) {
			if (iout) {
				monitor1(fp, msg2);
			}
			break;
		}
		if (iter == Solver.maxiter) {
			if (iout) {
				monitor1(fp, msg3);
			}
			break;
		}

		// update

		// b = (a / z) (r0*, r[k+1]) / (r0*, r[k]) = (a / z) (r0*, r) / dot1
		cuDoubleComplex l_dot5 = Zdotc(n1 - n0, &r0[n0], &r[n0]);
		cuDoubleComplex dot5 = make_cuDoubleComplex(comm_dsum(l_dot5.x), comm_dsum(l_dot5.y));
		if (cuCabs(cuCmul(zcg, dot1)) == 0) {
			if (iout) {
				sprintf(str, "%s (3) (%d %.3e %.3e %.3e %.3e)", msg4, iter, cuCabs(acg), cuCabs(zcg), cuCabs(dot1), cuCabs(dot5));
				monitor1(fp, str);
			}
			break;
		}
		cuDoubleComplex bcg = cuCdiv(cuCmul(acg, dot5), cuCmul(zcg, dot1));
		dot1 = dot5;

		// p = r + b(p - zAp) = r + b(p - zq)
/*
		Zaxpy(n, make_cuDoubleComplex(-zcg.x, -zcg.y), q, p);
		Zcopy(n, r, q);
		Zaxpy(n, bcg, p, q);
		Zcopy(n, q, p);
*/
		Zaxbycz(n, bcg, make_cuDoubleComplex(1, 0), cuCneg(cuCmul(bcg, zcg)), p, r, q);
	}

	// iterations
	niter[ifreq] = iter + 1;
}
