module addfld_mod
!
! Save fields to output history file.
!
  implicit none
  save
  contains
!-----------------------------------------------------------------------
  subroutine addfld(name,long_name,units,f,dname1,lb1,ub1,dname2,lb2,ub2,idx)
!
  use shr_kind_mod ,only: r8 => shr_kind_r8 ! 8-byte reals
  use params       ,only: max_fields,finit,iulog
  use geogrid      ,only: nlev,nilev,nlon,nlat
  use maggrid      ,only: nmlonp1,nmlev,nmlat
  use fields       ,only: fields_out,nfields,fillvalue
  use namelist     ,only: model_name
  implicit none
!
! Args:
  character(len=*),intent(in) :: dname1,dname2,name,long_name,units
  integer,intent(in) :: lb1,ub1,lb2,ub2,idx
  real(r8),intent(in) :: f(lb1:ub1,lb2:ub2)
!
! Local:
  integer :: ier,i,j,k,ix
  integer :: idim1,idim2,idim3,ihdim1,ihdim2,ihdim3
  logical :: mag,geo,firstcall
  character(len=8) :: dname3,model
!
! Valid dimension names on history (dname1 and dname2 must be one of these)
  integer,parameter :: nvalid=8
  character(len=8) :: valid_dnames(nvalid) = &
    (/'lev     ','ilev    ','lon     ','lat     ',  & ! geo dim names
      'mlev    ','imlev   ','mlon    ','mlat    '/)   ! mag dim names
!
! Dimensions of global fields_out(ix)%data corresponding to dim names:
  integer,save :: valid_dims(nvalid)
!
  model = trim(model_name)
  firstcall = .false.
  ix = strloc(fields_out%name,max_fields,name)
  if (ix==0) firstcall = .true. ! first call for this field
  if (firstcall) then
    nfields = nfields+1
    ix = nfields
    if (model=='TIMEGCM') then
      valid_dims = & ! repeat nmlev because nimlev is no longer used
      (/nlev      ,nlev      ,nlon+4    ,nlat         ,& ! geo dim sizes
        nmlev     ,nmlev     ,nmlonp1   ,nmlat     /)    ! mag dim sizes
    else
      valid_dims = & ! repeat nmlev because nimlev is no longer used
      (/nlev      ,nlev      ,nlon      ,nlat         ,& ! geo dim sizes
        nmlev     ,nmlev     ,nmlonp1   ,nmlat     /)    ! mag dim sizes
    endif
!   write(iulog,"('Addfld first call: model=',a,' nlev=',i4,' nlon=',i4,' nlat=',i4,' valid_dims(3)=',i4)") &
!     model,nlev,nlon,nlat,valid_dims(3)
  endif

! write(iulog,"('Enter addfld: field=',a,' firstcall=',l1,' nfields=',i3,' ix=',i3)") &
!   trim(name),firstcall,nfields,ix
! write(iulog,"('  dname1=',a,' lb1,ub1=',2i4,' dname2=',a,' lb2,ub2=',2i4,' idx=',i4)") &
!   dname1,lb1,ub1,dname2,lb2,ub2,idx
! write(iulog,"('  data min,max=',2e12.4)") minval(f(lb1:ub1,lb2:ub2)), &
!                                           maxval(f(lb1:ub1,lb2:ub2))
!
! Set idim1,idim2 (also validate dname1,2):
!
  idim1 = 0 ; idim2 = 0
  i = strloc(valid_dnames,nvalid,dname1)
  if (i > 0) idim1 = valid_dims(i)

  i = strloc(valid_dnames,nvalid,dname2)
  if (i > 0) idim2 = valid_dims(i)

  if (idim1==0) then
    write(iulog,"(/,'>>> addfld: invalid dname1 = ',a,' for field ',a)") &
      dname1,name
    write(iulog,"('Valid dim names: ',8a8)") valid_dnames
    call shutdown('addfld')
  endif
  if (idim2==0) then
    write(iulog,"(/,'>>> addfld: invalid dname2 = ',a,' for field ',a)") &
      dname2,name
    write(iulog,"('Valid dim names: ',8a8)") valid_dnames
    call shutdown('addfld')
  endif
!
! If idx==0, it must refer to lev dimension:
!
  if (idx==0.and. &
    ((index(dname1,'lev') > 0.and.index(dname2,'lon') > 0).or. &
     (index(dname1,'lev') > 0.and.index(dname2,'lat') > 0).or. &
     (index(dname1,'lon') > 0.and.index(dname2,'lev') > 0).or. &
     (index(dname1,'lat') > 0.and.index(dname2,'lev') > 0))) then
    write(iulog,"(/,5a)") '>>> addfld: idx can be zero only when saving',&
      ' f(lon,lat) or f(lat,lon): dname1,2=',dname1,' ',dname2
    call shutdown('addfld')
  endif
!
! If longitude dimension is global, assume task0_only, i.e., assume
! field is presented by master task only -- in this case gather2root
! must not be called on this field (sub mp_gather2root in my_mpi.F90)
!
      fields_out(ix)%task0_only = .false.
      if ((index(dname1,'lon') > 0 .and. lb1==1.and.ub1==idim1) .or.&
          (index(dname2,'lon') > 0 .and. lb2==1.and.ub2==idim2))    &
        fields_out(ix)%task0_only = .true.
!
! Mag or geo:
!
  geo = .true. ; mag = .false.
  i = strloc(valid_dnames(5:8),4,dname1)
  if (i > 0) then
    geo = .false.
    mag = .true.
  endif
!     
! Set full-domain secondary output array dimensions from dname1,2:
! (geographic or magnetic, lev,lat,or lon)
!        
  fields_out(ix)%dimnames = ' '
  fields_out(ix)%dimsizes = 0

  if (index(dname1,'lon') > 0) fields_out(ix)%dimsizes(1) = idim1
  if (index(dname1,'lat') > 0) fields_out(ix)%dimsizes(2) = idim1
  if (index(dname1,'lev') > 0) fields_out(ix)%dimsizes(3) = idim1
  if (index(dname2,'lon') > 0) fields_out(ix)%dimsizes(1) = idim2 
  if (index(dname2,'lat') > 0) fields_out(ix)%dimsizes(2) = idim2 
  if (index(dname2,'lev') > 0) fields_out(ix)%dimsizes(3) = idim2
  
  if (index(dname1,'lon') > 0) fields_out(ix)%dimnames(1) = dname1
  if (index(dname1,'lat') > 0) fields_out(ix)%dimnames(2) = dname1
  if (index(dname1,'lev') > 0) fields_out(ix)%dimnames(3) = dname1
  if (index(dname2,'lon') > 0) fields_out(ix)%dimnames(1) = dname2
  if (index(dname2,'lat') > 0) fields_out(ix)%dimnames(2) = dname2
  if (index(dname2,'lev') > 0) fields_out(ix)%dimnames(3) = dname2
!     
! Set third dimension of full-domain output array according to idx:
!    
  idim3 = 0                                ! 2d array
  fields_out(ix)%ndims = 2
  if (idx > 0) then                        ! 3d array
    fields_out(ix)%ndims = 3
    if (index(dname1,'lev')==0.and.index(dname2,'lev')==0) then
      idim3 = nlev ; if (mag) idim3 = nmlev
    elseif (index(dname1,'lat')==0.and.index(dname2,'lat')==0) then
      idim3 = nlat ; if (mag) idim3 = nmlat
    elseif (index(dname1,'lon')==0.and.index(dname2,'lon')==0) then
      idim3 = nlon
      if (model=='TIMEGCM') idim3 = nlon+4
      if (mag) idim3 = nmlonp1
    endif
    if (idim3==0) then
      write(iulog,"('>>> addfld: cannot determine idim3: dname1,2=',a,',',a)") &
        trim(dname1),trim(dname2)
      call shutdown('addfld')
    endif
    dname3 = '        '
    do i=1,nvalid
      if (idim3==valid_dims(i)) dname3 = valid_dnames(i)
    enddo
    do i=1,3
      if (len_trim(fields_out(ix)%dimnames(i))==0) then
        fields_out(ix)%dimsizes(i) = idim3
        fields_out(ix)%dimnames(i) = dname3
      endif
    enddo 
    if (idx > idim3) then
      write(iulog,"(/,3a,i4,3a,i4,a)") '>>> addfld: field ',trim(name),&
        ': bad idx=',idx,' dname3=',dname3,' idim3=',idim3,&
        ' (idx should not be > idim3)'
      call shutdown('addfld')
    endif
  endif

  fields_out(ix)%name = trim(name)

! write(iulog,"('addfld: name=',a,' ix=',i3,' idx=',i4,' ndims=',i3,' dimnames=',3a,' dimsizes=',3i3)") &
!   trim(fields_out(ix)%name),ix,idx,fields_out(ix)%ndims,fields_out(ix)%dimnames,fields_out(ix)%dimsizes

!
! Allocate fields_out(ix)%data, either 2d (lon,lat) or 3d (lon,lat,lev):
! (field is mag or geo)
! 
! 3d (i,j,k): 
  if (.not.associated(fields_out(ix)%data)) then
    ihdim1=fields_out(ix)%dimsizes(1)
    ihdim2=fields_out(ix)%dimsizes(2)
    ihdim3=fields_out(ix)%dimsizes(3)
    if (idx > 0) then
      allocate(fields_out(ix)%data(ihdim1,ihdim2,ihdim3),stat=ier)
      if (ier /= 0) then
        write(iulog,"(/,'>>> Error allocating 3d (ihdim1,2,3=',3i3,') for field ',a)") &
          ihdim1,ihdim2,ihdim3,trim(fields_out(ix)%name)
        call shutdown('addfld')
      else 
        write(iulog,"(/,'Allocated 3d sech field ',a,'(',a,'=',i3, ',',a,'=',i3,',',a,'=',i3,')')") &
          trim(fields_out(ix)%name),&
          trim(fields_out(ix)%dimnames(1)),ihdim1,&
          trim(fields_out(ix)%dimnames(2)),ihdim2,&
          trim(fields_out(ix)%dimnames(3)),ihdim3
      endif
!
! 2d (i,j):
    else ! idx==0 (no lev dimension)
      allocate(fields_out(ix)%data(ihdim1,ihdim2,1),stat=ier)
      if (ier /= 0) then
        write(iulog,"(/,'>>> Error allocating 2d (ihdim1,2=',2i3,' for field ',a)") &
          ihdim1,ihdim2,trim(fields_out(ix)%name)
        call shutdown('addfld')
      else
        write(iulog,"(/,'Allocated 2d sech field ',a,'(',a,'=',i3,',',a,'=',i3,')')") &
          trim(fields_out(ix)%name),&
          trim(fields_out(ix)%dimnames(1)),ihdim1,&
          trim(fields_out(ix)%dimnames(2)),ihdim2
      endif
    endif
    fields_out(ix)%data = finit ! init whole array to init value
  endif ! this field data not associated
!
! Assign data to fsech(ix)%data(i,j,k) from f, according to dim names:
!
! f(lev,lon), idx=j (this is the most common case)
  if (index(dname1,'lev') > 0.and.index(dname2,'lon') > 0) then
    do k=lb1,ub1
      do i=lb2,ub2
        fields_out(ix)%data(i,idx,k) = f(k,i)
      enddo
    enddo
!   write(iulog,"('addfld: ix=',i3,' idx=',i3,' lb1,ub1=',2i4,' lb2,ub2=',2i4,' data min,max=',2e12.4)") &
!     ix,idx,lb1,ub1,lb2,ub2,minval(fields_out(ix)%data(lb2:ub2,idx,lb1:ub1)),&
!                            maxval(fields_out(ix)%data(lb2:ub2,idx,lb1:ub1))
!
! f(lev,lat), idx=i
  elseif (index(dname1,'lev') > 0.and.index(dname2,'lat') > 0) then
    do k=lb1,ub1
      do j=lb2,ub2
        fields_out(ix)%data(idx,j,k) = f(k,j)
      enddo
    enddo
!
! f(lon,lev), idx=j
  elseif (index(dname1,'lon') > 0.and.index(dname2,'lev') > 0) then
    do k=lb2,ub2
      do i=lb1,ub1
        fields_out(ix)%data(i,idx,k) = f(i,k)
      enddo
    enddo
!
! f(lon,lat), idx = k or 0
  elseif (index(dname1,'lon') > 0.and.index(dname2,'lat') > 0) then
    if (idx==0) then
      do j=lb2,ub2
        do i=lb1,ub1
          fields_out(ix)%data(i,j,1) = f(i,j)
        enddo
      enddo
    else  ! 3d
      do j=lb2,ub2
        do i=lb1,ub1
          fields_out(ix)%data(i,j,idx) = f(i,j)
        enddo
      enddo
    endif
!
! f(lat,lev), idx = i
  elseif (index(dname1,'lat') > 0.and.index(dname2,'lev') > 0) then
    do i=lb2,ub2
      do j=lb1,ub1
        fields_out(ix)%data(idx,j,k) = f(j,k)
      enddo
    enddo
!
! f(lat,lon), idx = k or 0
  elseif (index(dname1,'lat') > 0.and.index(dname2,'lon') > 0) then
    if (idx==0) then
      do j=lb1,ub1
        do i=lb2,ub2
          fields_out(ix)%data(i,j,1) = f(j,i)
        enddo
      enddo
    else
      do j=lb1,ub1
        do i=lb2,ub2
          fields_out(ix)%data(i,j,idx) = f(j,i)
        enddo
      enddo
    endif
  endif
!
! Set other structure components:
  fields_out(ix)%long_name  = long_name
  fields_out(ix)%units = units
  fields_out(ix)%geo = geo
  fields_out(ix)%mag = mag
!
! Report to stdout:
  if (firstcall) then
    write(iulog,"(/,'Initialized history field ',a,' (ix=',i3,'):')") &
      trim(fields_out(ix)%name),ix
    write(iulog,"('  name       = ',a)") trim(fields_out(ix)%name)
    write(iulog,"('  long_name  = ',a)") trim(fields_out(ix)%long_name)
    write(iulog,"('  units      = ',a)") trim(fields_out(ix)%units)
    write(iulog,"('  geo        = ',l1)")  fields_out(ix)%geo
    write(iulog,"('  mag        = ',l1)")  fields_out(ix)%mag
    write(iulog,"('  dimnames   = ',3a)")  fields_out(ix)%dimnames
    write(iulog,"('  dimsizes   = ',3i4)") fields_out(ix)%dimsizes
    write(iulog,"('  ndims      = ', i2)") fields_out(ix)%ndims
    write(iulog,"('  task0_only = ', l5)") fields_out(ix)%task0_only
  endif ! first call for this field

  end subroutine addfld
!-----------------------------------------------------------------------
  integer function strloc(strarray,nstr,substr)
  implicit none
!
! Search for string str in string array strarray(nstr).
! Its important NOT to use comparisons such as "if (trim(str0)==trim(str1))"
!   to avoid memory leaks by Intel ifort compiler (this is not a problem
!   w/ PGI and xlf compilers).
!
! Args:
    integer,intent(in) :: nstr
    character(len=*),intent(in) :: strarray(nstr)
    character(len=*) :: substr
!
! Local:
    integer :: i
    character(len=len(substr)) :: subtrim
    character(len=len(strarray(1))) :: strtrim
!
    strloc = 0
    aloop: do i=1,nstr
      if (len_trim(strarray(i)) > 0) then
        subtrim = trim(substr)
        strtrim = trim(strarray(i))
        if (subtrim == strtrim) then
          strloc = i
          exit aloop
        endif
      endif
    enddo aloop
  end function strloc
!-----------------------------------------------------------------------
end module addfld_mod