/*  pdnmesh - a 2D finite element solver
    Copyright (C) 2001-2005 Sarod Yatawatta <sarod@users.sf.net>  
  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.
  
  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
  
  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  $Id: spmat.c,v 1.10 2005/02/11 18:11:46 sarod Exp $
*/

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif


#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "types.h"

extern RBT_node_type sentinel;
#define NIL &sentinel


/* data record type */
typedef struct record_{
    int key; /* column */
    MY_DOUBLE z; /* value */
} record_type;

static int
equ(const void*a,const void *b)
{ 
  record_type *p,*q;
  p=(record_type *)a;
  q=(record_type *)b;
  return(p->key==q->key);
}
static int
lt(const void*a,const void *b)
{ 
  record_type *p,*q;
  p=(record_type *)a;
  q=(record_type *)b;
  return(p->key<q->key);

}

static void
print_details(const void *a)
{
  record_type *p;
  p=(record_type *)a;
  printf("(%d)%lf\t",p->key,p->z);
}

static void
delete_record(void *a)
{
 free((record_type*)a);
}


/* intialize a matrix */
int
SPM_init( SPMat *sM, int N ) {
  int i;
  sM->N=N;
  /* allocate N trees */
	if ((sM->rt=(RBT_tree**)malloc(sizeof(RBT_tree*)*(size_t)N))==0) {
       fprintf(stderr,"%s: %d: no free memory",__FILE__,__LINE__);
       exit(1);
 }

 for (i=0; i<N; i++) {
	if ((sM->rt[i]=(RBT_tree*)malloc(sizeof(RBT_tree)))==0) {
       fprintf(stderr,"%s: %d: no free memory",__FILE__,__LINE__);
       exit(1);
 }
 }
/* initialize trees */
 for (i=0; i<N; i++) {
  RBT_init(sM->rt[i],&equ,&lt,&print_details,&delete_record);
 }
 return(0);
}

/* add z to i,j location of M */
/* initially assumed zero */
/* i,j has to be within [0,N-1] */
int
SPM_add( SPMat *sM, int i, int j,  MY_DOUBLE z) {
  record_type *rec,brec; 
  int status;
if (i<0 || i>= sM->N|| j<0 || j>= sM->N) return(-1); /* error in range */
brec.key = j;
brec.z=z;
rec=RBT_find(&brec, sM->rt[i]);
if (rec) {
      rec->z+=z;
      /* if record is zero do deletion */
     if (IS_ZERO(rec->z)) {
      RBT_delete((void*)rec,sM->rt[i]);
	    free(rec);
     }
} else {
    	if ((rec=(record_type *)malloc(sizeof(*rec)))==0) {
       fprintf(stderr,"%s: %d: no free memory",__FILE__,__LINE__);
       exit(1);
     }
     rec->key = brec.key;
     rec->z= z;
     status=RBT_insert(rec, sM->rt[i]);
     if (status==RBT_STATUS_OK) {
				 } else if (status==RBT_STATUS_KEY_FOUND) {
        free(rec);
				 } else {
      free(rec);
		   }
}

return(0);
}



/* myltiply i,j location of M  by z */
/* initially assumed zero */
/* i,j has to be within [0,N-1] */
int
SPM_mult( SPMat *sM, int i, int j,  MY_DOUBLE z) {
  record_type *rec,brec; 
if (i<0 || i>= sM->N|| j<0 || j>= sM->N) return(-1); /* error in range */
brec.key = j;
brec.z=z;
rec=RBT_find(&brec, sM->rt[i]);
if (rec) {
      rec->z*=z;
      /* if record is zero do deletion */
     if (IS_ZERO(rec->z)) {
      RBT_delete((void*)rec,sM->rt[i]);
	    free(rec);
     }
} /* if rec not found, no need to multiply because
  result will be zero anyway */ 
return(0);
}


/* returns size of M, or value of N */
int
SPM_dim(SPMat *sM) {
 return(sM->N);
}

/* delete matrix M */
int
SPM_destroy(SPMat *sM) {
 int i;
/* free all trees */
 for (i=0; i<sM->N; i++) {
  RBT_free(sM->rt[i]);
  free(sM->rt[i]);
 }
 free(sM->rt);
 
return(0);
}




/* local visit */
static void 
visit_node_copy(RBT_node_type *n,RBT_tree *t, MY_DOUBLE *v)
{
   record_type *p;
  if ( n != NIL ){
    visit_node_copy(n->left,t,v);
    p=(record_type *)n->rec;
    v[p->key]=p->z;
    visit_node_copy(n->right,t,v);
  }
}

/* local visit-product */
static MY_DOUBLE 
visit_node_product(RBT_node_type *n,RBT_tree *t, MY_DOUBLE *v)
{
   record_type *p;
  if ( n != NIL ){
    p=(record_type *)n->rec;
    return(v[p->key]*(p->z)+visit_node_product(n->left,t,v)+visit_node_product(n->right,t,v));
  } else {
   return(0);
  }
}

/* traverse tree i and */
/* copy row i of M to vector v */
/* v size N by 1 */
int 
SPM_row(SPMat *sM, MY_INT i, MY_DOUBLE *v) {
 if ( i<0 || i>=sM->N) return(-1); /* error range */
 visit_node_copy(sM->rt[i]->root,sM->rt[i],v);
 return(0);
}


/* Matrix vector product A.x*/
/* A-sparse matrix*/
/* x- vector*/
/* v- result = Ax */
int 
SPM_vec(SPMat *sM, MY_DOUBLE *x, MY_DOUBLE *v) {
 MY_INT i;
 for (i=0; i<sM->N; i++) {
 v[i]=visit_node_product(sM->rt[i]->root,sM->rt[i],x);
 }
 return(0);
}

/* returns the value at location (i,j)
 of the matrix or zero if none found */
MY_DOUBLE 
SPM_e(SPMat *sM, MY_INT i, MY_INT j) {
  record_type *rec,brec; 
if (i<0 || i>= sM->N|| j<0 || j>= sM->N) return(0); /* error in range */
brec.key = j;
rec=RBT_find(&brec, sM->rt[i]);
if (rec)
 return(rec->z);
/* else */
return(0);
}

/* copy column i of M to vector v */
/* v size N by 1 */
/* note column access is more expensive than row access */
int 
SPM_col(SPMat *sM, MY_INT i, MY_DOUBLE *v) {
 int k;
 if ( i<0 || i>=sM->N) return(-1); /* error range */
 for (k=0; k<sM->N; k++) {
  v[k]=SPM_e(sM,k,i);
 }
 return(0);
}


/* replaces the value at location (i,j)
 of the matrix with z */
int
SPM_eq(SPMat *sM, MY_INT i, MY_INT j,MY_DOUBLE z) {
  record_type *rec,brec; 
if (i<0 || i>= sM->N|| j<0 || j>= sM->N) return(0); /* error in range */
brec.key = j;
rec=RBT_find(&brec, sM->rt[i]);
if (rec) {
 rec->z=z;
    /* if record is zero do deletion */
    if (IS_ZERO(rec->z)) {
     RBT_delete((void*)rec,sM->rt[i]);
     free(rec);
    }
 return(0);
}
/* else */
 if (!IS_ZERO(z))
  SPM_add( sM, i, j,  z);
return(-1);
}


/* a function to solve a matrix equation by conjugate gradient */
/* method Ax=y , A - symmetric positive definite */
static void 
solve_conjugate_sparse( SPMat *A, MY_DOUBLE *x, MY_DOUBLE *y, int N) {
 MY_DOUBLE  *r_old, *p_old, *v;
 MY_DOUBLE a,b,alph,beta;
 MY_INT i,k;

#ifdef DEBUG
  int j;
  for ( i =0; i < N ; i++ ) {
    for ( j = 0 ; j <N; j++ )
      printf(MDF" ",SPM_e(A,i,j));
    printf("\n");
  }
#endif
  /* allocate memory for arrays */
 if ((r_old= (MY_DOUBLE*)calloc(N,sizeof(MY_DOUBLE)))==0){
  fprintf(stderr, "%s: %d: no free memory\n", __FILE__,__LINE__);
  exit(1);
 }
 if ((p_old= (MY_DOUBLE*)calloc(N,sizeof(MY_DOUBLE)))==0){
  fprintf(stderr, "%s: %d: no free memory\n", __FILE__,__LINE__);
  exit(1);
 }
 if ((v= (MY_DOUBLE*)calloc(N,sizeof(MY_DOUBLE)))==0){
  fprintf(stderr, "%s: %d: no free memory\n", __FILE__,__LINE__);
  exit(1);
 }

/* intialization */
 /* x_old = 0 */
 for (i=0; i<N; i++) {
  p_old[i]=r_old[i]=y[i];
  x[i]=0;
 }
 /* iteration  make iteration approx size of matrix */
 for (k=0; k < N+10; k++) {
  /* calculate Ap_old =>v*/
  SPM_vec(A, p_old, v);
  a=b=0;
  for (i=0; i<N; i++) {
   a+=r_old[i]*r_old[i]; 
   b+=p_old[i]*v[i];
  }
  alph=a/b;
  b=0;
  for (i=0; i<N; i++) {
   x[i]=x[i]+alph*p_old[i];
   r_old[i]=r_old[i]-alph*v[i];
   b+=r_old[i]*r_old[i];
  }
  beta=b/a;
  for (i=0; i<N; i++) {
  p_old[i]=r_old[i]+beta*p_old[i];
  }

#ifdef DEBUG
  printf("conj grad iteration %d\n",k);
  for (i=0; i<N; i++) {
   printf(MDF" ",x[i]);
  }
  printf("\n");
#endif
 }

  free(r_old);
  free(p_old);
  free(v);
}

/* a function to build the global matrix or the matrix */
/* equation for the system - Poisson Equation */
static void 
build_sparse( SPMat *P, MY_DOUBLE *q, MY_DOUBLE *x, MY_DOUBLE *y, MY_DOUBLE *phi, int NUN, int NE, int NN, MY_INT *renumber, MY_INT *re_renumber) 
{
  /* infd - input file descriptor                  */
  /* P[][] - global matrix to build   NUNxNUN        */
  /* q[] - global array to build     NUNx1          */
  /* x[] - x coordinates             NNx1          */
  /* y[] - y coordinates             NNx1          */
  /* phi[] - potentials              NNx1          */
  /* NUN - no of unknown potential nodes             */
  /* NE - no of elements                            */
  /* NN - no of nodes                               */

  /* now the running variables for each triangular element */
  int v[3]; /* to keep vertex number of each node */
  MY_DOUBLE pl[3][3];  /* local version of P[][] above */
  MY_DOUBLE ql[3];	/* local version of q[] above */
  MY_DOUBLE rho[3]; /* rho values */ /* conductivity */
  triangle *t;
 
  int    i,j;
  
  MY_DOUBLE 
    epsilon, /* permittivity */
    A, /* area of a triangulr element */
    k, /* value of the determinanat */
    yy0, yy1, yy2;
  
  MY_DOUBLE b[3]; /*local arrays */
  MY_DOUBLE c[3];

  

  update_status_bar("Begin building...");
	/* now read node data with unknown potentials */
  for ( i = 0; i < NUN; i++) {
    x[i]= (Mx(re_renumber[i],M))*g_xscale-g_xoff;
    y[i]= (My(re_renumber[i],M))*g_yscale-g_yoff;
  }
	
  /* now read node data with known potentials , which comes last */
  for ( i = NUN; i < NN; i++ ) { 
    x[i] = (Mx(re_renumber[i],M))*g_xscale-g_xoff;
    y[i] = (My(re_renumber[i],M))*g_yscale-g_yoff;
    phi[i] = (Mz(re_renumber[i],M));
  }
	
  /* now initialize P[][] and q[] */
 /* for ( i = 0; i < NUN; i++) {
    q[i] = 0;
  }*/
  
#ifdef DEBUG
	for  ( i = 0; i< NN; i++)
	    printf("build: %d "MDF", "MDF": "MDF"\n", i, x[i], y[i], phi[i]); 
#endif
	
	/* now read element data, whilst building the global matrices */
	DAG_traverse_list_reset(&dt);
	t=DAG_traverse_prune_list(&dt);
	while(t) {
	  if ( t &&t->status==LIVE ) {
	    
	    v[0] = renumber[t->p0];
	    v[1] = renumber[t->p1];
	    v[2] = renumber[t->p2]; /* ccw direction */
			if (bound.poly_array[t->boundary].rhoexp) {
	      rho[0] = ex(bound.poly_array[t->boundary].rhoexp,t->p0);
	      rho[1] = ex(bound.poly_array[t->boundary].rhoexp,t->p1);
	      rho[2] = ex(bound.poly_array[t->boundary].rhoexp,t->p2);
			} else {
	      rho[0] = bound.poly_array[t->boundary].rho;
	      rho[1] = bound.poly_array[t->boundary].rho;
	      rho[2] = bound.poly_array[t->boundary].rho;
			}

	    epsilon = 1/bound.poly_array[t->boundary].mu;
	    

	    /* now have read one triangle , add it to global */
	    /* first the coefficients */
	    /* some temporary values, neede several times */
	    yy0 = y[v[1]]-y[v[2]];
	    yy1 = y[v[2]]-y[v[0]];
	    yy2 = y[v[0]]-y[v[1]];
	    
	    /* area of the triangle */
	    A = 0.5*( x[v[0]]*yy0 + x[v[1]]*yy1 + x[v[2]]*yy2 ); 
	 /*   A = ABS(A); */ /* points are in ccw order to ensure this */
	    /* another constant */
	    k = 0.25/A;
	    /* now calculate b's and c's */
#ifdef DEBUG
	  printf("adding element with nodes %d %d and %d...\n\n", v[0], v[1], v[2]);
#endif
	  b[0]= yy0;
#ifdef DEBUG
	  printf("b0="MDF"  ", b[0]);
#endif
	  b[1]= yy1;
#ifdef DEBUG
	  printf("b1="MDF"  ", b[1]);
#endif
	  b[2]= yy2;
#ifdef DEBUG
	  printf("b2="MDF" \n", b[2]);
#endif	  
	  c[0] = ( x[v[2]]-x[v[1]] ); 
#ifdef DEBUG
	  printf("c0="MDF"  ", c[0]); 
#endif
	  c[1] = ( x[v[0]]-x[v[2]] );
#ifdef DEBUG
	  printf("c1="MDF"  ", c[1]);
#endif
	  c[2] = ( x[v[1]]-x[v[0]] );
#ifdef DEBUG
	  printf("c2="MDF" \n", c[2]);
#endif
	  
	  /* now build pl[] and ql */
	  for ( i = 0; i < 3; i++ ) 	
	    for ( j = 0; j < 3; j++ )
	      pl[i][j] = k*epsilon*(b[i]*b[j]+c[i]*c[j]); 
	  
	  ql[0] = A*(2*rho[0]+rho[1]+rho[2])*ONE_OVER_TWELVE;
	  ql[1] = A*(rho[0]+2*rho[1]+rho[2])*ONE_OVER_TWELVE;
	  ql[2] = A*(rho[0]+rho[1]+2*rho[2])*ONE_OVER_TWELVE;
#ifdef DEBUG 
	  printf("local p and q ..\n");
	  for ( i = 0; i <3; i++) {
	    printf("| ");
	    for ( j = 0; j < 3; j++)
	      printf(" "MDF" ", pl[i][j]);
	    printf("|| "MDF"  | \n", ql[i]); 
	  }	
#endif 
	  
	  
	  /* better to use multiplication than division */
	  /* now augment the global data set */
	  for ( i = 0; i < 3; i++)
	    if ( v[i] < NUN ) { /* we index from 0, not from 1 */
	      q[v[i]] = q[v[i]] + ql[i];
	      for ( j = 0; j < 3; j++ ) {
		if ( v[j] < NUN )
		  { 
					 if ( v[j] <= v[i] )
           SPM_add(P,v[i],v[j],pl[i][j]);
           SPM_eq(P,v[j],v[i],SPM_e(P,v[i],v[j]));
					 /* we build both halfs of the matrix:FIXME: this can be done later */
		  }
		else 
		  { 
					 q[v[i]] = q[v[i]] - (pl[i][j])*phi[v[j]];
		  } 
	      }
	    } /*else printf("rejecting %d\n", v[i]);*/
#ifdef DEBUG
	  printf("building global matrix ..\n");
#endif
	  }
			t=DAG_traverse_prune_list(&dt);
	}
	
#ifdef DEBUG	  
	  for ( i = 0; i < NUN; i++){
	    printf("|  ");
	    for( j = 0; j <= i; j++)
	      printf(" "MDF"  ", SPM_e(P,i,j));
	    printf("=  "MDF"  |\n", q[i]);
	    } 
	   for ( i = 0; i < NN; i++) 
		printf("node %d potential "MDF"\n", i, phi[i]); 
#endif
  update_status_bar("Done building.");
}

/* after genarating the mesh, need to re-number points */
/* so that unknown points are first */
/* then we solve the problem -poisson equation*/
int 
re_number_and_solve_poisson_sparse(void)
{
  unsigned MY_INT i,j;
 MY_INT * renumber; /* re-numbering array */
 MY_INT * re_renumber; /* re-renumbering array */
 /* used to map potentials back to the original points */


  unsigned MY_INT unknowns=0; /* unknown points */
  unsigned MY_INT total = 0; /* total points */
  MY_INT triangles = 0; /* known points */


  triangle *t;
  
  /* first the global node data, all total size */
  MY_DOUBLE *x ; /* x coordinates */
  MY_DOUBLE *y ; /* y coordinates */
  MY_DOUBLE *phi ; /* potentials at each node */
  
  /* the global variables */
  SPMat P ; /* sparse NN x NN matrix, only the unknowns */
  MY_DOUBLE *q ; /* unknowns sizesame here */

#ifdef DEBUG
 printf("starting sparse solver\n");
#endif

  /* allocate memory for arrays */
  if ((renumber = (MY_INT *)calloc(M.count,sizeof(MY_INT )))==0){
   fprintf(stderr, "%s: %d: no free memory\n", __FILE__,__LINE__);
   exit(1);
 }

  if((re_renumber = (MY_INT *)calloc(M.count,sizeof(MY_INT)))==0){
   fprintf(stderr, "%s: %d: no free memory\n", __FILE__,__LINE__);
   exit(1);
 }

  /* need to solve for only points referred by triangles */
  /* SO -- */
		DAG_traverse_list_reset(&dt);
		t=DAG_traverse_prune_list(&dt);
  while ( t ) {
    if ( t && t->status==LIVE ) {
      /* mark the points in this triangle */
      /* use re_renumber[] array for this */
      re_renumber[t->p0] = 1;
      re_renumber[t->p1] = 1;
      re_renumber[t->p2] = 1;

      triangles++;
    }
		t=DAG_traverse_prune_list(&dt);
  }

  /* now find unknowns, or INSIDE points */
  j = 0;
  for ( i = 0; i < M.count; i++ )
    if (  Mval(i,M) == VAR  && re_renumber[i] == 1) {
      renumber[i] = j++; 
#ifdef DEBUG
						printf("renumber_and_solve: renumbering variable point %d as %d\n",i,renumber[i]);
#endif
      unknowns++;
    }
  /* now renumber the known points */
  total = unknowns;
  for ( i = 0; i < M.count; i++ )
    if ( Mval(i,M) == FX  && re_renumber[i]==1 ) {
      renumber[i] = j++; 
#ifdef DEBUG
						printf("renumber_and_solve: renumbering fixed point %d as %d\n",i,renumber[i]);
#endif
      total++;
    }
 /* finally give numbers to points 0,1,2 */
	/* we will not use these points */
		renumber[0]=j++;
		renumber[1]=j++;
		renumber[2]=j++;

/* now construct re-renmbering array */
		  i = 0;
				/* do an insertion sort */
				while ( i < M.count ) {
						for ( j = 0; j < M.count; j++ ) {
						   if ( renumber[j] ==(int) i )
									  break;
						}
										re_renumber[i] = j;
										i++;
				}

#ifdef DEBUG
printf("build: renumbering %d points with %d unknowns out of %d total\n",total,unknowns,M.count);
#endif
  /*renumber_points(renumber, re_renumber, total, unknowns); */
  
#ifdef DEBUG
  printf("i,  renum, rerenum\n");
  for ( i = 0; i < total; i++ )
    printf("%d %d %d\n",i,renumber[i],re_renumber[i]);
  for ( i = 0; i < total; i++ )
    printf("%d  "MDF"  "MDF" "MDF"\n",i,Mx(i,M),My(i,M),Mz(i,M));
#endif

  
  /* we free these two arrays only after solving the problem */
  /* now solve the problem */
  /* first allocate memory for the arrays */
  if((x=(MY_DOUBLE *)calloc(total, sizeof(MY_DOUBLE)))==0){
   fprintf(stderr, "%s: %d: no free memory\n", __FILE__,__LINE__);
   exit(1);
 }

  if((y=(MY_DOUBLE *)calloc(total, sizeof(MY_DOUBLE)))==0){
   fprintf(stderr, "%s: %d: no free memory\n", __FILE__,__LINE__);
   exit(1);
 }

  if((phi=(MY_DOUBLE *)calloc(total, sizeof(MY_DOUBLE)))==0){
   fprintf(stderr, "%s: %d: no free memory\n", __FILE__,__LINE__);
   exit(1);
 }

  if((q=(MY_DOUBLE*)calloc(unknowns, sizeof(MY_DOUBLE)))==0){
   fprintf(stderr, "%s: %d: no free memory\n", __FILE__,__LINE__);
   exit(1);
 }

	
	/* now the matrix */
  SPM_init(&P,unknowns);

	/* Poisson equation */
	/* now the memory allocation is over */
  build_sparse( &P, q, x, y, phi, unknowns, triangles, total, renumber, re_renumber);
  /* if user has chosen so, solve by conj. grad method*/
  /* by default solve by cholevsky decomposition */
  solve_conjugate_sparse(&P,phi,q,unknowns);
  
	/* after solution, array phi[] will contain the new potentials */
  /* now put them back to the point data */
  /* we also note the maximum and minimum values for contour plotting */
  g_maxpot = -1000;
  g_minpot = 1000;
  for ( i = 0; i < total; i++ ) {
    Mz(re_renumber[i],M) = phi[i];
    if ( phi[i] > g_maxpot )
      g_maxpot = phi[i];
    if ( phi[i] < g_minpot )
      g_minpot = phi[i];
  }



  /* now we can free the renumber[] array */
  free(renumber);
  /* now free this array too */
  free(re_renumber);
  
  /* free everything else */
  
  /* first the arrays */
  free(x);
  free(y);
  free(phi);
  free(q);
  
  /* next the matrix */
  SPM_destroy(&P); 
  /* now printout a summery */
#ifdef DEBUG
  for ( i = 0; i < total; i++ )
    printf("renumber_and_solve: %d  "MDF"  "MDF" "MDF"\n",i,Mx(i,M),My(i,M),Mz(i,M));
#endif

 return(0);
}