/*----------------------------------------------------------------
Tridiagonal Matrix Inversion Source Code
Copyright © 2020 James R. Craig, University of Waterloo
----------------------------------------------------------------*/
#include <math.h>

using namespace std;

//////////////////////////////////////////////////////////////////
/// \brief Inverts Tridiagonal matrix with diagonals cc,a, and b using algorithm from
/// Usmani,R. A. (1994). "Inversion of a tridiagonal jacobi matrix". Linear Algebra and its Applications. 212-213: 413–414.
/// assumes b[N] is zero, and all three indices correspond to matrix row - performs a shift to standard tridiagonal form
/// \param cc   [in ] shifted left diagonal  (cc[i+1]=A[i][i-1]) [length=size]
/// \param a    [in ] diagonal       (a[i]=A[i][i]) [length=size]
/// \param b    [in ] right diagonal (b[i]=A[i][i+1] [length=size]
/// \param Ainv [out] resultant inverse matrix (assumes memory is already allocated in array of doubles)
/// \param size [in]  N, size of NxN square matrix 
///
void InvertTridiagonal(const double *cc,const double *a,const double *b,
                       double **Ainv,const int     size)
{
  int i,j,k;
  int N=size;
  double th_im1,th_jm1,ph_ip1,ph_jp1;
  double bprod,cprod;
  double *theta=new double[N];
  double *phi  =new double[N];
  double *c    =new double[N];
  for(i=0;i<N;i++) { c[i]=cc[i+1]; } //requires a shift of the off-diagonal 

  theta[0]=a[0];
  theta[1]=a[1]*theta[0]-b[0]*c[0];
  for(i=2;i<size;i++) { theta[i]=a[i]*theta[i-1]-b[i-1]*c[i-1]*theta[i-2]; }
  phi[N-1]=a[N-1];
  phi[N-2]=a[N-2]*phi[N-1]-b[N-2]*c[N-2];
  for(i=N-3;i>=0;i--) { phi  [i]=a[i]*phi  [i+1]-b[i  ]*c[i  ]*phi  [i+2]; }
  for(i=0;i<N;i++) 
  {
    if(i==0)   { th_im1=1.0; }
    else       { th_im1=theta[i-1]; }
    if(i==N-1) { ph_ip1=1.0; }
    else       { ph_ip1=phi[i+1]; }
    for(j=0;j<N;j++) 
    {
      if(j==0)   { th_jm1=1.0; }
      else       { th_jm1=theta[j-1]; }
      if(j==N-1) { ph_jp1=1.0; }
      else       { ph_jp1=phi[j+1]; }

      if(i<j) {
        bprod=1.0; for(k=i;k<=j-1;k++) { bprod*=b[k]; }
        Ainv[i][j]=pow(-1.0,i+j+2)*bprod*th_im1*ph_jp1/theta[N-1];
      }
      else if(i==j) {
        Ainv[i][j]=th_im1*ph_jp1/theta[N-1];
      }
      else {
        cprod=1.0; for(k=j;k<=i-1;k++) { cprod*=c[k]; }
        Ainv[i][j]=pow(-1.0,i+j+2)*cprod*th_jm1*ph_ip1/theta[N-1];
      }
    }
  }
  delete[] theta;
  delete[] phi;
  delete[] c;
}