MODULE prec_m
  INTEGER,PARAMETER:: wp=SELECTED_REAL_KIND(12)
END MODULE prec_m


!
! tridiag. matrix     | B1 A2        |
!                     | A2 B2 A3     |
!                     |    A3 B3     |
!                     |        ...   |
!                     |        AN BN |
!
! is stored in  a(:,:) as follows:
!
!               B1 B2 B3 .... BN
!               *  A2 A3 .... AN
!
MODULE tridiag_m
  USE prec_m
  IMPLICIT NONE
CONTAINS
  SUBROUTINE tridiag_cholesky_factor( a, info )
    INTEGER,  INTENT(OUT)  :: info
    REAL(wp), INTENT(INOUT):: a(:,:)
    REAL(wp):: tmp
    INTEGER:: i, n
    info = -1
    n = SIZE(a,2)
    IF ( a(1,1) < 0.0 ) RETURN
    a(1,1) = SQRT(a(1,1))
    DO i=2,n
       a(2,i) = a(2,i)/a(1,i-1)
       tmp = a(1,i) - a(2,i)*a(2,i)
       IF ( tmp < 0.0 ) RETURN
       a(1,i) = SQRT( tmp )
    END DO
    info = 0
    RETURN
  END SUBROUTINE tridiag_cholesky_factor

  SUBROUTINE tridiag_cholesky_solve( a, x, b )
    REAL(wp), INTENT(IN) :: a(:,:), b(:)
    REAL(wp), INTENT(OUT):: x(:)
    INTEGER:: i, n
    n = SIZE(a,2)
    x(1) = b(1)/a(1,1)
    DO i=2,n
       x(i) = (b(i)-a(2,i)*x(i-1))/a(1,i)
    END DO
    x(n) = x(n)/a(1,n)
    DO i=n-1,1,-1
       x(i) = (x(i)-a(2,i+1)*x(i+1))/a(1,i)
    END DO
  END SUBROUTINE tridiag_cholesky_solve

  SUBROUTINE tridiag_mxv( a, x, ax )
    REAL(wp), INTENT(IN) :: a(:,:), x(:)
    REAL(wp), INTENT(OUT):: ax(:)
    INTEGER:: i, n
    n = SIZE(a,2)
    ax(1) = a(1,1)*x(1) + a(2,1)*x(2)
    DO i=2,n-1
       ax(i) = a(2,i)*x(i-1) + a(1,i)*x(i) + a(2,i+1)*x(i+1)
    END DO
    ax(n) = a(2,n)*x(n-1) + a(1,n)*x(n)
  END SUBROUTINE tridiag_mxv
END MODULE tridiag_m



!
! Solves 2D heat equation in a unit square
!
!     u_t - c*(u_xx + u_yy) = 0   0<x<1, 0<y<1, 0<t<=T
!     u=0   on the boundary
!     u(x,y,0)=u_0(x,y)
!
! using alternate direction implicit (ADI) method.
!
!   n = number of spatial grid points in x and y direction
!   nt = number of time steps
!
MODULE adi_m
  USE prec_m
  USE tridiag_m
  IMPLICIT NONE
CONTAINS

  SUBROUTINE adi(u0,u,c,T,n,nt)
    REAL(wp), INTENT(IN) :: u0(:,:)
    REAL(wp), INTENT(OUT):: u(:,:)
    REAL(wp), INTENT(IN) :: c, T
    INTEGER,  INTENT(IN) :: n, nt
!
    REAL(wp):: a1(2,n-2), a2(2,n-2), tmpvec1(n-2), tmpvec2(n-2)
!    
    REAL(wp):: r, h, dt
    INTEGER:: i, j, k, info
!
    h = 1.0_wp/(n-1)
    dt = T/nt
    r = c*dt/(2*h*h)
    u = u0
    a1(1,:) = (1+2*r)
    a1(2,:) = -r
    a2(1,:) = (1-2*r)
    a2(2,:) = r
    CALL tridiag_cholesky_factor( a1, info )
    IF ( info < 0 ) STOP "Fatal error: Matrix not pos. def"
    DO k=1,nt
       !$omp parallel private (tmpvec1, tmpvec2)
       !$omp do
       DO j=2,n-1
          tmpvec1 = u(j,2:n-1)
          CALL tridiag_mxv(a2,tmpvec1,tmpvec2)
          CALL tridiag_cholesky_solve( a1, tmpvec1, tmpvec2 )
          u(j,2:n-1) = tmpvec1
       END DO
       !$omp end do
       !$omp do
       DO i=2,n-1
          tmpvec1 = u(2:n-1,i)
          CALL tridiag_mxv(a2,tmpvec1,tmpvec2)
          CALL tridiag_cholesky_solve( a1, tmpvec1, tmpvec2 )
          u(2:n-1,i) = tmpvec1
       END DO
       !$omp end do
       !$omp end parallel
    END DO
  END SUBROUTINE adi
END MODULE adi_m

PROGRAM main
  USE adi_m
  USE omp_lib
  INTEGER:: n, nt, i, j
  REAL(wp):: h, c=0.1_wp, T=0.1_wp, t2
  REAL(wp), ALLOCATABLE:: u(:,:), u0(:,:)
  INTEGER:: ij(2)
  WRITE(*,*) "Enter n and nt"
  READ(*,*) n, nt
  ALLOCATE( u(n,n), u0(n,n) )
  h=1.0_wp/(n-1)
  CALL initial_value(u0)
!$  t2 = omp_get_wtime()
  CALL adi(u0,u,c,T,n,nt)
!$  write(*,*) "Elapsed time =",omp_get_wtime()-t2
  ij = MAXLOC(u)
  WRITE(*,'(A,2F5.2,A,F6.3)') 'max u at',ij(1)*h,ij(2)*h,' = ',u(ij(1),ij(2))
  CALL write_sol(u)
  !
  DEALLOCATE( u0, u )
CONTAINS
  SUBROUTINE initial_value(u0)
    REAL(wp):: u0(:,:)
    REAL(wp):: x, y
    INTEGER:: i,j
    DO j=1,n
       y = (j-1)*h
       DO i=1,n
          x = (i-1)*h
          IF ( x < 0.5 .AND. x > 0.2 .AND. y > 0.1 .AND. y < 0.4 ) THEN
             u0(i,j)=1.0_wp
          ELSE
             u0(i,j)=0.0_wp
          END IF
       END DO
    END DO
  END SUBROUTINE initial_value

  SUBROUTINE write_sol(u)
    REAL(wp):: u(:,:)
    INTEGER:: i,j
    OPEN( 10, FILE="u.dat" )
    DO i=1,n
       DO j=1,n
          WRITE(10,"(3E15.5)") (i-1)*h, (j-1)*h, u(i,j)
       END DO
       WRITE(10,*) " "
    END DO
    CLOSE(10)
  END SUBROUTINE write_sol
END PROGRAM main