!------------------------------------------------------
!   Thermal Convection
!------------------------------------------------------
module convection_data
    integer, parameter :: imx=64, jmx=32, im0=imx-1, jm0=jmx-1, im1=imx+1, jm1=jmx+1
    real   , parameter :: pi=3.1415926535

    real :: psi(im1,jm1), zeta(im0,jm0), temp(imx,jmx), x(im1), z(jm1)
    real :: xmax = 1., zmax = 0.5,  dx,  dz

    real, private ::wx(imx*3+15), wz(jmx*3+15), fac(im0, jm0)

    contains
!------------------- setup data ----------------------
    subroutine setup

      zeta = 0.
      temp = 0.

      dx = xmax/imx 
      dz = zmax/jmx

      x = (/ (dx*i, i=0, imx) /)
      z = (/ (dz*j, j=0, jmx) /)

      do i=1, im0
        xk = pi/xmax*i
        do j=1, jm0	
          zl = pi/zmax*j
          fac(i,j) = -1. /(xk*xk + zl*zl)/(4.*imx*jmx)
        end do
      end do

      call sinti(im0,wx)
      call sinti(jm0,wz)

    end subroutine

!------------------ average_x ------------------------
    function a_x (var)
      real, dimension(:,:) :: var
      real, dimension(size(var,1)-1, size(var,2)) :: a_x

      im  = size(var,1)
      a_x = (var(1:im-1,:) + var(2:im,:))/2.
    end function a_x

!------------------ average_z ------------------------
    function a_z (var)
      real, dimension(:,:) :: var
      real, dimension(size(var,1), size(var,2)-1) :: a_z

      jm  = size(var,2)
      a_z = (var(:,1:jm-1) + var(:,2:jm))/2.
    end function a_z

!--------------- derivative_x ------------------
    function d_x (var)
      real, dimension(:,:) :: var
      real, dimension(size(var,1)-1, size(var,2)) :: d_x

      im  = size(var,1)
      d_x = (var(2:im,:) - var(1:im-1,:))/dx 
    end function d_x
    
!--------------- derivative_z ------------------
    function d_z (var)
      real, dimension(:,:) :: var
      real, dimension(size(var,1), size(var,2)-1) :: d_z

      jm  = size(var,2)
      d_z = (var(:,2:jm) - var(:,1:jm-1))/dz 
    end function d_z
    
!------------------ expand_x ------------------------
    function e_x(var, fac)
      integer :: fac
      real, dimension(:,:) :: var
      real, dimension(size(var,1)+2, size(var,2)) :: e_x

      im = size(var,1)
      e_x(2:im+1, : ) =  var
      e_x(     1, : ) =  var( 1,:)*fac
      e_x(  im+2, : ) =  var(im,:)*fac
    end function e_x

!------------------ expand_z ------------------------
    function e_z(var, fac)
      integer :: fac
      real, dimension(:,:) :: var
      real, dimension(size(var,1), size(var,2)+2) :: e_z

      jm = size(var,2)
      e_z( :, 2:jm+1) =  var
      e_z( :,      1) =  var(:, 1)*fac
      e_z( :,   jm+2) =  var(:,jm)*fac
    end function e_z
    
!----------------- poisson ------------------------
   function poisson(zz)
      real :: poisson(im1,jm1), zz(im0,jm0)

      poisson = 0.
      poisson(2:imx, 2:jmx) = fft_ss(fac*fft_ss(zz))
   end function

!---------------- fft (sin-sin) ---------------------
   function fft_ss(var)
      real, dimension(im0, jm0) :: var, fft_ss

      fft_ss = var
      do j=1, jm0
        call sint(im0, fft_ss(:,j), wx)
      end do
      do i=1, im0
        call sint(jm0, fft_ss(i,:), wz)
      end do
    end function fft_ss
   
end module
!-----------------------------------------------------
!   Thermal convection
!-----------------------------------------------------
module convection_eq
   use convection_data
    real :: p_time=0., dt, alpha=1., beta=-20., nu=5.e-3, kappa=5.e-3

   contains

!------------- Runge-Kutta ------------
    subroutine runge_kutta
      real, dimension(im0,jm0) :: dz1  , dz2  , dz3  , dz4
      real, dimension(imx,jmx) :: dtmp1, dtmp2, dtmp3, dtmp4

      call drv(zeta       , temp         , dz1, dtmp1)
      call drv(zeta+dz1/2., temp+dtmp1/2., dz2, dtmp2)
      call drv(zeta+dz2/2., temp+dtmp2/2., dz3, dtmp3)
      call drv(zeta+dz3   , temp+dtmp3   , dz4, dtmp4)
     
      zeta = zeta + (dz1   + 2.*dz2   + 2.*dz3   + dz4  ) / 6.
      temp = temp + (dtmp1 + 2.*dtmp2 + 2.*dtmp3 + dtmp4) / 6.
      p_time = p_time + dt
    end subroutine

!------------- Vorticity Eq. ------------
    subroutine drv(zz, tt, dz, dtmp)
      real, dimension(im0,jm0) :: zz  , dz
      real, dimension(imx,jmx) :: tt, dtmp
      real :: u(im1, jmx), v(imx, jm1)

      psi = poisson(zz)
      u = -d_z(psi)
      v =  d_x(psi)

      dz = dt* (alpha*d_x(a_z(tt)) &
         &  + nu*(d_x(d_x(e_x(zz,0))) + d_z(d_z(e_z(zz,0)))))

      dtmp = dt*( -d_x(u*a_x(e_x(tt,1))) - d_z(v*a_z(e_z(tt,-1))) - beta*a_z(v) &
         &  + kappa*(d_x(d_x(e_x(tt,1))) + d_z(d_z(e_z(tt,-1)))))
    end subroutine

    function t_all()
      real, dimension(im1,jm1) :: t_all

      t_all = spread(beta*z, 1, im1)  + a_x(a_z(e_x(e_z(temp,-1),1)))

     end function
end module
!-----------------------------------------------------
program main
  use DCL
  use convection_eq
    call setup

!    write(*,*) 'input dt'
!    read  (*,*) dt

    dt = 0.005
    temp (32, 16) = 1.e-1

!  ---- time integration ----

    call DclSetParm('ALTERNATE' , .true.)
    call DclSetParm('WAIT'      , .false.)
    call DclSetParm('USE_FULL_WINDOW', .true.)

    call DclOpenGraphics  (1)
    call DclSetAspectRatio(4., 3.)              ! 縦横比を4:3にする
    call DclSetAxisFactor (0.7)                 ! 座標軸の文字を0.7倍する

    do ii=1, 100
      do j=1, 20
        call runge_kutta
      end do
      psi = poisson(zeta)
      call graph
    end do
    call DclCloseGraphics()

  contains
!------------------ graphic output ------------------
  subroutine graph
    character msg*16

    call DclNewFrame
    call DclSetWindow   ( 0., xmax, 0., zmax )
    call DclSetViewPort (0.07, 0.97, .15, 0.6)
    call DclSetTransFunction

    call DclSetColorRange  (-zmax*abs(beta), 0.)
    call DclPaintData      (t_all())
    call DclDrawScaledAxis
    call DclDrawContour    (psi)

    write(msg,'(a,f10.2)') 't = ', p_time  
    call  DclDrawTitle ('t', msg, 0.0)

  end subroutine
end program