/************************************************************************
* This program solves a tridiagonal linear set of equations M * U = R   *
* --------------------------------------------------------------------- *
* SAMPLE RUN:                                                           *
*                                                                       *
* Input data file (tridiag.dat):                                        *
*                                                                       *
*  5                                                                    *
*  0.0   1.0   -2.0                                                     *
* 10.0   5.0   -7.5                                                     *
*  2.22 -3.25   0.00456                                                 *
*  2.0   4.0   -3.3                                                     *
* -1.0   3.0    0.0                                                     *
*                                                                       *
* (The first line contains the size of system                           *
* Then the three columns give the three main diagonals of matrix M, the *
* first column begins with a zero and the third column ends with zero)  *
*                                                                       *
* Output file (tridiag.lst):                                            *
*                                                                       *
* --------------------------------------------------------------------  *
*   Tridiagonal system of equations                                     *
* --------------------------------------------------------------------  *
*  upper subdiagonal:                                                   *
* -2.000000 -7.500000  0.004560 -3.300000  0.000000                     *
*  main diagonal:                                                       *
*  1.000000  5.000000 -3.250000  4.000000  3.000000                     *
*  lower subdiagonal:                                                   *
*  0.000000  10.00000  2.220000  2.000000 -1.000000                     *
*  right side member:                                                   *
*  1.000000  1.000000  1.000000  1.000000  1.000000                     *
* --------------------------------------------------------------------  *
*   System solution:                                                    *
*  -0.136497 -0.568249 -0.694162  1.202870  0.734290                    *
* --------------------------------------------------------------------  *
*                                                                       *
* Reference: "Numerical Recipes By W.H. Press, B.P. Flannery, S.A. Teu- *
*             kolsky, W.T. Vetterling, Cambridge University Press, 1987 *
*             [BIBLI 08].                                               *
*                                                                       *
*                      C++ version from FORTRAN by J-P Moreau, Paris    *
*                                      (www.jpmoreau.fr)                *
************************************************************************/
#include <basis.h>
#include <vmblock.h>

#define NMAX 100


int  TRIDAG(REAL *A,REAL *B,REAL *C,REAL *R,REAL *U,int N)  {
  // Solves for a vector U of length N the tridiagonal linear set
  // M U = R, where A, B and C are the three main diagonals of matrix
  // M(N,N), the other terms are 0. R is the right side vector.
  
  REAL BET,GAM[NMAX];
  int j;

  if (B[1]==0.0)  return 1;

  BET=B[1];
  U[1]=R[1]/BET;
  for (j=2; j<N+1; j++)  {          // Decomposition and forward substitution
    GAM[j]=C[j-1]/BET;
    BET=B[j]-A[j]*GAM[j];
    if (BET==0.0) return 2;         // Algorithm fails
    U[j]=(R[j]-A[j]*U[j-1])/BET;
  }

  for (j=N-1; j>0; j--)             // Back substitution
    U[j]-=GAM[j+1]*U[j+1];
  
  return 0;
}


int main()  {

  REAL *A, *B, *C, *R, *U;
  int i, n, rc;
  void *vmblock = NULL; 

  FILE *fp1,*fp2;

  fp1 = fopen("tridiag5.dat","r");
  fp2 = fopen("tridiag5.lst","w");
  
  fscanf(fp1,"%d",&n);

  vmblock = vminit();  
  A       = (REAL *) vmalloc(vmblock, VEKTOR,  n+1, 0);  
  B       = (REAL *) vmalloc(vmblock, VEKTOR,  n+1, 0);
  C       = (REAL *) vmalloc(vmblock, VEKTOR,  n+1, 0);  
  R       = (REAL *) vmalloc(vmblock, VEKTOR,  n+1, 0);
  U       = (REAL *) vmalloc(vmblock, VEKTOR,  n+1, 0);  

  if (! vmcomplete(vmblock))
  {
    LogError ("No Memory", 0, __FILE__, __LINE__);
    return 1;
  }

  for (i=1; i<n+1; i++)  {
    fscanf(fp1,"%lf %lf %lf",&A[i],&B[i],&C[i]);
	R[i]=ONE;
  } 

  fclose(fp1);

  WriteHead(fp2,"  Tridiagonal system of equations");
  fprintf(fp2,"  upper subdiagonal:");
  WriteVec1(fp2,n,C);
  fprintf(fp2,"  main diagonal    :");
  WriteVec1(fp2,n,B);
  fprintf(fp2,"  lower subdiagonal:");
  WriteVec1(fp2,n,A);
  fprintf(fp2,"  right side member:");
  WriteVec1(fp2,n,R); 

  
  rc = TRIDAG(A,B,C,R,U,n);

  if (rc != 0) 
    fprintf(fp2,"  ERROR IN TRIDAG : RETURN CODE=%d\n",rc);
  else  {
    WriteHead(fp2,"  System solution:");
	WriteVec1(fp2,n,U);
  }
  WriteEnd(fp2);
  
  printf("\n Results in tridiag5.lst.\n\n");  
  fclose(fp2);
  
  return 0;
}
    
// End tridiag.cpp