      program stps
      implicit none

cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
c NOTE :
c Daehyun Kim modified Dr. Matthew C. Wheeler's code for the MJO
c Simulation Metrics calculation system
c Note that any bugs or errors might be caused by modifications.  
c Contact : Daehyun Kim (kim@climate.snu.ac.kr)
cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc
cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc

c This program calculates 2-dim (or wavenumber-frequency) power spectrum
c for a "test" dataset consisting of specified waves and other things.
c
c--------------------------------------------------------------------
      integer lrecIn1, lrecIn2
      integer ic,i,t,n,pt,pn,NT,NL
      integer nn, tt, ny, iy
      parameter (NT=364,NL=144)
      parameter (ny=27)

      integer year1, leap, linux, dmax
      parameter (year1 = 1979, leap = 1)
      parameter (linux = 4)   
      parameter (dmax = 365)
      integer yy, dd, nod, it

      real WSAVE1(4*NL+15),WSAVE2(4*NT+15)
      real EEr(NL,NT),PEE(NL+1,NT+1),globmean,totvar
c     real OEE(NT+1,NL/2+1)
      real OEE(39,11)
      real ave_EEr(NL)
      real zm_EEr(NT)
      real Rd,eif,s,xx,xxm,ll,Pi,g,ps,k,ss(NL+1),ff(NT+1)
      real rand,s2,ps2,k2,eif2

      character*7 FOUT
      data FOUT /'u850_n1'/

c EEr(n,t) = the initial (real) data set that is a function of space and time.
c PEE(_,_) = the power spectrum that is a function of wavenumber and freq.
c ss() and ff() are arrays of wavenumbers and frequencies (cycles per day)
c          corresponding to PEE(,)
c NL = Number of longitudes in 'test' dataset
c NT = Number of times in 'test' dataset

c**** Quantities for determination of "test" dataset (consisting of waves).
c**** dimensional quantities are in MKS units.
c s = planetary zonal wavenumber
c k = zonal wavenumber = 2 Pi s / L  (m-1)  
c ps = dimensional phase speed w.r.t. the earth 
C ll = distance around latitude circle (m)
C xx = longitude in degrees
c xxm = longitudinal distance from Greenwich in metres.
c id = time in days  (integer)
c tt = time in seconds (real)
c eif = dimensional frequency of test wave
c--------------------------------------------------------------------
c------------------------MAIN PROGRAM--------------------------------

      Pi = 3.1415926
      g = 9.81
      ll = 2.*Pi*6.37E06           ! at the equator
      globmean = 0.
      totvar = 0.

      inquire(IOLENGTH=lrecIn1) ave_EEr
      inquire(IOLENGTH=lrecIn2) OEE

      open (1, file=
     &'/home/enver/work/programs/MJOWG/msd/level_1/u850_n1/data/'//
     &'/daily.10S10N.19790101_20051231.gdat', form='unformatted',
c     &access='direct',recl=nl*linux,status='old')
     &access='direct',recl=lrecIn1,status='old')
      open (2, file='u850_n1', form='unformatted',
c    &access='direct',recl=(NT+1)*(NL/2+1),status='unknown')
c    &access='direct',recl=39*11*linux,status='unknown')
     &access='direct',recl=lrecIn2,status='unknown')

      OEE = 0.

      yy = year1 - 1
      dd = 0

      do iy = 1, ny
      yy = yy + 1

      globmean = 0.
      ave_EEr = 0.
      zm_EEr = 0.
      totvar = 0.

      if (leap.eq.1) then
       nod = 365
       if (mod(yy,4).eq.0.and.mod(yy,100).ne.0) nod = 366
       if (mod(yy,400).eq.0) nod = 366
      elseif (leap.eq.0) then
       nod = dmax
      endif

      do t=1,NT

       if (leap.eq.1) then
        it = dd + t
       elseif (leap.eq.0) then
        it = nod*(iy-1)+t
       endif

      read (1, rec=it) (EEr(n,t),n=1,NL)

      do n = 1, NL
      globmean = globmean + EEr(n,t)/float(NT*NL)
      ave_EEr(n) = ave_EEr(n) + EEr(n,t)/float(NT)
      zm_EEr(t) = zm_EEr(t) + EEr(n,t)/float(NL)
      enddo

      enddo ! t

      do t=1,NT
      do n=1,NL
c      EEr(n,t) = EEr(n,t) - ave_EEr(n) - zm_EEr(t)
c      EEr(n,t) = EEr(n,t) - globmean
       EEr(n,t) = EEr(n,t) - ave_EEr(n)
      enddo
      enddo

      do t=1,NT
      do n=1,NL
c      totvar = totvar + ((EEr(n,t)-globmean)**2)/float(NT*NL)
       totvar = totvar + ((EEr(n,t))**2)/float(NT*NL)
      enddo
      enddo
      print*,'total variance of EEr is',totvar,' globmean is',globmean

c--------------------------------------------------------------------
c------------------COMPUTING SPACE-TIME SPECTRUM---------------------

      call cffti(NL,WSAVE1)      ! Initialize FFT for longitude direction
      call cffti(NT,WSAVE2)      ! Initialize FFT for time direction

      call powerspec2dim(EEr,NL,NT,WSAVE1,WSAVE2,PEE)

c The newly created PEE(NL+1,NT+1) contains the power spectrum.

c The corresponding frequencies, ff, and wavenumbers, ss, are:-
      do t=1,NT+1
       ff(t) = float(t-1-NT/2)/float(NT)    ! in cycles per day.
      enddo
      do n=1,NL+1
       ss(n) = float(n-1-NL/2)          ! in planetary wavenumber
      enddo

!     do n = NL/2+1, NL+1
!      nn = n-NL/2
!     do t = 1, NT+1
!      OEE(t,nn) = PEE(n,t)
!     enddo
!     enddo

      do n = NL/2+1, NL/2+1+10
       nn = n-NL/2
      do t = NT/2+1-19, NT/2+1+19
       tt = -(NT/2+1)+20+t
       OEE(tt,nn) = OEE(tt,nn)+PEE(n,t)/float(ny)
      enddo ! t
      enddo ! n

      if (leap.eq.1) then
       dd = dd + nod
      endif

      enddo ! iy

      write(2, rec=1) OEE

c--------------------------------------------------------------------
c------------             GRADS OUTPUT       ------------------------
c--------------------------------------------------------------------
      open(12,file='u850_n1.ctl',status='unknown')
      write  (12,1234) FOUT,39,(ff(i),i=nt/2+1-19,nt/2+1+19)
     *       ,11,(ss(i),i=nl/2+1,nl/2+1+10)
1234  format ('DSET ^',a7/
     &        '*'/
     &        'UNDEF  -999.'/
     &        'TITLE SPCTIME OUTPUT'/
     &        '*'/
     &        'XDEF ',i4,1x,'LEVELS ',39(1x,f7.3) /
     &        '*'/
     &        'YDEF ',i4,1x,'LEVELS ',11(1x,f7.3) /
     &        '*'/
     &        'ZDEF  1 LEVELS   1000.      '/
     &        '*'/
     &        'TDEF  1 LINEAR 1jan1979 1dy'/
     &        '*'/
     &        'VARS 1'/
     &        'power      0    99   power'/
     &        'ENDVARS')

      END

c-----------------------------------------------------------------------
      subroutine powerspec2dim(EEr,NL,NT,WSAVE1,WSAVE2,PEE) 

c Use NCAR's FFTpack routines to compute the two-dimensional power
c spectrum of the input dataset EEr.
c
c The input array EEr(n,t) is a function of two dimensions, e.g.,
c space (n=1 to NL) and time (t=1 to NT).
c
c The output array PEE(pn,pt) is also a function of two dimensions, e.g.,
c wavenumber (pn) and frequency (pt).
c pn varies from 1 to NL+1, and pt from 1 to NT+1.
c And the array PEE will be symmetric about pn=NL/2+1 and pt=NT/2+1.
c Hence only half of PEE needs to be plotted to see all the information.
c
c If the input array EE has as its two dimensions longitude and time,
c then the following applies to the arrangement of the output spectrum.
c   - Power for negative wavenumbers will be for pn=1 to NL/2
c   - Power for the zonal mean (wavenumber 0) will be at pn=NL/2+1
c   - Power for the positive wavenumbers will be for pn=NL/2+2 to NL+1
c   - Power for the Nyquist wavenumber will be at pn=1 and pn=NL+1
c   - Power for negative frequencies will be for pt=1 to NT/2+1
c   - Power for the time mean (frequency=0) will be at pt=NT/2+1
c   - Power for the positive frequencies will be for pt=NT/2+2 to NT+1
c   - Power for the Nyquist frequency will be at pt=1 and pt=NT+1
c   - Westward moving waves will be located in the quadrants defined by
c     +ve frequency and -ve wavenumber, or -ve frequency and +ve wavenumber.
c   - Eastward moving waves are for +ve/+ve, or -ve/-ve.
c For example, the wavenumber 1 eastward moving wave with the frequency of
c the first harmonic will be located at (NL/2+2,NT/2+2) or (NL/2,NT/2)
c
c If the "space" dimension is in the meridional direction, then replace
c "westward" with "southward", and "eastward" with "northward" in the above.
c
c If the two dimensions (n and t) are both spatial, then the above
c will apply to "left" and "right" sloping wave-forms instead.
c
c Prior to calling this subroutine, the following calls to initialization
c routines must be made.
c      call cffti(NL,WSAVE1)
c      call cffti(NT,WSAVE2)   
c These routines only need to be called once for each different value of
c NL or NT for which this power spectrum routine is called.
c
c "totvar" will contain the total variance of the input dataset as computed
c   by the FFTs.
c "mean" will contain the global mean of the input array.

      integer NL,NT,n,t,pn,pt
      real EEr(NL,NT),WSAVE1(4*NL+15),WSAVE2(4*NT+15)
      complex EE(NL,NT),CEE1(NL),CEE2(NT)
      real PEE(NL+1,NT+1)   ! note the unusual dimensions
      real totvar1,totvar2,mean,ptstart,ptend

c*** Big loop over "t", calling FFT
      do 200 t=1,NT
       do 150 n=1,NL
        EE(n,t) = CMPLX(EEr(n,t)) 
        CEE1(n) = EE(n,t)
 150   continue
       call cfftf(NL,CEE1,WSAVE1)
       do 170 n=1,NL
        EE(n,t) = CEE1(n)/float(NL)
 170   continue
 200  continue

c*** Big loop over "n", calling FFT.
      do 300 n=1,NL
       do 250 t=1,NT
        CEE2(t) = EE(n,t)
 250   continue
       call cfftf(NT,CEE2,WSAVE2)
       do 270 t=1,NT
        EE(n,t) = CEE2(t)/float(NT)
 270   continue
 300  continue

c Now the array EE(n,t) contains the (complex) space-time spectrum.

c Create array PEE(NL+1,NT/2+1) which contains the (real) power spectrum.
c  Note how the PEE array is arranged into a different order to EE.
      do 191 pt=1,NT+1
       do 189 pn=1,NL+1
        if(pn.le.NL/2) then
         n=NL/2+2-pn
         if(pt.le.NT/2) then
          t=NT/2+pt
         else
          t=pt-NT/2
         endif
        elseif(pn.ge.NL/2+1) then
         n=pn-NL/2
         if (pt.le.NT/2+1) then
          t=NT/2+2-pt
         else
          t=NT+NT/2+2-pt
         endif
        endif
        PEE(pn,pt) = (CABS(EE(n,t)))**2
 189   continue
 191  continue

c**************************
c The following computations of mean and total variance are useful checks.
c To save processing time, they may be skipped.
c     goto 4000 

c Global mean value is contained in this position of EE array.
      mean = REAL(EE(1,1))
      print*,'Global Mean from the EE(1,1) is',mean

c The total variance of the input dataset equals the sum over
c (almost) all of PEE.
       totvar1=0.
       do pt=1,NT         ! only going to NT, not NT+1
        do pn=1,NL        ! only going to NL, not NL+1
         if(pt.eq.NT/2+1.and.pn.eq.NL/2+1) then 
          continue        ! don't include the global mean in sum
         else
          totvar1=totvar1+PEE(pn,pt)
         endif
        enddo
       enddo
c Alternatively, the total variance equals twice the sum over almost half
c of PEE.
       totvar2=0.
       do pt=1,NT/2+1
        do pn=1,NL        ! only going to NL, not NL+1
         if(pt.eq.NT/2+1.and.pn.eq.NL/2+1) then 
          continue        ! don't include the global mean in sum
         elseif(pt.ne.1.and.pt.ne.NT/2+1)
     #       then
c$           Count variance for Nyquist frequency and time mean only once.
c$           Everything else gets multiplied by 2 when I do this sum of the
c$           variance only over half the bins. For this sum, I am taking
c$           advantage of the symmetrical nature of PEE.
          totvar2=totvar2+2.*PEE(pn,pt)
         else
          totvar2=totvar2+PEE(pn,pt)
         endif
        enddo
       enddo
       print*,'Total variance is, two different ways',totvar1,totvar2

4000   continue
c**************************

      END
