/*
bicgstab.c (MPI)
*/

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

void bicgstab(int ifreq, int64_t n, d_complex_t *a, d_complex_t *x, d_complex_t *b, d_complex_t *r, d_complex_t *p, d_complex_t *q, d_complex_t *t, d_complex_t *u, d_complex_t *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\n", commSize, commRank); fflush(stdout);
	//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));

	// == initial value ==

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

	// r = b - Ax = b - q
	Zcopy(n, b, r);
	Zaxpy(n, d_complex(-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)
	d_complex_t l_dot1 = Zdotc(n1 - n0, &r0[n0], &r[n0]);
	d_complex_t dot1 = comm_zsum(l_dot1);

	// == 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)
		d_complex_t l_dot2 = Zdotc(n1 - n0, &r0[n0], &q[n0]);
		d_complex_t dot2 = comm_zsum(l_dot2);
		if (d_abs(dot2) == 0) {
			if (iout) {
				sprintf(str, "%s (1) (%d %e %e)", msg4, iter, d_abs(dot1), d_abs(dot2));
				monitor1(fp, str);
			}
			break;
		}
		d_complex_t acg = d_div(dot1, dot2);
		//printf("a : %d %d %d %e %e %e\n", commSize, commRank, i, acg.r, acg.i, d_abs(acg));  fflush(stdout);

		// t = r - aAp = r - aq
		Zaxpy(n, d_rmul(-1, acg), 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)
		d_complex_t l_dot3 = Zdotc(n1 - n0, &u[n0], &t[n0]);
		d_complex_t l_dot4 = Zdotc(n1 - n0, &u[n0], &u[n0]);
		d_complex_t dot3 = comm_zsum(l_dot3);
		d_complex_t dot4 = comm_zsum(l_dot4);
		if (d_abs(dot4) == 0) {
			if (iout) {
				sprintf(str, "%s (2) (%d %e %e)", msg4, iter, d_abs(dot3), d_abs(dot4));
				monitor1(fp, str);
			}
			break;
		}
		d_complex_t zcg = d_div(dot3, dot4);
		//printf("z : %d %d %d %e %e %e\n", commSize, commRank, i, zcg.r, zcg.i, d_abs(zcg)); fflush(stdout);

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

		// r = t - zAt = t - zu
		Zaxpy(n, d_rmul(-1, zcg), 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

		//printf("%d %d %d %f\n", commSize, commRank, i, iter[i]); fflush(stdout);
		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
		d_complex_t l_dot5 = Zdotc(n1 - n0, &r0[n0], &r[n0]);
		d_complex_t dot5 = comm_zsum(l_dot5);
		if (d_abs(d_mul(zcg, dot1)) == 0) {
			if (iout) {
				sprintf(str, "%s (3) (%d %.3e %.3e %.3e %.3e)", msg4, iter, d_abs(acg), d_abs(zcg), d_abs(dot1), d_abs(dot5));
				monitor1(fp, str);
			}
			break;
		}
		d_complex_t bcg = d_div(d_mul(acg, dot5), d_mul(zcg, dot1));
		//printf("b : %d %d %d %e %e %e\n", commSize, commRank, i, bcg.r, bcg.i, d_abs(bcg)); fflush(stdout);
		dot1 = dot5;

		// p = r + b(p - zAp) = r + b(p - zq)
/*
		Zaxpy(n, d_rmul(-1, zcg), q, p);
		Zcopy(n, r, q);
		Zaxpy(n, bcg, p, q);
*/
		Zaxbycz(n, bcg, d_complex(1, 0), d_neg(d_mul(bcg, zcg)), p, r, q);
	}

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