program main
use def
use fft
use init
use scalar
use diff
use advec
use press
use output

complex, dimension(bs) :: uh=0.0,vh=0.0,rhoh=0.0,ph=0.0

call MPI_INIT(ierr)
call MPI_COMM_RANK(MPI_COMM_WORLD,rank,ierr)
call MPI_COMM_SIZE(MPI_COMM_WORLD,nproc,ierr)

call init_fftw_plans()

if (mod(n,nproc)/=0) then
   print *, '***Number of processors must divide evenly into n***'
   stop
elseif (bs.ne.n*n/nproc) then
   print *, '***block size is not correct***'
   stop
elseif (rank==root) then
   print *, '*****************************************************'
   print *, '* 2D Parallel Stratified Navier stokes'
   print *, '* written by James Sandham'
   print *, '*****************************************************'
   print *, '* N = ',n
   print *, '* dt = ',dt
   print *, '* nt: ',nt
   print *, '* Re = ',Re
   print *, '* Thermal diffusivity: ',kappa
   print *, '* Number of processors: ',nproc
   print *, '* Run type: ',run_type
   print *, '*****************************************************'
   call cpu_time(tstart)
endif

call create_netcdf_file()

call initialization(u,v,rho,ik,x,y)
call MPI_BARRIER(MPI_COMM_WORLD,ierr)


call fft1(u,uh)
call fft1(v,vh)
call fft1(rho,rhoh)

do tt=1,nt

  call density(uh,vh,rhoh,sscr1,sscr2,ik) 
  call advection(uh,vh,ascr1,ascr2,ascr3,ascr4,ik)
  call diffusion(uh,vh)
  call pressure(uh,vh,ph,ik)
  call update_pressure(uh,vh,ph,ik)
  call update_density(vh,rhoh)


  if (mod(tt,plot_freq)==0) then
     call ifft1(rhoh,rho)
     call write_to_netcdf_file(real(rho),tt/plot_freq)
     call fft1(rho,rhoh)
     if (rank==root) then
        print *,' Timestep # ',tt,' Model time: ', dt*tt, 'val: ', uh(n+50)
     endif
  endif

enddo

call ifft1(uh,u)
call ifft1(vh,v)
call ifft1(rhoh,rho)
call ifft1(ph,p)
!call create_netcdf_file(real(rho))

if (rank==root) then
   call cpu_time(tfinish)
   telapsed = tfinish-tstart
   print *, '***Total Elapsed Time***: ',telapsed
endif

call dest_fftw_plans()
call MPI_FINALIZE(ierr)
end program main
