#include <alberta.h>
#include <alberta_util.h>
#if (defined(ALBERTA_VER) && (ALBERTA_VER==2))
//# include <alberta_inlines.h>
//#include <alberta_util_inlines.h>
# include "alberta1_wrapper.h"
#endif
#include "newalberta.hh"
#include "solve_multiphase.h"

#define USE_THREADS (1)
#ifdef USE_THREADS
#include <pthread.h>
#endif
/****************************************************************************/
/*   y <- Ax   coupled vector valued problems with matrix                   */
/*                                                                          */
/*             |A_00   A_01 ... A_0n|                                       */
/*             |A_01^t A_11 ....A_1n|                                       */
/*             | ...            ... |                                       */
/*             |A_0n^t ...      A_nn|                                       */
/*                                                                          */
/****************************************************************************/
#define GETBLOCK_MULTI(n,m) ((n)*nn+(m))

#ifdef USE_THREADS
/* every phase is handled in its own thread*/

// wie in phase.hh 
#ifndef NROFPHASES
# define NROFPHASES (8)
#endif

pthread_mutex_t mutex[NROFPHASES];
pthread_t thread[NROFPHASES];

struct mv_arg 
{ 
  int phase;
  int nrofphases;
  const DOF_MATRIX * const * const a; 
  const DOF_REAL_VEC * const * const x;
  DOF_REAL_VEC * const * const y;
  mv_arg(int phase, int nrofphases, const DOF_MATRIX * const * const a,
	 const DOF_REAL_VEC * const * const x,
	 DOF_REAL_VEC * const * const y)
    : phase(phase), nrofphases(nrofphases), a(a), x(x), y(y)
  {
    pthread_mutex_init(&mutex[phase],NULL);
    pthread_mutex_lock(&mutex[phase]);
  }; 
  ~mv_arg() 
  {
    // Deadlock droht
    /*    pthread_mutex_destroy(&mutex[phase]); */
  };
};
// pthreads brauchen parameter-casts
void *mv_phase(void *ptr)
{
  mv_arg *arg=(mv_arg*)ptr;
  int n=arg->phase; 
  const int nn=arg->nrofphases; // fuers macro
 
  const dof_matrix *suba=arg->a[GETBLOCK_MULTI(n,n)];

  /*  if((suba)) */
  dof_mv(NoTranspose, suba, arg->x[n], arg->y[n]);
  pthread_mutex_unlock(&mutex[n]);

  /* off-diagonal terms -- counting down */
  for(int m=arg->nrofphases-1; m>n ; m--)
    {
      if((suba=arg->a[GETBLOCK_MULTI(n,m)]))
       {
	 pthread_mutex_lock(&mutex[n]);
	 dof_gemv(NoTranspose, 1.0, suba, arg->x[m], 1.0, arg->y[n]);
	 pthread_mutex_unlock(&mutex[n]);

/* symmetry is lost for Dirichlet problems */
         suba=arg->a[GETBLOCK_MULTI(m,n)];
	 pthread_mutex_lock(&mutex[m]);
	 dof_gemv(NoTranspose, 1.0, suba, arg->x[n], 1.0, arg->y[m]);
	 pthread_mutex_unlock(&mutex[m]);
       }
    }
    pthread_exit(ptr);
}
void dof_mv_multi (int nn, const DOF_MATRIX * const * const a, 
		   const DOF_REAL_VEC * const * const x, 
		   DOF_REAL_VEC * const * const y)
{
  int    n;
  const  dof_matrix *suba;
  mv_arg *argp;
 
  pthread_mutex_init(&mutex[nn-1],NULL);
  pthread_mutex_lock(&mutex[nn-1]);

  for (n=nn-2; n>=0; n--)
    {
      argp = new mv_arg(n,nn,a,x,y);
      
      pthread_create(&thread[n],NULL,mv_phase,(void *)argp);
    }
  /* handle last phase here*/
  n=nn-1;
  if((suba=a[GETBLOCK_MULTI(n,n)]))
     dof_mv(NoTranspose, suba, x[n], y[n]);
  pthread_mutex_unlock(&mutex[n]);
  /* collect threads */
  for (n=nn-2; n>=0; n--)
    {
      pthread_join(thread[n],(void **)&argp);
      delete argp;
    }
  /* free mutexes */
  for(n=0; n<nn; n++) pthread_mutex_destroy(&mutex[n]);
  return;
}
#else
void dof_mv_multi (int nn, const DOF_MATRIX * const * const a, 
		      const DOF_REAL_VEC * const * const x, DOF_REAL_VEC * const * const y)
{
  int             n, m;
  const dof_matrix *suba;

/* diagonal terms */
  for (n = 0; n < nn; n++)
  {
    if((suba=a[GETBLOCK_MULTI(n,n)]))
     dof_mv(NoTranspose, suba, x[n], y[n]);
     
  }
/* off-diagonal terms */
  for (n = 0; n < nn; n++)
  {
    for (m = n+1; m < nn; m++)
    {
     if((suba=a[GETBLOCK_MULTI(n,m)]))
       {
	 dof_gemv(NoTranspose, 1.0, suba, x[m], 1.0, y[n]);
	 /* use transpose for symmetric data  
	 dof_gemv(Transpose, 1.0, suba, x[n], 1.0, y[m]);
	 */
/* symmetry is lost for Dirichlet problems */
         suba=a[GETBLOCK_MULTI(m,n)];
	 dof_gemv(NoTranspose, 1.0, suba, x[n], 1.0, y[m]);
	 
       }
    }
  }
}
#endif
/* ohne wurzel */
REAL dof_nrm_multi(int n, const DOF_REAL_VEC * const * const x)
{
  int i;
  REAL nrm = 0.0;

  for (i=0;i<n;i++)
    {
      REAL h=dof_nrm2(x[i]); 
      nrm += h*h; 
    }
  return nrm;
}

REAL dof_dot_multi(int n, const DOF_REAL_VEC * const * const x, const DOF_REAL_VEC * const * const y)
{
  int i;
  REAL sum = 0.0;

  for (i=0;i<n;i++)
    {
      sum += dof_dot(x[i],y[i]);  
    }
  return sum;
}

void dof_axpy_multi(int n, REAL alfa, const DOF_REAL_VEC * const * const x, DOF_REAL_VEC * const * const y)
{
  int i;

  for (i=0;i<n;i++)
    {
      dof_axpy(alfa, x[i], y[i]);  
    }
}

void dof_xpay_multi(int n, REAL alfa, const DOF_REAL_VEC * const * const x, DOF_REAL_VEC * const * const y)
{
  int i;

  for (i=0;i<n;i++)
    {
      dof_xpay(alfa, x[i], y[i]);  
    }
}

void dof_set_multi(int n, REAL alfa, DOF_REAL_VEC * const * const x)
{
  int i;
  for (i=0;i<n;i++)
    {
      dof_set(alfa, x[i]);  
    }
}

void dof_copy_multi(int n, const DOF_REAL_VEC * const * const x, DOF_REAL_VEC * const * const y)
{
  int i;

  for (i=0;i<n;i++)
    {
      dof_copy(x[i], y[i]);  
    }
}

/* muessen wir oem_cg nachbauen? Vielleicht abgespeckte Version
   ohne precon, was wohl etwa wie folgt aussehen wuerde */

/* aus : alberta-1.2/SOLVER/src/cg.c */ 
/* arbeiten mit dof_* typen: das mag langsam aber dafuer durchschaubar sein */
/* beachte reihenfolge: A, b, x !*/
int cg_multi(const SOLVE_DATA *sd, int n, const DOF_MATRIX *const*const a, 
		      const DOF_REAL_VEC *const*const b, DOF_REAL_VEC *const*const x)
{
#ifndef USE_THREADS
  FUNCNAME("cg_multi");
#else
  FUNCNAME("cg_multi(threaded)");
#endif
  REAL  rho, hd_2, delta, gamma;
  DOF_REAL_VEC  *d[n], *h[n], *r[n];
  int   i,j,iter=0;
  const int nn=n;
  const REAL TOL = 1.e-30;
  char key[128];
  
  INFO(sd->info,2,"with tolerance %12.5le\n",sd->tolerance);
  if(sd->info>=10)
  {
   MSG("nr of components %d\n",n);
   for(i=0;i<n;i++)
   {
    print_dof_real_vec(b[i]);
    print_dof_real_vec(x[i]);
    
   }
   for(i=0;i<n;i++) for(j=0;j<n;j++) print_dof_matrix(a[GETBLOCK_MULTI(i,j)]);
  }
  /* test for b==0 */
  if (sqrt(dof_nrm_multi(n, b)) < TOL)
  {
    INFO(sd->info,2,"b == 0, x = 0 is the solution of the linear system\n");
    dof_set_multi(n, 0.0, x);
    return(0);
  }
 
  /* allocate WS for d,h,r */
  for(i=0; i<n; i++) 
    {
      const FE_SPACE *fe_space = x[i]->fe_space;
      sprintf (key,"d%d",i);
      d[i]=get_dof_real_vec( key, fe_space);
      sprintf (key,"r%d",i);
      r[i]=get_dof_real_vec( key, fe_space);
      sprintf (key,"h%d",i);
      h[i]=get_dof_real_vec( key, fe_space);
    }
 
  dof_mv_multi(n, a, (const DOF_REAL_VEC **)x, r); // r=Ax
  dof_axpy_multi(n, -1.0, b, r);      // r=r-b
  // skipping precon
/*--------------------------------------------------------------------------*/
/*---  check initial residual  ---------------------------------------------*/
/*--------------------------------------------------------------------------*/

  delta=dof_nrm_multi(n, (const DOF_REAL_VEC **)r);
  INFO(sd->info,2,"iter. |     residual\n");
  INFO(sd->info,2,"%5d | %12.5le\n",0,delta);
  if(sqrt(delta)>sd->tolerance) {

  dof_copy_multi(n, (const DOF_REAL_VEC **)r, d);  // d = r
/*--------------------------------------------------------------------------*/
/*---  Iteration  ----------------------------------------------------------*/
/*--------------------------------------------------------------------------*/

  for (iter = 1; iter <= sd->max_iter; iter++)
    {
      dof_mv_multi(n, a, (const DOF_REAL_VEC **)d, h); // h=Ad
/*--------------------------------------------------------------------------*/
/*---  decent parameter rho = delta/(h,d)_2  -------------------------------*/
/*--------------------------------------------------------------------------*/
      hd_2 = dof_dot_multi(n, (const DOF_REAL_VEC **)h, (const DOF_REAL_VEC **)d); // <h,d>
      if (ABS(hd_2) < TOL)
	{
	  break;
	}
/*--------------------------------------------------------------------------*/
/*---  update x and r  -----------------------------------------------------*/
/*--------------------------------------------------------------------------*/
    rho = delta/hd_2;
    dof_axpy_multi(n, -rho, (const DOF_REAL_VEC **)d, x); // x = x- rho*d
    dof_axpy_multi(n, -rho, (const DOF_REAL_VEC **)h, r); // r = r- rho*h
    // skipping precon
    gamma = 1.0/delta;
    
    delta=dof_nrm_multi(n, (const DOF_REAL_VEC **)r);
    if(!(iter%20)) { INFO(sd->info,2,"%5d | %12.5le\n",iter,delta); }
    if(sqrt(delta)<sd->tolerance) break;
/*--------------------------------------------------------------------------*/
/*---  update of the decent direction  -------------------------------------*/
/*--------------------------------------------------------------------------*/
    gamma *= delta;
    dof_xpay_multi(n, gamma, (const DOF_REAL_VEC **)r, d); // d= r+gamma*d
  }
  } //ENDIF
  /* initial_res=0 or delta<tol or max_iter reached */ 
  INFO(sd->info,2,"%5d | %12.5le\n",iter,delta); 
  
  if(sd->info>=10)
  {
   MSG("nr of components %d\n",n);
   for(i=0;i<n;i++)
   {
    print_dof_real_vec(x[i]);
   }
  }
   
  /* free ws */
  for (i=0; i<n; i++) 
    { 
      free_dof_real_vec(d[i]);
      free_dof_real_vec(h[i]);
      free_dof_real_vec(r[i]); 
    } 
  return iter;
}
