#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <iostream>
using namespace std;

void ThomasAlgorithm_P(int mynode, int numnodes,
                   int N, double *b, double *a, double *c, double *x, double *q){
      int i;
      int rows_local,local_offset;
      double S[2][2],T[2][2],s1tmp,s2tmp;
      double *l,*d,*y;
      MPI_Status status;

      l = (double*)malloc(sizeof(double) * N);
      d = (double*)malloc(sizeof(double) * N);
      y = (double*)malloc(sizeof(double) * N);

      for(i=0;i<N;i++)
        l[i] = d[i] = y[i] = 0.0;

      S[0][0] = S[1][1] = 1.0;
      S[1][0] = S[0][1] = 0.0;

      rows_local = (int) floor((double)N/numnodes);
      local_offset = mynode*rows_local;

      if(mynode==0){
        s1tmp = a[local_offset]*S[0][0];
        S[1][0] = S[0][0];
        S[1][1] = S[0][1];
        S[0][1] = a[local_offset]*S[0][1];
        S[0][0] = s1tmp;
        for(i=1;i<rows_local;i++){
          s1tmp = a[i+local_offset]*S[0][0] - b[i+local_offset-1]*c[i+local_offset-1]*S[1][0];
          s2tmp = a[i+local_offset]*S[0][1] - b[i+local_offset-1]*c[i+local_offset-1]*S[1][1];
          S[1][0] = S[0][0];
          S[1][1] = S[0][1];
          S[0][0] = s1tmp;
          S[0][1] = s2tmp;
        }
      }
      else{
        for(i=0;i<rows_local;i++){
          s1tmp = a[i+local_offset]*S[0][0] - b[i+local_offset-1]*c[i+local_offset-1]*S[1][0];
          s2tmp = a[i+local_offset]*S[0][1] - b[i+local_offset-1]*c[i+local_offset-1]*S[1][1];
          S[1][0] = S[0][0];
          S[1][1] = S[0][1];
          S[0][0] = s1tmp;
          S[0][1] = s2tmp;
        }
      }

      for(i=0; i<=log2(numnodes);i++){
        if((mynode+pow(2.0,i)) < numnodes){
          MPI_Send(S,4,MPI_DOUBLE,int(mynode+pow(2.0,i)),0,MPI_COMM_WORLD);
        }
        if(mynode-pow(2.0,i)>=0){
          MPI_Recv(T,4,MPI_DOUBLE,int(mynode-pow(2.0,i)),0,MPI_COMM_WORLD,&status);
          s1tmp = S[0][0]*T[0][0] + S[0][1]*T[1][0];
          S[0][1] = S[0][0]*T[0][1] + S[0][1]*T[1][1];
          S[0][0] = s1tmp;
          s1tmp = S[1][0]*T[0][0] + S[1][1]*T[1][0];
          S[1][1] = S[1][0]*T[0][1] + S[1][1]*T[1][1];
          S[1][0] = s1tmp;
        }
      }

      d[local_offset+rows_local-1] = (S[0][0] + S[0][1])/(S[1][0] + S[1][1]);

      if(mynode == 0){
        MPI_Send(&d[local_offset+rows_local-1],1,MPI_DOUBLE,1,0,MPI_COMM_WORLD);
      }
      else{
        MPI_Recv(&d[local_offset-1],1,MPI_DOUBLE,mynode-1,0,MPI_COMM_WORLD,&status);
        if(mynode != numnodes-1)
          MPI_Send(&d[local_offset+rows_local-1],1,MPI_DOUBLE,mynode+1,0,MPI_COMM_WORLD);
      }


      if(mynode == 0){
        l[0] = 0;
        d[0] = a[0];
        for(i=1;i<rows_local-1;i++){
          l[local_offset+i] = b[local_offset+i-1]/d[local_offset+i-1];
          d[local_offset+i] = a[local_offset+i] - l[local_offset+i]*c[local_offset+i-1];
        }
        l[local_offset+rows_local-1] = b[local_offset+rows_local-2]/d[local_offset+rows_local-2];
      }
      else{
        for(i=0;i<rows_local-1;i++){
          l[local_offset+i] = b[local_offset+i-1]/d[local_offset+i-1];
          d[local_offset+i] = a[local_offset+i] - l[local_offset+i]*c[local_offset+i-1];
        }
        l[local_offset+rows_local-1] = b[local_offset+rows_local-2]/d[local_offset+rows_local-2];
      }

      /***********************************************************************************/

      if(mynode>0)
        d[local_offset-1] = 0;

      double * tmp = new double[N];
      for(i=0;i<N;i++)
        tmp[i] = d[i];
      MPI_Allreduce(tmp,d,N,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
      for(i=0;i<N;i++)
        tmp[i] = l[i];
      MPI_Allreduce(tmp,l,N,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
      delete[] tmp;

      if(mynode ==0){
        /* Forward Substitution [L][y] = [q] */
        y[0] = q[0];
        for(i=1;i<N;i++)
          y[i] = q[i] - l[i]*y[i-1];

        /* Backward Substitution [U][x] = [y] */
        x[N-1] = y[N-1]/d[N-1];
        for(i=N-2;i>=0;i--)
          x[i] = (y[i] - c[i]*x[i+1])/d[i];

      }

     if(mynode ==0){
          for(i=0;i<40;i++){
            cout << x[i];
            cout << endl;
          }
      }


      free(l);
      free(y);
      free(d);
      return;
    }

int main(int argc,char *argv[])
{
     int numnodes, mynode, sendcount, recvcount, source;
     int N,i;
    double b[] = {0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048};
    double c[] = {0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048,0.048};
    double a[] = {-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096,-1.096};
    double q[] = {-0.206311,-0.408774,-0.603633,-0.787322,-0.95655,-1.10838,-1.2403,-1.35031,-1.43691,-1.49921,-1.53691,-1.55031,-1.5403,-1.50838,-1.45655,-1.38732,-1.30363,-1.20877,-1.15431};
    double x[] = {0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0};

     MPI_Init(&argc,&argv);
     MPI_Comm_rank(MPI_COMM_WORLD, &mynode);
     MPI_Comm_size(MPI_COMM_WORLD, &numnodes);

    N = 19;

    ThomasAlgorithm_P(mynode,numnodes,N,b,a,c,x,q);

     MPI_Finalize();
}