! **************************************************************

program mainprogram
        implicit none
        
        integer :: Nx,Ny
        parameter (Nx = 201, Ny = 101)
               
        double precision :: dx,dy,x_val,y_val,Lx,Ly,dt,t_final
        double precision :: err1,A0
        double precision, dimension(1:Nx,1:Ny) :: f_source,C,C_temp,RHS
        
        integer :: n_timesteps,iteration_time,sor_iteration,file_period_time, &
               max_iteration
        parameter(max_iteration=50)
        
        integer :: i,j,im1,ip1
        
        character*60 filename
        
! **************************************************************
! Simulation parameters

        Lx=2.d0
        Ly=1.d0
        
        A0=10.d0
        
        dx=Lx/dble(Nx-1)
        dy=Ly/dble(Ny-1)
        
        dt=1.d-3
        t_final=0.4d0
        n_timesteps=floor(t_final/dt)
        file_period_time=10
        
! **************************************************************
! Compute source, initialise guess

        f_source=0.d0
        C=0.d0
        C_temp=0.d0
        
        write (*,*) 'initializing concentration'
        call get_C_init(C,Nx,Ny,dx,dy,Lx,Ly)
        write (*,*) 'done'
        
        write (*,*) 'getting source'
        call get_f_periodic(f_source,Nx,Ny,dx,dy,Lx,Ly,A0)
        write (*,*) 'done'
        
! **************************************************************
! Enter into time loop now
        
        do iteration_time=1,n_timesteps

          call get_RHS(RHS,C,dt,dx,dy,Nx,Ny)
          RHS=RHS+dt*f_source
          
          do sor_iteration=1,max_iteration
            C_temp=C
            call do_sor_C(C,RHS,dt,dx,dy,Nx,Ny)
            
            ! Implement Dirichlet conditions at y=0,y=L_y.
            C(:,1)=C(:,2)
            C(:,Ny)=C(:,Ny-1)
          end do
        
          call get_diff(C,C_temp,Nx,Ny,err1)
          write(*,*) iteration_time, ' Residual is ', err1
          
          ! ****************************************************
          ! write result to file
        
          if(mod(iteration_time,file_period_time).eq.0)then
            write(*,*) 'writing to file'
            write(filename,'(A,I9,A)')'diffusion_',iteration_time,'.dat'
            write(*,*) filename
            open(unit=9,file=filename,status='unknown')
          
            do j=1,Ny
              do i=1,Nx
                x_val=(i-1)*dx
                y_val=(j-1)*dy
                Write(9,*) x_val, y_val, C(i,j)
              end do
            end do
            
            close(9)
            write(*,*) 'done'
          end if
          
          ! ****************************************************
          ! Step out of time loop now
          
        end do
      
end program mainprogram


! **************************************************************
! **************************************************************
! Diffusion subroutines

        subroutine get_RHS(RHS,C,dt,dx,dy,Nx,Ny)
        implicit none
  
        integer ::i,j,ip1,im1,Nx,Ny
        double precision :: dx,dy,ax,ay,dt
        double precision :: RHS(1:Nx,1:Ny),C(1:Nx,1:Ny),Diffusion(1:Nx,1:Ny)
        
        ax=dt/(dx*dx)
        ay=dt/(dy*dy)
        
        do j=2,Ny-1
          do i=1,Nx
          
            if(i.eq.1)then
              im1=Nx-1
            else
              im1=i-1
            end if

            if(i.eq.Nx)then
              ip1=2
             else
               ip1=i+1
             end if
          
            Diffusion(i,j)=ax*( C(ip1,j)+C(im1,j)-2.d0*C(i,j))+ay*( C(i,j+1)+C(i,j-1)-2.d0*C(i,j) )
          end do
        end do
        
        RHS=C+0.5d0*Diffusion
        
        end subroutine get_RHS

! ***************************************************************

        subroutine do_sor_C(C,RHS,dt,dx,dy,Nx,Ny)
        implicit none
  
        integer ::i,j,ip1,im1,Nx,Ny
        double precision :: dx,dy,dt,ax,ay
        double precision :: RHS(1:Nx,1:Ny),C(1:Nx,1:Ny)
        
        double precision :: relax,diag_val,temp_val
        
        relax=1.5d0
        
        ax=dt/(dx*dx)
        ay=dt/(dy*dy)
        diag_val=1+ax+ay
        
        do j=2,Ny-1
          do i=1,Nx
          
            if(i.eq.1)then
              im1=Nx-1
            else
              im1=i-1
            end if

            if(i.eq.Nx)then
              ip1=2
             else
               ip1=i+1
             end if
             
             temp_val=(ax/2.d0)*(C(ip1,j)+C(im1,j))+(ay/2.d0)*(C(i,j+1)+C(i,j-1))+RHS(i,j)
             C(i,j)=C(i,j)*(1.d0-relax)+relax*temp_val/diag_val
           end do
         end do
         
         end subroutine do_sor_C
             

! **************************************************************
! **************************************************************

  subroutine get_f_periodic(f_src,Nx,Ny,dx,dy,Lx,Ly,A0)
  implicit none
  
  integer ::i,j,Nx,Ny
  double precision :: dx,dy,Lx,Ly,x_val,y_val, pi=3.1415926535
  double precision :: kx0,ky0,kx,ky,A0
  double precision :: f_src(1:Nx,1:Ny)
  
  kx0=2.d0*pi/Lx
  ky0=pi/Ly
  
  kx=kx0
  ky=3.d0*ky0
  
  f_src=0.d0
  do j=1,Ny
    do i=1,Nx
      x_val=(i-1)*dx
      y_val=(j-1)*dy
      f_src(i,j)=-A0*cos(kx*x_val)*cos(ky*y_val)
    end do
  end do
  
  return
  end subroutine get_f_periodic

! **************************************************************

  subroutine get_C_init(C,Nx,Ny,dx,dy,Lx,Ly)
  
  implicit none
  
  integer ::i,j,Nx,Ny
  double precision :: dx,dy,Lx,Ly,x_val,y_val, pi=3.1415926535
  double precision :: kx0,ky0,kx,ky
  double precision :: C(1:Nx,1:Ny)
  
  kx0=2.d0*pi/Lx
  ky0=pi/Ly
    
  C=0.d0
  do j=1,Ny
    do i=1,Nx
      x_val=(i-1)*dx
      y_val=(j-1)*dy
      C(i,j)=cos(kx0*x_val)*cos(ky0*y_val)+cos(2*kx0*x_val)*cos(ky0*y_val)+cos(kx0*x_val)*cos(4*ky0*y_val)
    end do
  end do
  
  return
  end subroutine get_C_init
  
! **************************************************************

  subroutine get_diff(C,C_old,Nx,Ny,diff)
  implicit none

  double precision :: diff,sum
  integer :: Nx,Ny,i,j
  double precision, dimension(1:Nx,1:Ny) :: C, C_old

  sum = 0.0D0
  Do j = 1, Ny
    Do i = 1, Nx
      sum = sum + (C(i,j)-C_old(i,j))**2
    End Do
  End Do
  diff = sum
  
  Return
  End subroutine get_diff

! **************************************************************
! **************************************************************