/*-------------------------------------------------
  rsdpa_cg.cpp
-------------------------------------------------*/

// デバッグが必要

#include <rsdpa_cg.h>

int rCg::m      = 0;
int rCg::mb2    = 0;
int rCg::maxmp2 = 0;
int rCg::maxnp2 = 0;
int rCg::ictxt2 = 0;
int rCg::iam    = 0;
int rCg::nprocs = 0;
int rCg::myrow2 = 0;
int rCg::mycol2 = 0;
int rCg::nprow2 = 0;
int rCg::npcol2 = 0;

rCg::rCg()
{
  p = NULL;
  q = NULL;
  r = NULL;
}

rCg::~rCg()
{
  if (p!=NULL) {
    delete[] p;
  }
  p = NULL;
  if (q!=NULL) {
    delete[] q;
  }
  q = NULL;
  if (r!=NULL) {
    delete[] r;
  }
  r = NULL;
}

int rCg::initialize(int m, double* A, int* descA, double* x,
		    int* descx)
{
  // ここまでに, A,x の領域を確保しておくこと
  this->A     = A;
  this->descA = descA;
  this->x     = x;
  this->descx = descx;
  this->m     = m;

  iam    = rMpiSt::iam;
  mb2    = rMpiSt::mb2;
  ictxt2 = rMpiSt::ictxt2;
  nprocs = rMpiSt::nprocs;
  myrow2 = rMpiSt::myrow2;
  mycol2 = rMpiSt::mycol2;
  nprow2 = rMpiSt::nprow2;
  npcol2 = rMpiSt::npcol2;

  maxmp2 = max(1,numroc_(&m,&mb2,&myrow2,&IZERO,&nprow2));
  maxnp2 = max(1,numroc_(&m,&mb2,&mycol2,&IZERO,&npcol2));

  p = NULL;
  p = new double[maxmp2*1];
  if (p == NULL) {
    rError("memory exhausted");
  }
  q = new double[maxmp2*1];
  if (q == NULL) {
    rError("memory exhausted");
  }
  r = new double[maxmp2*1];
  if (r == NULL) {
    rError("memory exhausted");
  }

  int info = 0;
  descinit_(descp,&m,&IONE,&mb2,&IONE,&IZERO,&IZERO,
	    &ictxt2,&maxmp2,&info);
  if (info!=0) {
    rMessage("info << " << info);
  }
  descinit_(descq,&m,&IONE,&mb2,&IONE,&IZERO,&IZERO,
	    &ictxt2,&maxmp2,&info);
  if (info!=0) {
    rMessage("info << " << info);
  }
  descinit_(descr,&m,&IONE,&mb2,&IONE,&IZERO,&IZERO,
	    &ictxt2,&maxmp2,&info);
  if (info!=0) {
    rMessage("info << " << info);
  }
  return _SUCCESS;
}


int rCg::solve(double epsilon, int  maxiter, int* iter)
{
  int ISTDOUT = 6;
  // double tmpwork[maxmp2];
  // 初期化
  int k = 0;
  // r = x;
  pdcopy_(&m,x,&IONE,&IONE,descx,&IONE,
	  r,&IONE,&IONE,descr,&IONE);
  // pdlaprnt_(&m,&IONE,r,&IONE,&IONE,descr,&IZERO,&IZERO,
  //	    "r", &ISTDOUT,tmpwork,1);
  // x = 0;
  catlas_dset(maxmp2,0.0,x,1);
  double r_norm_1,r_norm_2,alpha,beta;
  pdnrm2_(&m,&r_norm_1,r,&IONE,&IONE,descr,&IONE);
  r_norm_1 = r_norm_1*r_norm_1;
  // pdlaprnt_(&m,&IONE,r,&IONE,&IONE,descr,&IZERO,&IZERO,
  //	    "r", &ISTDOUT,tmpwork,1);
  // ここでは r_norm_1 はすくなくとも p,q,r をもっている
  // プロセスには正確に伝わるはず
  // つまり、mycol2 == 0 に情報はきている
  if (mycol2 == 0) {
    Cdgebs2d(ictxt2,"Row"," ",1,1,(char*)&r_norm_1,1);
  }
  else {
    Cdgebr2d(ictxt2,"Row"," ",1,1,(char*)&r_norm_1,1,myrow2,0);
  }
  // pdlaprnt_(&m,&IONE,r,&IONE,&IONE,descr,&IZERO,&IZERO,
  //	    "r", &ISTDOUT,tmpwork,1);
  // rMessage("r_norm_1 = " << r_norm_1);
  Cblacs_barrier(ictxt2,"All");
  
  while (k < maxiter && r_norm_1 > epsilon*epsilon) {
    k++;
    if (k == 1) {
      // p = r;
      pdcopy_(&m,r,&IONE,&IONE,descr,&IONE,
	      p,&IONE,&IONE,descp,&IONE);
      // pdlaprnt_(&m,&IONE,p,&IONE,&IONE,descp,&IZERO,&IZERO,
      //   "p", &ISTDOUT,tmpwork,1);
    }
    else {
      beta = r_norm_1 / r_norm_2;
      // p = r + beta*p;

      // １度 p を q に退避
      pdcopy_(&m,p,&IONE,&IONE,descp,&IONE,
	      q,&IONE,&IONE,descq,&IONE);
      // p = r
      pdcopy_(&m,r,&IONE,&IONE,descr,&IONE,
	      p,&IONE,&IONE,descp,&IONE);
      // p = p + beta*q;
      pdaxpy_(&m,&beta,q,&IONE,&IONE,descq,&IONE,
	      p,&IONE,&IONE,descp,&IONE);
    }

    // pdlaprnt_(&m,&IONE,p,&IONE,&IONE,descp,&IZERO,&IZERO,
    //	"p", &ISTDOUT,tmpwork,1);
    // pdlaprnt_(&m,&m,A,&IONE,&IONE,descA,&IZERO,&IZERO,
    //	"A", &ISTDOUT,tmpwork,1);
    // q = A*p;
    pdsymv_("Lower",&m,&DONE,A,&IONE,&IONE,descA,
	    p,&IONE,&IONE,descp,&IONE,
	    &DZERO,q,&IONE,&IONE,descq,&IONE);
    // pdlaprnt_(&m,&IONE,q,&IONE,&IONE,descq,&IZERO,&IZERO,
    //	      "q", &ISTDOUT,tmpwork,1);
    // alpha = r_norm_1/(p'*q);
    pddot_(&m,&alpha,p,&IONE,&IONE,descp,&IONE,
	   q,&IONE,&IONE,descq,&IONE);
    if (mycol2 == 0) {
      Cdgebs2d(ictxt2,"Row"," ",1,1,(char*)&alpha,1);
    }
    else {
      Cdgebr2d(ictxt2,"Row"," ",1,1,(char*)&alpha,1,myrow2,0);
    }
    // rMessage("alpha = " << alpha);
    alpha = r_norm_1 / alpha;
    // rMessage("alpha = " << alpha);

    // x = x + alpha*p;
    // pdlaprnt_(&m,&IONE,x,&IONE,&IONE,descx,&IZERO,&IZERO,
    //    "x", &ISTDOUT,tmpwork,1);
    pdaxpy_(&m,&alpha,p,&IONE,&IONE,descp,&IONE,
	    x,&IONE,&IONE,descx,&IONE);
    // pdlaprnt_(&m,&IONE,x,&IONE,&IONE,descx,&IZERO,&IZERO,
    //      "x", &ISTDOUT,tmpwork,1);
    // r = r - alpha*q;
    alpha = -alpha;
    pdaxpy_(&m,&alpha,q,&IONE,&IONE,descq,&IONE,
	    r,&IONE,&IONE,descr,&IONE);
    // pdlaprnt_(&m,&IONE,r,&IONE,&IONE,descr,&IZERO,&IZERO,
    //      "r", &ISTDOUT,tmpwork,1);
    
    r_norm_2 = r_norm_1;
    
    // r_norm_1 = r'*r;
    pdnrm2_(&m,&r_norm_1,r,&IONE,&IONE,descr,&IONE);
    if (mycol2 == 0) {
      Cdgebs2d(ictxt2,"Row"," ",1,1,(char*)&r_norm_1,1);
    }
    else {
      Cdgebr2d(ictxt2,"Row"," ",1,1,(char*)&r_norm_1,1,myrow2,0);
    }
    r_norm_1 = r_norm_1*r_norm_1;
    // rMessage("r_norm_1 = " << r_norm_1);

  }

  // これでちゃんと解き終わっているはず
  // pdlaprnt_(&m,&IONE,x,&IONE,&IONE,descx,&IZERO,&IZERO,
  //	    "x", &ISTDOUT,tmpwork,1);
  
  *iter = k;
  if (rMpiSt::iam == 0) {
    rMessage(k << "iterations in linear solver");
  }
  return 0;
}
