program heat2D
!$use OMP_LIB
implicit none

	integer :: niter,i,j,t,rank
	real(8), allocatable, dimension(:,:) :: U0,U1
	real(8) :: alpha,totaltime,Dt,Stab,DtAlpha
	integer, dimension(1:2) :: Nx
	real(8), dimension(1:2) :: Lx,Dx,DxDx,InvDxDx
	character(len=4) :: charbuff

	niter = 8000

	Nx(:) = 700
	alpha = 0.4d0
	totaltime = 0.005d0
	Dt = totaltime/((niter-1)*1.0d0)
	Lx(:) = 1.0d0
	Dx(:) = Lx(:)/((Nx(:)-1)*1.0d0)
	InvDxDx(:) = 1.0d0/(Dx(:)*Dx(:))
	DxDx(:) = Dx(:)*Dx(:)
	Stab = alpha*Dt*(InvDxDx(1)) + alpha*Dt*(InvDxDx(2))
	DtAlpha = Dt*alpha
	print *, "Stability factor (need < 0.5): ",stab
	if(stab > 0.5d0) stop
	print *,"Calculating, please wait..."
	allocate(U0(0:Nx(1)+1,0:Nx(2)+1),U1(0:Nx(1)+1,0:Nx(2)+1))

	U0(:,:) = 1.0d0
	U0(100:150,200:250) = -10.0d0
	U0(300:350,400:450) = +12.0d0
	U1(:,:) = U0(:,:)

	rank = 0
	write(charbuff,'(I4.4)') rank
	open(10+rank,file=trim(adjustl('IN'//'.'//trim(adjustl(charbuff))//'.dat')))
	do j=1,Nx(2),10
		do i=1,Nx(1),10
			write(10,*) i,j+(rank*Nx(2)),U0(i,j)
		end do
	end do
	close(10+rank)

!$OMP PARALLEL DEFAULT(SHARED) PRIVATE(i,j,t)
	do t=1,niter
		! print *,t

!$OMP DO SCHEDULE(RUNTIME)
		do j=1,Nx(2)
			do i=1,Nx(1)
				U1(i,j) = U0(i,j) + DtAlpha* &
				     & ( (U0(i+1,j)-2.0*U0(i,j)+U0(i-1,j))*InvDxDx(1) + &
			         &   (U0(i,j+1)-2.0*U0(i,j)+U0(i,j-1))*InvDxDx(2) )

			end do
		end do
!$OMP END DO

!$OMP DO SCHEDULE(RUNTIME)
		do j=1,Nx(2)
			do i=1,Nx(1)
				U0(i,j) = U1(i,j) + DtAlpha* &
				     & ( (U1(i+1,j)-2.0*U1(i,j)+U1(i-1,j))*InvDxDx(1) + &
				     &   (U1(i,j+1)-2.0*U1(i,j)+U1(i,j-1))*InvDxDx(2) )

			end do
		end do
!$OMP END DO

	end do
!$OMP END PARALLEL

	rank = 0
	write(charbuff,'(I4.4)') rank
	open(10+rank,file=trim(adjustl('OUT'//'.'//trim(adjustl(charbuff))//'.dat')))
	do j=1,Nx(2),10
		do i=1,Nx(1),10
			write(10,*) i,j+(rank*Nx(2)),U0(i,j)
		end do
	end do
	close(10+rank)

	deallocate(U0,U1)

end program heat2D