!=============================================================================!
! IPP NUBEAM Driver                                                           !
!                                                                             !
! by Michael Kraus (michael.kraus@ipp.mpg.de)                                 !
! Extended: Giovanni Tardini January 2010                                     !
!                                                                             !
! Comments:                                                                   !
!   - variables tagged '(nml)' are defined in the namelist                    !
!   - ps means plasmastate                                                    !
!                                                                             !
!=============================================================================!

PROGRAM mpi_nubeam_driver

use plasma_state_mod ! Plasma State (PS) module definitions & methods
use nbi_pstypes      ! NUBEAM control data structures & methods
use nbi_types        ! definition of NUBEAM output data types
use logmod           ! A set of utilities for writing to a log file
use netcdf
use nbdrive_common

implicit none

!======================== 
! Variables' declaration
!========================

character(len=128) :: arg, argval  ! temporary argument strings
integer :: command_argument_count,npar,ipar  ! command line argument

!==============================  
! Run configuration:

character(len=32)  :: mrunid               ! (command line onput)
character(len=128) :: namelist_file = ''   ! (from runid)
character(len=32)  :: shot_str             ! (parsed from mrunid)
character(len=32)  :: tail                 ! (parsed from mrunid)

integer, parameter :: nisp  = 1   ! number of main ion species
integer :: nshot                  ! (convert shot_str)
integer :: ABeam = 2              ! NB ion mass (nml) [a.m.u.]
integer :: ZBeam = 1              ! NB ion Z    (nml)
integer, parameter :: n_beams = 8 ! maximum number of beam lines
integer    :: debug = 0           ! debug level (nml)
real(KIND=rspec) :: nbi_kvolt(n_beams)     ! beam voltage [kV]  (nml)
real(KIND=rspec) :: nbi_full(n_beams)      ! full energy fraction  (nml)
real(KIND=rspec) :: nbi_half(n_beams)      ! half energy fraction  (nml)
real(KIND=rspec) :: curtot=1e6,czeff=1,zimp=1     ! total plasma current [A] (nml)
real, parameter :: min_ne=1.e18,min_te=1.e-1,min_ti=1.e-1,min_vt=-1.e10,min_zef=1.
real, parameter :: min_ipl=1.e5,min_rb=-40.0,min_p=1.e3,min_q=1.0
real, parameter :: max_ne=2.e20,max_te=5.e1,max_ti=5.e1,max_vt= 1.e10,max_zef=10.
real, parameter :: max_ipl=2.e6,max_rb=40.0,max_p=1.e7,max_q=15.0
real :: max_var
real :: outtim(9)=(/ 0.,0.,0.,0.,0.,0.,0.,0.,0. /)

!==============================  
! Filenames:

character(len=120) :: workpath   ! path to load/save all I/O data
character(len=120) :: d_path     ! location of label diles
character(len=60) :: home_dir
character(len=120) :: out1d,out2d
character(len=100) :: psfile_in ! Output file of initial plasma state (nml)
character(len=100) :: psfile_out! Output file of final plasma state   (nml)
character(len=100) :: fbm_out   ! Output file of distribution function
character(len=1)   :: dummy
character(len=80)  :: temp_dir  ! temporary directory for output
character(len=80)  :: temp_root ! temporary pathname variable
character(len=100) :: temp_out  ! temporary filename variable

integer, parameter :: outf = 74,inf  = 75, u1d = 76, u2d = 77 ! File handler
integer    :: ios               ! file I/O status

!==============================  
! UFiles:

character(len=4)   :: idev          ! Device name
character(len=120) :: inputdir=' '  ! path to Ufiles
character(len=16)  :: precur,prerbz,prerrz,prezrz,preqpr,prepre
character(len=16)  :: prenb2,preter,prener,preti2,prevp2=' ',prezef=' '
character(len=32)  :: extcur,extrbz,extrrz,extzrz,extqpr,extpre
character(len=32)  :: extnb2,extter,extner,extti2,extvp2=' ',extzef=' '
character(len=3)   :: dianam='EQI'
character(len=4)   :: expnam='EQXY'

integer, parameter :: nrho_in=51
integer :: ntime, nrho_eq, nth_eq, nrho, nrho_nbi ! number of time, rho and theta steps
integer  :: iochan            ! to read Ufiles
integer  :: nscal             ! Temporary Ufile scalars
integer  :: nx1, nx2, nx3     ! temporary grid sizes for reading/writing UFiles

real, dimension(:),     allocatable :: time_u, rho_tor_u, th_u, beam_label
real, dimension(:),     allocatable :: d1d
real, dimension(:,:),   allocatable :: d2d
real, dimension(:,:,:), allocatable :: d3d
REAL(KIND=rspec), ALLOCATABLE :: time(:), rho(:), rho_nbi(:)   ! time, rho grid (I, O profiles)
REAL(KIND=rspec) :: rho_in(nrho_in)
REAL(KIND=rspec), ALLOCATABLE :: rho_eq(:), th_eq(:)   ! rho, theta grids (equilibrium)

!==============================  
! Input control(s):

type (nbi_ctrl_init)   :: nb_init   ! initialization control object
type (nbi_ctrl_update) :: nb_update ! update control object

!==============================  
! Output objects:

type(nbotype_cptimes)    :: zcptimes    ! CPU time usage
type(nbotype_powers)     :: zpower_ids  ! Xplasma IDs for detailed power profiles (Plasma State contains total heating Pbe, Pbi, etc.).
type(nbotype_deposition) :: zdeposition ! Xplasma deposition profiles
type(nbotype_density)    :: zdensity    ! Xplasma current profiles
type(nbotype_currents)   :: zcurrents   ! Xplasma current profiles
type(nbotype_powers)     :: zpowers     ! Xplasma power profiles
type(nbitype_profiles)   :: zprofiles 

integer      :: inubeam

!==============================  
! runtime variables:

integer :: status        ! return code: 0=OK
integer :: iErr          ! error code:  0=OK
integer :: it, irho, ith, isp, ib, j_fbm=1  ! time, rho, theta, ion species, beam, nubeam, fbm-output loop indices
integer :: tind          ! temporary string index (used for command argument list parsing)

real(KIND=rspec) :: tinit, ftime   ! start and stop time (nml)
real(KIND=rspec) :: timestep       ! timestep (nml)
integer :: irhob
integer :: n_beams_use   
!==============================  
! temporary variables for output profiles

real(KIND=rspec), ALLOCATABLE   :: fs_area(:), fs_vol(:), fs_data(:)

!==============================  
! physical constants
real(KIND=rspec), parameter :: ev = 1.6022e-19   ! unit charge in C

!==============================  
! NetCDF variables

character(len=120) :: netcdf_out  ! filename of NetCDF output file

integer :: ncid           ! id of NetCDF output file
integer :: dimids1(1),dimids2(2),dimids3(3)  ! array to contain dim ids for 1D,2D,3D data

character (len = *), parameter :: UNITS   = "units"
character (len = *), parameter :: DESC   = "description"
character (len = *), parameter :: TIME_NAME  = "TIME"
character (len = *), parameter :: RHO_NAME  = "RHO"
character (len = *), parameter :: RHO_EQ_NAME = "RHO_EQ"
character (len = *), parameter :: TH_EQ_NAME = "TH_EQ"
character (len = *), parameter :: RHO_NBI_NAME = "RHO_NBI"
character (len = *), parameter :: RSURF_NAME = "RSURF"
character (len = *), parameter :: ZSURF_NAME = "ZSURF"
character (len = *), parameter :: TIME_UNITS = "s"
character (len = *), parameter :: RHO_UNITS  = "-"
character (len = *), parameter :: RHO_EQ_UNITS = "-"
character (len = *), parameter :: TH_EQ_UNITS = "rad"
character (len = *), parameter :: RHO_NBI_UNITS = "-"
character (len = *), parameter :: RSURF_UNITS = "m"
character (len = *), parameter :: ZSURF_UNITS = "m"
character (len = *), parameter :: RSURF_DESC = "Major radius surf(t,rho,theta)"
character (len = *), parameter :: ZSURF_DESC = "z  surf(t,rho,theta)"

integer :: time_dimid, rho_dimid, rho_eq_dimid, th_eq_dimid, rho_nbi_dimid
integer :: time_varid, rho_varid, rho_eq_varid, th_eq_varid, rho_nbi_varid
integer :: area_varid, vol_varid, darea_varid,  dvol_varid
integer :: varid, rsurf_varid, zsurf_varid
integer :: te_varid, ti_varid, timp_varid, ne_varid, ni_varid, nimp_varid
integer :: q_varid,  p_varid,  rbphi_varid
integer :: ib_varid, uib_varid, itot_varid, eheat_varid, iheat_varid, heat_varid
integer :: bdep_varid, pbe_varid, pbi_varid, curb_varid, ucurb_varid
integer :: bdens_varid,ufastpa_varid,ufastpp_varid

character (len = 10) :: var_name,var_coord
character (len = 32) :: var_desc
character (len = 16) :: var_unit

real(KIND=rspec), ALLOCATABLE :: TE_IN(:,:), TI_IN(:,:), TIMP_IN(:,:), VTOR_IN(:,:)
real(KIND=rspec), ALLOCATABLE :: NE_IN(:,:), NI_IN(:,:), NIMP_IN(:,:), ZEFF_IN(:,:)
real(KIND=rspec), ALLOCATABLE :: IPL_IN(:), PNBI_IN(:,:)
real(KIND=rspec), ALLOCATABLE :: RSURF_IN(:,:,:),  ZSURF_IN(:,:,:)
real(KIND=rspec), ALLOCATABLE :: Q_IN(:,:),  P_IN(:,:),  RBPHI_IN(:,:)
real(KIND=rspec), ALLOCATABLE :: TE_OUT(:,:), TI_OUT(:,:), TIMP_OUT(:,:), VTOR_OUT(:,:)
real(KIND=rspec), ALLOCATABLE :: NE_OUT(:,:), NI_OUT(:,:), NIMP_OUT(:,:), ZEFF_OUT(:,:)
real(KIND=rspec), ALLOCATABLE :: Q_OUT(:,:),  P_OUT(:,:),  RBPHI_OUT(:,:)
real(KIND=rspec), ALLOCATABLE :: AREA_OUT(:,:), DAREA_OUT(:,:), VOL_OUT(:,:), DVOL_OUT(:,:)
real(KIND=rspec), ALLOCATABLE :: IB_OUT(:), UIB_OUT(:), ITOT_OUT(:)
real(KIND=rspec), ALLOCATABLE :: EHEAT_OUT(:), IHEAT_OUT(:), HEATTOT_OUT(:)
real(KIND=rspec), ALLOCATABLE :: BDEP_OUT(:,:), PBE_OUT(:,:), PBI_OUT(:,:)
real(KIND=rspec), ALLOCATABLE :: BDENS_OUT(:,:), UFASTPA_OUT(:,:), UFASTPP_OUT(:,:)
real(KIND=rspec), ALLOCATABLE :: CURB_OUT(:,:), UCURB_OUT(:,:)
!misc
integer :: nbsfinish ! dealing with beam fast ion birth source file
                     ! if set to '1' -- subroutine to finish up 
                     ! related files on disk in workpath directory
                     ! will be called

integer :: i_trk_dep0
character*300 :: suffix_del !suffix to the file/s to be deleted

!==============================  
! MPI variables

#ifdef __MPI
 include 'mpif.h'
 character*(MPI_MAX_PROCESSOR_NAME) proc_name
#endif

integer :: myid,nbi_comm
integer :: iresult,numprocs
character*200 :: env_preactdir,env_adasdir,log_file
character*2   :: myid_str

!==============================  
! namelist parameters

NAMELIST /run/ &
   & tinit, ftime, timestep, debug, outtim

NAMELIST /plasma/ curtot,czeff

NAMELIST /nbi/ ABeam, ZBeam, nbi_kvolt, nbi_full, nbi_half

NAMELIST /preext/ precur,prerbz,prerrz,prezrz,preqpr,prepre, &
                  prenb2,preter,prener,preti2,prevp2,        &
                  extcur,extrbz,extrrz,extzrz,extqpr,extpre, &
                  extnb2,extter,extner,extti2,extvp2

!==================
! End declarations
!==================

!==================
! MPI initialisation
!==================

#ifdef __MPI
 nbi_comm = MPI_COMM_WORLD
 call MPI_INIT(ierr)
 if(ierr.ne.0) call errmsg_exit(' MPI_INIT failure! ')

 call MPI_COMM_RANK(nbi_comm,myid,ierr)
 call MPI_COMM_SIZE(nbi_comm,numprocs,ierr)

 call MPI_GET_PROCESSOR_NAME(proc_name,iresult,ierr)
 write(0,*) ' nubeam_driver_mpi: MPI, myid=',myid,' proc_name = "',trim(proc_name),'"'
#else
 nbi_comm = -1
 myid=0
#endif

call sget_env('ADASDIR',env_adasdir)
if (env_adasdir.eq.' ') then
  write(0,*) ' ?nubeam_driver: environment variable ADASDIR not found!'
  call bad_exit
endif
call sget_env('PREACTDIR',env_preactdir)
if (env_preactdir.eq.' ') then
  write(0,*) ' ?nubeam_driver: environment variable PREACTDIR not found!'
  call bad_exit
endif

!==============================  
! parse command line arguments
!==============================  

npar = command_argument_count()
workpath=' '
mrunid=' '
idev=' '
DO ipar = 1, npar,2
  call get_command_argument(ipar, arg)
  if (trim(arg) .eq. '-id') then
    call get_command_argument(ipar+1, argval)
    READ(argval,*) mrunid
  else if (trim(arg) .eq. '-idir') then
    call get_command_argument(ipar+1, inputdir)
  else if (trim(arg) .eq. '-wdir') then
    call get_command_argument(ipar+1, workpath)
  else if (trim(arg) .eq. '-tok') then
    call get_command_argument(ipar+1, idev)
  else
    WRITE(6,*) 'Arguments are:'
    WRITE(6,*) '-id    <RunId> '
    WRITE(6,*) '-tok   <Device> '
    WRITE(6,*) '-idir  <path to Ufiles> '
    WRITE(6,*) '-wdir  <path for Output> '
    WRITE(6,*) 'e.g.:'
    WRITE(6,*) 'nubeam_driver -id 80200C13 -tok ITER -idir $DATADIR/ITER -wdir $WORKDIR/ITER'
    WRITE(6,*) '   '
    WRITE(6,*) 'To print this info:'
    WRITE(6,*) 'nubeam_driver  [help -help -h]'
    WRITE(6,*) '   '
    stop
  endif
ENDDO

IF (trim(mrunid) .EQ. '') THEN
  WRITE(6,*) 'Missing runid'
  stop
ENDIF

WRITE(6,*) ' RunID: ' // mrunid

if (len_trim(mrunid) .eq. 8) then
   shot_str=mrunid(1:5)
   tail=mrunid(6:8)
else if (len_trim(mrunid) .eq. 9) then
  shot_str=mrunid(1:6)
  tail=mrunid(7:9)
endif

read(shot_str,*) nshot

IF (trim(idev) .EQ. '') THEN
  idev='AUGD'
endif
call sget_env('HOME',home_dir)
IF (trim(workpath) .EQ. '') THEN
  workpath = trim(home_dir) // '/nb_client/' // trim(idev) // '/' // &
             trim(shot_str) // '/' // trim(tail)
  out1d='/afs/ipp/home/n/nubeam/output/scalars.txt'
  out2d='/afs/ipp/home/n/nubeam/output/profiles.txt'
  CALL system('mkdir -p ' // trim(workpath))
else
  call sget_env('TRANSP_LOCATION',d_path)
  if (trim(d_path) .eq. '') then
    out1d='scalars.txt'
    out2d='profiles.txt'
  else
    out1d=trim(d_path)//'/scalars.txt'
    out2d=trim(d_path)//'/profiles.txt'
  endif
endif

IF (trim(inputdir) .EQ. '') THEN
  inputdir = trim(home_dir) // '/nb_client/' // trim(idev) // '/' // &
             trim(shot_str) // '/' // trim(tail)
ENDIF
namelist_file=trim(inputdir) // '/' // trim(mrunid) // '.nml'

!===============
WRITE(6,*) 'Reading namelist '// namelist_file
!===============

open(inf, file=trim(namelist_file), delim='apostrophe')
read(inf, nml=run,   iostat=ios)
read(inf, nml=plasma,iostat=ios)
read(inf, nml=nbi,   iostat=ios)
read(inf, nml=preext,iostat=ios)
close(inf)

ntime      = (ftime - tinit) / timestep + 1
DO irho=1,nrho_in
  rho_in(irho)=real((irho-1._rspec)/(nrho_in-1._rspec),KIND=rspec)
ENDDO

IF (debug .GT. 0) ps_debug = 3

IF (myid .EQ. 0) THEN

!=============
  WRITE(6,*) 'Reading UFiles'
!=============

  ALLOCATE(time(ntime),IPL_IN(ntime),PNBI_IN(ntime,n_beams), STAT=iErr)
  DO it=1,ntime
    time(it) = tinit + real(it-1,KIND=rspec) * timestep
  ENDDO

!=====
  WRITE(6,*) 'I_pl'
  call managerd_u(trim(inputdir),trim(mrunid),trim(precur),trim(extcur),&
                  iochan,nx1,nx2,nx3,nscal,iErr)
  allocate(d1d(nx1),time_u(nx1),STAT=iErr)
  CALL read_u1d(iochan,idev,time_u,nx1,d1d,nscal,iErr)
  CALL MQUIT(iErr,nbi_comm)
  CALL INTERP1D(nx1,time_u,d1d(1:nx1),ntime,time,IPL_IN)
  max_var=maxval(IPL_IN)
  IF ( ( max_var .LE. min_ipl) .OR. (max_var .GE. max_ipl) ) THEN
    write(6,*) 'Maximum value of IPL_IN is ',max_var
    write(6,'(2A,i5,2A)') 'Please check units in ', precur, nshot, '.', extcur
    STOP
  endif
  deallocate(d1d,time_u)
  
!=================
  WRITE(6,*) 'P_NBI'
  call managerd_u(trim(inputdir),trim(mrunid),trim(prenb2),trim(extnb2),&
                  iochan,nx1,nx2,nx3,nscal,iErr)
  allocate(d2d(nx1,nx2),time_u(nx1),beam_label(nx2), STAT=iErr)
  CALL read_u2d(iochan,idev,time_u,nx1,beam_label,nx2,d2d,nscal,iErr)
  n_beams_use=nx2
  IF (iErr .ne. 0) THEN
    WRITE(6,*) ' Error reading P_NBI UFile.'
    STOP
  ELSE
    CALL INTERP2D(nx2,nx1,time_u,d2d(1:nx1,1:nx2),ntime,time,PNBI_IN)
  ENDIF
  deallocate(d2d,time_u,beam_label)
  
!=================
  WRITE(6,*) 'R(t,rho_tor,th)'
  call managerd_u(trim(inputdir),trim(mrunid),trim(prerrz),trim(extrrz),&
                  iochan,nx1,nx2,nx3,nscal,iErr)
  allocate(d3d(nx1,nx2,nx3),time_u(nx1),rho_tor_u(nx2),th_u(nx3), STAT=iErr)
  CALL read_u3d(iochan,idev,time_u,nx1,rho_tor_u,nx2,th_u,nx3,d3d,nscal,iErr)
  if (iErr .ne. 0) then
    WRITE(6,*) ' Error reading R(rho_tor,theta) UFile.'
    CALL MQUIT(iErr,nbi_comm)
  endif
  nrho_eq = nx2
  nth_eq  = nx3

  ALLOCATE(rho_eq(nrho_eq),th_eq(nth_eq), STAT=iErr)
  ALLOCATE(RSURF_IN(ntime,nrho_eq,nth_eq),ZSURF_IN(ntime,nrho_eq,nth_eq), STAT=iErr)

  rho_eq = real(rho_tor_u(1:nrho_eq),KIND=rspec)
  th_eq  = real(th_u(1:nth_eq),KIND=rspec)
! TODO: remove this workaround for inaccurate theta start/end values
  th_eq(nth_eq) = 2.*3.1415926535897931e0_rspec
  
  if (iErr .ne. 0) then ! array could not be allocated
    WRITE(6,*) " Could not allocate array 'RSURF_IN(ntime,nrho_eq,nth_eq)'"
    CALL MQUIT(iErr,nbi_comm)
  endif
!  CALL INTERP3D(nx2,nx3,nx1,time_u,d3d(1:nx1,1:nx2,1:nx3),ntime,time,RSURF_IN)
  CALL APPROX3D(nx2,nx3,nx1,time_u,d3d(1:nx1,1:nx2,1:nx3),ntime,time,RSURF_IN)
  deallocate(d3d,time_u,rho_tor_u,th_u)
  
!=================
  WRITE(6,*) 'Z(t,rho_tor,th)'
  call managerd_u(trim(inputdir),trim(mrunid),trim(prezrz),trim(extzrz),&
                  iochan,nx1,nx2,nx3,nscal,iErr)
  allocate(d3d(nx1,nx2,nx3),time_u(nx1),rho_tor_u(nx2),th_u(nx3), STAT=iErr)
  CALL read_u3d(iochan,idev,time_u,nx1,rho_tor_u,nx2,th_u,nx3,d3d,nscal,iErr)
  if (iErr .ne. 0) then
    WRITE(6,*) ' Error reading Z(rho_tor,theta) UFile.'
    CALL MQUIT(iErr,nbi_comm)
  endif
  if (nx2 .ne. nrho_eq) then
    WRITE(6,*) ' Number of rho steps of Zsurf(ntime,nrho_eq,nth_eq) inconsistent.'
    stop
  endif
  if (nx3 .ne. nth_eq) then
    WRITE(6,*) ' Number of theta steps of Zsurf(ntime,nrho_eq,nth_eq) inconsistent.'
    stop
  endif
  
!  CALL INTERP3D(nx2,nx3,nx1,time_u,d3d(1:nx1,1:nx2,1:nx3),ntime,time,ZSURF_IN)
  CALL APPROX3D(nx2,nx3,nx1,time_u,d3d(1:nx1,1:nx2,1:nx3),ntime,time,ZSURF_IN)
  deallocate(d3d,time_u,rho_tor_u,th_u)
  
!=================
  WRITE(6,*) 'q,P,R*B_phi(t,rho_tor_eq)'

  ALLOCATE(Q_IN(ntime,nrho_eq),P_IN(ntime,nrho_eq),RBPHI_IN(ntime,nrho_eq),STAT=iErr)
  
  call read_prof(trim(inputdir),ierr,trim(preqpr),trim(extqpr),trim(mrunid),'Q',&
                 1.e0_rspec,min_q,max_q,ntime,time,nrho_eq,rho_eq,Q_IN,idev)
  call read_prof(trim(inputdir),ierr,trim(prepre),trim(extpre),trim(mrunid),'P',&
                 1.e0_rspec,min_p,max_p,ntime,time,nrho_eq,rho_eq,P_IN,idev)
  call read_prof(trim(inputdir),ierr,trim(prerbz),trim(extrbz),trim(mrunid),'R*Bphi',&
                 1.e0_rspec,min_rb,max_rb,ntime,time,nrho_eq,rho_eq,RBPHI_IN,idev)
  
!=================
  WRITE(6,*) 'ne,Te,Ti,v_tor,Zeff(t,rho_tor)'
  
  ALLOCATE(NE_IN(ntime,nrho_in),TE_IN(ntime,nrho_in),TI_IN(ntime,nrho_in),&
           VTOR_IN(ntime,nrho_in),ZEFF_IN(ntime,nrho_in),TIMP_IN(ntime,nrho_in),&
           NI_IN(ntime,nrho_in),NIMP_IN(ntime,nrho_in), STAT=iErr)
  
  call read_prof(trim(inputdir),ierr,trim(prener),trim(extner),trim(mrunid),'Ne',&
                 1.e6_rspec,min_ne,max_ne,ntime,time,nrho_in,rho_in,NE_IN,idev)
  call read_prof(trim(inputdir),ierr,trim(preter),trim(extter),trim(mrunid),'Te',&
                 1.e-3_rspec,min_te,max_te,ntime,time,nrho_in,rho_in,TE_IN,idev)
  call read_prof(trim(inputdir),ierr,trim(preti2),trim(extti2),trim(mrunid),'Ti',&
                 1.e-3_rspec,min_ti,max_ti,ntime,time,nrho_in,rho_in,TI_IN,idev)
  
  IF (trim(prezef) .NE. '') THEN
    call read_prof(trim(inputdir),ierr,trim(prezef),trim(extzef),trim(mrunid),'Zeff',&
                   1._rspec,min_zef,max_zef,ntime,time,nrho_in,rho_in,ZEFF_IN,idev)
  ELSE
    DO it=1,ntime
      DO irho=1,nrho_in
        ZEFF_IN(it,irho)=czeff
      ENDDO
    ENDDO
  ENDIF
  
  IF (trim(prevp2) .NE. '' ) THEN
    call read_prof(trim(inputdir),ierr,trim(prevp2),trim(extvp2),trim(mrunid),'Vt',&
              1.e-2_rspec,min_vt,max_vt,ntime,time,nrho_in,rho_in,VTOR_IN,idev)
  ELSE
    DO it=1,ntime
      DO irho=1,nrho_in
        VTOR_IN(it,irho)=0_rspec
      ENDDO
    ENDDO
  ENDIF
  
!========================================
! read PS parameters from namelist_file:
! machine description BEFORE shot config!
!========================================

  WRITE(6,*) '  PS_MDESCR_READ'
  CALL PS_MDESCR_READ (namelist_file, iErr)!, g_filename = ps_equilibrium)
  CALL MQUIT(iErr,nbi_comm)
  
  WRITE(6,*) 'Setting grid sizes and allocating ps arrays...'
  ps%nrho_eq  = nrho_eq
  ps%nth_eq  = nth_eq
  ps%nrho   = nrho_in
  ps%nrho_eq_geo = nrho_eq
  ps%t0 = 0.0
  ps%t1 = 0.0

  WRITE(6,*) '  PS_ALLOC_PLASMA_STATE'
  CALL PS_ALLOC_PLASMA_STATE(iErr, ps)
  CALL MQUIT(iErr,nbi_comm)
  
  WRITE(6,*) '  PS_SCONFIG_READ'
  CALL PS_SCONFIG_READ(namelist_file, iErr)
  CALL MQUIT(iErr,nbi_comm)
  
  WRITE(6,'(/A)') 'Grid sizes / timesteps:'
  WRITE(6,'(A10,I5)') ' ntime   ', ntime
  WRITE(6,'(A10,I5)') ' nspec_th', ps%nspec_th
  WRITE(6,'(A10,I5)') ' nrho    ', ps%nrho
  WRITE(6,'(A10,I5)') ' nrho_eq ', ps%nrho_eq
  WRITE(6,'(A10,I5)') ' nth_eq  ', ps%nth_eq
  WRITE(6,'(A10,I5)') ' nR_eq   ', ps%nR
  WRITE(6,'(A10,I5)') ' nZ_eq   ', ps%nZ
  WRITE(6,*) 

! Output files
  netcdf_out = trim(workpath) // '/' // trim(mrunid) // '.cdf'
  psfile_in  = trim(mrunid) // '_psfile_in.cdf'
  psfile_out = trim(mrunid) // '_psfile_out.cdf'
  
!==============================  
! set some fundamental variables into ps:
  ps%runid = mrunid
  ps%Global_Label = 'AUGD NUBEAM Driver plasma_state_test'
  ps%rho   = rho_in! rho for density/temperature profiles
  ps%rho_eq  = rho_eq! rho for equilibrium
  ps%th_eq  = th_eq  ! theta for equilibirum
  ps%rho_eq_geo = rho_eq ! use same grid as above
  
! set beam energy
  do ib = 1,ps%nbeam
    ps%kvolt_nbi(ib) = nbi_kvolt(ib)
    ps%frac_full(ib) = nbi_full(ib)
    ps%frac_half(ib) = nbi_half(ib)
  enddo
  
  zimp=ps%q_S(2)/ev
  
!=================
! Impurities
  DO it=1,ntime
    DO irho=1,nrho_in
      TIMP_IN(it,irho) = TI_IN(it,irho)
      if (zimp .NE. 1.) then
        NIMP_IN(it,irho) = NE_IN(it,irho)*(ZEFF_IN(it,irho)-1._rspec)/zimp/(zimp-1._rspec)
      else
        NIMP_IN(it,irho) = 0._rspec
      endif
      NI_IN(it,irho) = NE_IN(it,irho)-zimp*NIMP_IN(it,irho) !n_D
    ENDDO
  ENDDO

ENDIF ! myid eq 0

!===========
! Time loop
!===========

IF (debug .GE. 2) THEN
  IF (myid .LT. 10) THEN
    write(myid_str,'(A,i1)') '0',myid
  ELSE
    write(myid_str,'(i2)') myid
  ENDIF
  write(log_file,'(A)') trim(workpath)//'/nubeam_comp'//trim(myid_str)//'.log'
  write(6,'(2A)') 'Log file is ',trim(log_file)
  CALL openLog(file=trim(log_file),status='RENAME',iostat=ios,comm_world=nbi_comm)
ENDIF

WRITE(6,'(/,A,i5)') 'Starting time loop for processor ',myid

DO it = 1, ntime
  nbsfinish=0 ! new time step starts from fresh
!
  IF (myid .EQ. 0) THEN
! remove all <RUNID>_birth* files in WORK directory as they interfere with
! new 'birth' data storing
!
! NOTR: Keep in mind that <RUNID> _nubeam_birth.cdf<nouttim> files will be rewritten
!
     suffix_del=trim(mrunid)//'_birth'//'*'
     call fwc_delete(trim(workpath),trim(suffix_del),iErr) !delete fiels <runid>_birth_cpu*
!
    WRITE(6,'(/,A,f9.4,/)') 'Time ',time(it)

    WRITE(6,*) '  PS_UPDATE_RZSURF'
    CALL PS_UPDATE_RZSURF(Rsurf_IN(it,:,:), Zsurf_IN(it,:,:), 1, iErr)
    CALL MQUIT(iErr,nbi_comm)

    WRITE(6,*) 'Writing equilibrium data into ps... (q_eq, P_eq, g_eq)'
    ps%q_eq = Q_IN(it,:)
    ps%P_eq = P_IN(it,:)
    ps%g_eq = abs(RBphi_IN(it,:)) !MG

    WRITE(6,*) '  PS_MHDEQ_VERIFY'
    call PS_MHDEQ_VERIFY(iErr, update_psi = .TRUE., update_phit = .TRUE., pcur = IPL_IN(it))
    CALL MQUIT(iErr,nbi_comm)

    WRITE(6,*) '  PS_MHDEQ_DERIVE'
    call PS_MHDEQ_DERIVE("everything", ierr)
    CALL MQUIT(iErr,nbi_comm)
    do ib = 1,n_beams_use
      ps%power_nbi(ib) = PNBI_IN(it,ib)
    enddo
     do irho=2,nrho_in
      irhob=irho-1
      ps%ns(irhob,0) = (NE_IN(it,irho)+NE_IN(it,irho-1))/2.
      ps%Ts(irhob,0) = (TE_IN(it,irho)+TE_IN(it,irho-1))/2.
      ps%ns(irhob,1) = (NI_IN(it,irho)+NI_IN(it,irho-1))/2.
      ps%Ts(irhob,1) = (TI_IN(it,irho)+TI_IN(it,irho-1))/2.
      ps%ns(irhob,2) = (NIMP_IN(it,irho)+NIMP_IN(it,irho-1))/2.
      ps%Ts(irhob,2) = (TIMP_IN(it,irho)+TIMP_IN(it,irho-1))/2.
      ps%omegat(irhob) = (VTOR_IN(it,irho)+VTOR_IN(it,irho-1))/1.65/2.
    enddo

    IF (it .EQ. 1) THEN

!==============================
! save initial ps as netcdf file

      WRITE(6,*) 'Saving input ps to cdf... (PS_STORE_PLASMA_STATE)'
      CALL PS_STORE_PLASMA_STATE(iErr, trim(workpath) // '/' // psfile_in)
      if (debug .GT. 0) WRITE(6,*) ' error code ',iErr

!==============================
! initialise NUBEAM: 

! from nml     
      WRITE(6,*) '  NBI_CTRL_INIT_READ)'
      CALL NBI_CTRL_INIT_READ(namelist_file, nb_init, iErr)
      CALL MQUIT(iErr,nbi_comm)

      WRITE(6,*) '  NUBEAM_SET_WORKPATH'
      CALL NUBEAM_SET_WORKPATH(trim(workpath), iErr)  ! set workpath for NUBEAM files

! from ps     
      WRITE(6,*) '  NUBEAM_CTRL_INIT'
      CALL NUBEAM_CTRL_INIT(ps, nb_init, iErr)
      CALL MQUIT(iErr,nbi_comm)

      WRITE(6,*) '  NBI_CTRL_UPDATE_DFLT' 
      CALL NBI_CTRL_UPDATE_DFLT(nb_update)! initialise NUBEAM update state
      CALL MQUIT(iErr,nbi_comm)

    ENDIF ! if it .EQ. 1
  ENDIF   ! if myid EQ 0

!====================
! NUBEAM calculation
!====================

! Make sure all processor are in the same step and see the same
! input PS 
#ifdef __MPI
 CALL MPI_BARRIER(nbi_comm,ierr)
#endif

  WRITE(6,'(/,2(A,I3))') ' Starting NUBEAM control step for itime = ',it, &
          ' for processor',myid
  IF (myid .EQ. 0) THEN
    ps%t0 = ps%t1
    ps%t1 = time(it)
    status = 0
    i_trk_dep0 = inbi_logical(time(it).gt.outtim(j_fbm))
    i_trk_dep0 =i_trk_dep0*nb_init%nltrk_dep0
    if(i_trk_dep0.eq.1)then
      nbsfinish=nbsfinish+1
    endif	
    nb_update%nltrk_dep0 = i_trk_dep0
    call nubeam_ctrl_step(ps, ps, nb_update, status) ! NUBEAM step
  ELSE
    status=0
    call nubeam_ctrl_step_child(status) ! NUBEAM step for MPI
    CALL MQUIT(status,nbi_comm)
  ENDIF ! myid

#ifdef __MPI
    call MPI_BCAST(nbsfinish,1,MPI_INTEGER,0,nbi_comm,ierr) ! all CPUs have nbsfinish now
#endif
    !dealing with files related to beam ion birth source (if any action needed)
     call nbi_avg_bsource(nbsfinish,.true.) !keep files in workpath directory
     if(myid.eq.0.and.nbsfinish.ge.1) then
        call nubeam_avg_bsource(1,.true.) ! keep bs files in workpath directory, embed nouttim=j_fbm
     endif


!===============================================================
! set output grid and allocate temporary output data structures
  IF (myid .EQ. 0) THEN
    IF (it .EQ. 1) THEN
      nrho_nbi = ps%nrho_nbi
      nrho = nrho_nbi-1

      WRITE(6,*) "Setting output grid, sized ",nrho
      ALLOCATE(rho(nrho), STAT=iErr)
      ALLOCATE(rho_nbi(nrho_nbi), STAT=iErr)
      rho_nbi=ps%rho_nbi
      CALL MQUIT(iErr,nbi_comm)
      DO irho = 1, nrho
        rho(irho) = 0.5*(rho_nbi(irho) + rho_nbi(irho+1))
      ENDDO

      WRITE(6,*) "Allocating output data structures..."
      ALLOCATE(IB_OUT(ntime),UIB_OUT(ntime),ITOT_OUT(ntime), &
               EHEAT_OUT(ntime),IHEAT_OUT(ntime),HEATTOT_OUT(ntime), STAT=iErr)
      ALLOCATE(AREA_OUT(nrho,ntime),DAREA_OUT(nrho,ntime), &
               VOL_OUT(nrho,ntime),DVOL_OUT(nrho,ntime),STAT=iErr)
      ALLOCATE(TE_OUT(nrho,ntime),TI_OUT(nrho,ntime),TIMP_OUT(nrho,ntime), &
               NE_OUT(nrho,ntime),NI_OUT(nrho,ntime),NIMP_OUT(nrho,ntime), &
               VTOR_OUT(nrho,ntime),ZEFF_OUT(nrho,ntime), STAT=iErr)
      ALLOCATE(Q_OUT(nrho_eq,ntime),P_OUT(nrho_eq,ntime),RBPHI_OUT(nrho_eq,ntime), STAT=iErr)
      ALLOCATE(BDEP_OUT(nrho,ntime),PBE_OUT(nrho,ntime), &
               PBI_OUT(nrho,ntime),CURB_OUT(nrho,ntime), &
               UCURB_OUT(nrho,ntime), STAT=iErr)
      ALLOCATE(BDENS_OUT(nrho,ntime),UFASTPA_OUT(nrho,ntime), &
               UFASTPP_OUT(nrho,ntime), STAT=iErr)
      ALLOCATE(fs_data(nrho_nbi),fs_area(nrho_nbi),fs_vol(nrho_nbi))

    ENDIF! IF it .EQ. 1

    CALL nbo_get_cptimes(zcptimes)
    call nbo_get_deposition(zdeposition)
    call nbo_get_density(zdensity)
    call nbo_get_currents(zcurrents)
    call nbo_get_powers(zpowers)
    call nbi_get_profiles(zprofiles)
    
    call mindex_nbfi(ZBeam,ABeam,1, inubeam)
    if (debug .GT. 0) WRITE(6,*) ' inubeam = ',inubeam
 
!================================================================
! Calculate differential flux surface cross sections and volumes
!================================================================
    
    call eq_volume(nrho_nbi, rho_nbi, 0, fs_vol, ierr)
    if (iErr .ne. 0) then
      WRITE(6,*) '!!! eq_volume error: ', iErr
    endif

    call eq_area(nrho_nbi, rho_nbi, 0, fs_area, ierr)
    if (iErr .ne. 0) then
      WRITE(6,*) '!!! eq_area error: ', iErr
    endif
    
    do irho = 1,nrho
      VOL_OUT(irho,it)   = fs_vol(irho+1)
      AREA_OUT(irho,it)  = fs_area(irho+1)
      DVOL_OUT(irho,it)  = fs_vol(irho+1)  - fs_vol(irho)
      DAREA_OUT(irho,it) = fs_area(irho+1) - fs_area(irho)
    enddo


! beam particle deposition
    CALL eq_rgetf(ps%nrho_nbi, ps%rho_nbi, zdeposition%io_bdeps(inubeam), 0, fs_data, iErr)
    if (iErr .ne. 0)  WRITE(6,*) '!!! eq_rgetf error: ', iErr
    do irho=1,nrho
      BDEP_OUT(irho,it) = (fs_data(irho+1) - fs_data(irho))/DVOL_OUT(irho,it)
    enddo

! Fast ion density
    CALL eq_rgetf(ps%nrho_nbi, ps%rho_nbi, zdensity%io_bdenss(inubeam), 0, fs_data, iErr)
    if (iErr .ne. 0)  WRITE(6,*) '!!! eq_rgetf error: ', iErr
    do irho=1,nrho
      BDENS_OUT(irho,it) = (fs_data(irho+1) - fs_data(irho))/DVOL_OUT(irho,it)
    enddo

! Fast ion parallel energy density
    CALL eq_rgetf(nrho_nbi, rho_nbi, zdensity%io_udenspl(inubeam), 0, fs_data, iErr)
    if (iErr .ne. 0)  WRITE(6,*) '!!! eq_rgetf error: ', iErr
    do irho=1,nrho
      UFASTPA_OUT(irho,it) = (fs_data(irho+1) - fs_data(irho))/DVOL_OUT(irho,it)
    enddo

! Fast ion perpendicular energy density
    CALL eq_rgetf(nrho_nbi, rho_nbi, zdensity%io_udenspp(inubeam), 0, fs_data, iErr)
    if (iErr .ne. 0)  WRITE(6,*) '!!! eq_rgetf error: ', iErr
    do irho=1,nrho
      UFASTPP_OUT(irho,it) = (fs_data(irho+1) - fs_data(irho))/DVOL_OUT(irho,it)
    enddo

! electron heating
    CALL eq_rgetf(nrho_nbi, rho_nbi, zpowers%io_pbe, 0, fs_data, iErr)
    if (iErr .ne. 0)  WRITE(6,*) '!!! eq_rgetf error: ', iErr
    EHEAT_OUT(it) =  fs_data(nrho_nbi)
    do irho=1,nrho
      PBE_OUT(irho,it) = (fs_data(irho+1) - fs_data(irho)) / DVOL_OUT(irho,it)
    enddo
    
! ion heating
    CALL eq_rgetf(nrho_nbi, rho_nbi, zpowers%io_pbi, 0, fs_data, iErr)
    if (iErr .ne. 0)  WRITE(6,*) '!!! eq_rgetf error: ', iErr
    IHEAT_OUT(it) =  fs_data(nrho_nbi)
    do irho=1,nrho
      PBI_OUT(irho,it) = (fs_data(irho+1) - fs_data(irho)) / DVOL_OUT(irho,it)
    enddo
    
    HEATTOT_OUT(it) = EHEAT_OUT(it) + IHEAT_OUT(it)
    
! nbi driven current (shielded)
    CALL eq_rgetf(nrho_nbi, rho_nbi, zcurrents%io_curjs(inubeam), 0, fs_data, iErr)
    if (iErr .ne. 0)  WRITE(6,*) '!!! eq_rgetf error: ', iErr
    IB_OUT(it) =  fs_data(nrho_nbi)
    do irho=1,nrho
      CURB_OUT(irho,it) = (fs_data(irho+1) - fs_data(irho)) / DAREA_OUT(irho,it)
    enddo
    
! nbi driven current (unshielded)
    CALL eq_rgetf(nrho_nbi, rho_nbi, zcurrents%io_ucurjs(inubeam), 0, fs_data, iErr)
    if (iErr .ne. 0)  WRITE(6,*) '!!! eq_rgetf error: ', iErr
    UIB_OUT(it) =  fs_data(nrho_nbi)
    do irho=1,nrho
      UCURB_OUT(irho,it) = (fs_data(irho+1) - fs_data(irho)) / DAREA_OUT(irho,it)
    enddo

! Output directly from input
    ITOT_OUT(it) = IPL_IN(it)
!=============================


    if (j_fbm .LE. 9) then
      if (time(it) .GE. outtim(j_fbm) .AND. outtim(j_fbm) .GT. 0. ) then

        WRITE(6,*) 'Write FBM data, distribution function on MC grid'
        write(fbm_out,'(4A,i1,A)') trim(workpath), '/', trim(mrunid), '_fi_',j_fbm, '.cdf'
        ps%dist_fun(ps%snbi_to_alla(inubeam)) = trim(fbm_out)
        write(6,'(2(A,f9.4))') 'time=',time(it),'  t_FBM=',outtim(j_fbm)
        call nbdrive_load(ps,ierr)
        if (ierr .NE. 0) then
          write(6,*) 'Problems with nbdrive_load', j_fbm
          CALL MQUIT(iErr,nbi_comm)
        else
! FBM
          call nbdrive_fbm_write(trim(fbm_out),ierr)
          if (iErr .ne. 0) then
            write(6,*) 'Problems with nbdrive_fbm_write ',j_fbm
            CALL MQUIT(iErr,nbi_comm)
          else
            write(6,'(A,/,2X,A)') 'Writing FBM distribution function in ',trim(fbm_out)
            j_fbm=j_fbm+1
          endif
        endif
      endif ! if time=OUTTIM

    endif

  ENDIF ! myid

enddo
!===============
! End time loop
!===============

!===============
! NetCDF output 
!===============

IF (myid .EQ. 0) THEN

  TE_OUT   = transpose(TE_IN)
  TI_OUT   = transpose(TI_IN)
  TIMP_OUT = transpose(TIMP_IN)
  NE_OUT   = transpose(NE_IN)
  NI_OUT   = transpose(NI_IN)
  NIMP_OUT = transpose(NIMP_IN)
  VTOR_OUT = transpose(VTOR_IN)
  ZEFF_OUT = transpose(ZEFF_IN)
  P_OUT    = transpose(P_IN)
  Q_OUT    = transpose(Q_IN)
  RBPHI_OUT= transpose(RBPHI_IN)

  WRITE(6,*) '  PS_STORE_PLASMA_STATE'
  CALL PS_STORE_PLASMA_STATE(iErr, trim(workpath) // '/' // psfile_out)
  if (iErr .GT. 0) WRITE(6,*) ' error code ',iErr
  
  !===============
  ! NetCDF output 
  
  WRITE(6,*) "  NetCDF output file"
  call check( nf90_create(netcdf_out, nf90_clobber, ncid) )
  
  WRITE(6,*) "   Defining dimensions..."
  WRITE(6,*) "      time:    ",ntime
  WRITE(6,*) "      rho:     ", nrho
  WRITE(6,*) "      rho_eq:  ", nrho_eq
  WRITE(6,*) "      th_eq:   ", nth_eq
  WRITE(6,*) "      rho_nbi: ", nrho_nbi
  
  call check( nf90_def_dim(ncid, TIME_NAME,    ntime,     time_dimid)    )
  call check( nf90_def_dim(ncid, RHO_NAME,     nrho,  rho_dimid)     )
  call check( nf90_def_dim(ncid, RHO_EQ_NAME,  nrho_eq,   rho_eq_dimid)  )
  call check( nf90_def_dim(ncid, TH_EQ_NAME,   nth_eq,    th_eq_dimid)   )
  call check( nf90_def_dim(ncid, RHO_NBI_NAME, nrho_nbi,  rho_nbi_dimid) )
  
  WRITE(6,*) "   Defining coordinate variables..."
  call check( nf90_def_var(ncid,TIME_NAME,   NF90_REAL,time_dimid,   time_varid)    )
  call check( nf90_def_var(ncid,RHO_NAME,    NF90_REAL,rho_dimid,    rho_varid)     )
  call check( nf90_def_var(ncid,RHO_EQ_NAME, NF90_REAL,rho_eq_dimid, rho_eq_varid)  )
  call check( nf90_def_var(ncid,TH_EQ_NAME,  NF90_REAL,th_eq_dimid,  th_eq_varid)   )
  call check( nf90_def_var(ncid,RHO_NBI_NAME,NF90_REAL,rho_nbi_dimid,rho_nbi_varid) )
  
  WRITE(6,*) "   Assigning attributes to coordinate variables..."
  call check( nf90_put_att(ncid, time_varid,    UNITS, TIME_UNITS)    )
  call check( nf90_put_att(ncid, rho_varid,     UNITS, RHO_UNITS)     )
  call check( nf90_put_att(ncid, rho_eq_varid,  UNITS, RHO_EQ_UNITS)  )
  call check( nf90_put_att(ncid, th_eq_varid,   UNITS, TH_EQ_UNITS)   )
  call check( nf90_put_att(ncid, rho_nbi_varid, UNITS, RHO_NBI_UNITS) )
  
  ! Write signals to NetCDF
  ! Time traces
  
  ios=0
  dimids1 = (/ time_dimid /)
  open(u1d,file=out1d)
  read(u1d,'(/A)')
  DO WHILE (ios .EQ. 0)
    read(u1d,'(A10,X,A32,X,A16)',iostat=ios) var_name,var_desc,var_unit
    IF (ios .EQ. 0) THEN
      write(6,'( 3(A,"%") )') var_name,var_desc,var_unit
      call check( nf90_def_var(ncid,var_name,NF90_REAL,dimids1,varid) )
      call check( nf90_put_att(ncid,varid,UNITS,var_unit) )
      call check( nf90_put_att(ncid,varid,DESC, var_desc) )
      IF (var_name .EQ. 'EHEAT     ') eheat_varid=varid
      IF (var_name .EQ. 'HEATTOT   ') heat_varid =varid
      IF (var_name .EQ. 'IBEAM     ') ib_varid   =varid
      IF (var_name .EQ. 'IHEAT     ') iheat_varid=varid
      IF (var_name .EQ. 'ITOT      ') itot_varid =varid
      IF (var_name .EQ. 'UIBEAM    ') uib_varid  =varid
    ENDIF
  ENDDO
  close(u1d)
  
  ! Profiles
  
  ios=0
  open(u2d,file=out2d)
  read(u2d,'(/A)')
  DO WHILE (ios .EQ. 0)
    read(u2d,'(A10,X,A32,X,A16,X,A10)',iostat=ios) var_name,var_desc,var_unit,var_coord
    IF (ios .EQ. 0) THEN
      write(6,'( 4(A,"%") )') var_name,var_desc,var_unit,var_coord
      IF (trim(var_coord) .EQ. 'RHO') then
        dimids2=(/ rho_dimid ,  time_dimid /)
      ELSE IF (trim(var_coord) .EQ. 'RHO_NBI') THEN
        dimids2 = (/ rho_nbi_dimid,  time_dimid /)
      ELSE IF (trim(var_coord) .EQ. 'RHO_EQ') THEN
        dimids2 = (/ rho_eq_dimid, time_dimid /)
      ENDIF
      call check( nf90_def_var(ncid,var_name,NF90_REAL,dimids2,varid) )
      call check( nf90_put_att(ncid,varid,UNITS,var_unit) )
      call check( nf90_put_att(ncid,varid,DESC, var_desc) )
      IF (var_name .EQ. 'TE        ') te_varid   =varid
      IF (var_name .EQ. 'TI        ') ti_varid   =varid
      IF (var_name .EQ. 'TIMP      ') timp_varid =varid
      IF (var_name .EQ. 'NE        ') ne_varid   =varid
      IF (var_name .EQ. 'NI        ') ni_varid   =varid
      IF (var_name .EQ. 'NIMP      ') nimp_varid =varid
      IF (var_name .EQ. 'Q         ') q_varid    =varid
      IF (var_name .EQ. 'P         ') p_varid    =varid
      IF (var_name .EQ. 'RBPHI     ') rbphi_varid=varid
      IF (var_name .EQ. 'BDEP      ') bdep_varid =varid
      IF (var_name .EQ. 'BDENS     ') bdens_varid =varid
      IF (var_name .EQ. 'UFASTPA   ') ufastpa_varid =varid
      IF (var_name .EQ. 'UFASTPP   ') ufastpp_varid =varid
      IF (var_name .EQ. 'PBE       ') pbe_varid  =varid
      IF (var_name .EQ. 'PBI       ') pbi_varid  =varid
      IF (var_name .EQ. 'CURB      ') curb_varid =varid
      IF (var_name .EQ. 'UCURB     ') ucurb_varid=varid
      IF (var_name .EQ. 'AREA      ') area_varid =varid
      IF (var_name .EQ. 'DAREA     ') darea_varid=varid
      IF (var_name .EQ. 'VOL       ') vol_varid  =varid
      IF (var_name .EQ. 'DVOL      ') dvol_varid =varid
    ENDIF
  ENDDO
  
  dimids3 = (/ time_dimid, rho_eq_dimid, th_eq_dimid /)
  call check( nf90_def_var(ncid, RSURF_NAME, NF90_REAL, dimids3, rsurf_varid) )
  call check( nf90_def_var(ncid, ZSURF_NAME, NF90_REAL, dimids3, zsurf_varid) )
  
  call check( nf90_put_att(ncid,rsurf_varid,UNITS,rsurf_units) )
  call check( nf90_put_att(ncid,zsurf_varid,UNITS,zsurf_units) )
  call check( nf90_put_att(ncid,rsurf_varid,DESC, rsurf_desc)  )
  call check( nf90_put_att(ncid,zsurf_varid,DESC, zsurf_desc)  )
  
  ! End define mode.
  call check( nf90_enddef(ncid) )
  
  ! Writing NetCDF data
  write(6,*) "   Writing grid data..."
  call check( nf90_put_var(ncid, time_varid,   time)    )
  call check( nf90_put_var(ncid, rho_varid,    rho)     )
  call check( nf90_put_var(ncid, rho_eq_varid, rho_eq)  )
  call check( nf90_put_var(ncid, th_eq_varid,  th_eq)   )
  call check( nf90_put_var(ncid, rho_nbi_varid,rho_nbi) )
  write(6,*) "   Writing time traces..."
  call check( nf90_put_var(ncid,eheat_varid,EHEAT_OUT)  )
  call check( nf90_put_var(ncid,heat_varid, HEATTOT_OUT))
  call check( nf90_put_var(ncid,ib_varid,   IB_OUT)     )
  call check( nf90_put_var(ncid,iheat_varid,IHEAT_OUT)  )
  call check( nf90_put_var(ncid,itot_varid, ITOT_OUT)   )
  call check( nf90_put_var(ncid,uib_varid,  UIB_OUT)    )
  write(6,*) "   Writing profiles..."
  call check( nf90_put_var(ncid,te_varid,   TE_OUT)    )
  call check( nf90_put_var(ncid,ti_varid,   TI_OUT)    )
  call check( nf90_put_var(ncid,timp_varid, TIMP_OUT)  )
  call check( nf90_put_var(ncid,ne_varid,   NE_OUT)    )
  call check( nf90_put_var(ncid,ni_varid,   NI_OUT)    )
  call check( nf90_put_var(ncid,nimp_varid, NIMP_OUT)  )
  call check( nf90_put_var(ncid,q_varid,    Q_OUT)     )
  call check( nf90_put_var(ncid,p_varid,    P_OUT)     )
  call check( nf90_put_var(ncid,rbphi_varid,RBPHI_OUT) )
  call check( nf90_put_var(ncid,bdep_varid, BDEP_OUT) )
  call check( nf90_put_var(ncid,bdens_varid, BDENS_OUT) )
  call check( nf90_put_var(ncid,ufastpa_varid, UFASTPA_OUT) )
  call check( nf90_put_var(ncid,ufastpp_varid, UFASTPP_OUT) )
  call check( nf90_put_var(ncid,pbe_varid,  PBE_OUT)  )
  call check( nf90_put_var(ncid,pbi_varid,  PBI_OUT)  )
  call check( nf90_put_var(ncid,curb_varid, CURB_OUT) )
  call check( nf90_put_var(ncid,ucurb_varid,UCURB_OUT))
  call check( nf90_put_var(ncid,area_varid, AREA_OUT) )
  call check( nf90_put_var(ncid,darea_varid,DAREA_OUT))
  call check( nf90_put_var(ncid,vol_varid,  VOL_OUT)  )
  call check( nf90_put_var(ncid,dvol_varid, DVOL_OUT) )
  write(6,*) "   R,z(t,rho,theta)"
  call check( nf90_put_var(ncid, rsurf_varid,  RSURF_IN))
  call check( nf90_put_var(ncid, zsurf_varid,  ZSURF_IN))
  
  ! Close the NetCDF file.
  call check( nf90_close(ncid) )

ENDIF !myid

#ifdef __MPI
 call MPI_FINALIZE(ierr)
#endif

!================

CONTAINS

!================
SUBROUTINE CHECK(status)

implicit none

integer, intent(in) :: status

if(status /= nf90_noerr) then
  print *, trim(nf90_strerror(status))
  CALL MQUIT(2,nbi_comm)
endif

END SUBROUTINE CHECK
!================
SUBROUTINE MQUIT(iErr,nbi_comm)

implicit none

integer, intent(in) :: iErr,nbi_comm         ! exit code

logical, parameter :: QOE = .TRUE.  ! turn quit on error on/off

if (QOE) then
  IF (iErr .NE. 0) THEN
    WRITE(6,*) 'Quitting in MQUIT due to error ',iErr
#ifdef __MPI
      CALL MPI_ABORT(nbi_comm,status)
#else
      CALL exit(iErr)
#endif
  ENDIF
endif

END SUBROUTINE MQUIT
!==================

END PROGRAM
