program catchment

  ! UW Land Surface Hydrology Group implementation of the Catchment model
  ! modified from NLDAS and SAC/NOAH implementation (originally by Ted Bohn)
  ! author: Kostas Andreadis, kostas@hydro.washington.edu

  ! driverMod contains definitions of all global driver variables
  USE driverMod
  USE process_cat
  USE catchment_model, ONLY: calc_soil_moist
  USE catch_constants, ONLY: ALHE

  implicit none

  INTEGER :: num_fsteps, tsteplen_save, FSTEP_COUNT
  INTEGER :: MODEL_STEP
  INTEGER :: FORCING_STEP
  INTEGER :: OUTPUT_STEP
  INTEGER :: MODEL_STEP_COUNT
  INTEGER :: tsteps_per_day,FORCING_STEP_OF_DAY
  INTEGER :: days_per_year,day_of_year
  INTEGER :: year,month,day,hour,minute,sec,julday
  INTEGER :: year_new,month_new,day_new,hour_new,minute_new,sec_new,day_of_year_new
  INTEGER :: timestp, K, I, J
  LOGICAL :: forc_exist
  INTEGER :: klo, khi
  REAL :: wlo, whi, akday(1:13)
  REAL, ALLOCATABLE :: srfexc_tmp(:), rzexc_tmp(:), catdef_tmp(:)
  REAL, ALLOCATABLE :: sfmc(:), rzmc(:), prmc(:)

!!$  DATA akday /-15.,16.,45.,75.,105.,136.,166.,197.,228.,258., &
!!$       289.,319.,350.,381./
  DATA akday /0.,31.,59.,90.,120.,151.,181.,212.,243.,273.,304.,334.,366./

  ! Initialize error status
  status = 0

  ! READ CONTROL FILE
  !IF(IARGC().NE.1)THEN
  !   PRINT*, 'USAGE:  catchment <control_file>'
  !   STOP
  !ENDIF
  CALL GETARG(1,CNTRFL)
  WRITE(*,*) 'Reading control file ',TRIM(CNTRFL)
  CALL READCNTL  

  ! This is necessary for calculation of IDT below
  IF (MODEL_DT < 3600) THEN
     WRITE(*,*) 'ERROR: model dt must be at least 3600 sec'
     STOP
  END IF

  ! Convert date/time information
  MODEL_DT_REAL = REAL(MODEL_DT)
  OUTPUT_DT_REAL = REAL(OUTPUT_DT)
  CALL GREG2JUL(YEAR0,MONTH0,DAY0,JULDAY0)
  CALL BUILD_DATE_STRING(YEAR0,MONTH0,DAY0,0,0,0,MODEL_START_TIME)

  ! READ MODEL GEOMETRY
  WRITE(*,*)'Reading model geometry ',TRIM(LSC)
  CALL GET_GRID

  ! Allocate arrays based on the model geometry
!!$  IF (nbands < 1 .or. SNOWBAND_FILE == '') THEN
!!$    nbands = 1
!!$  END IF
  CALL ALLOCATE_ARRAYS
  ALLOCATE(srfexc_tmp(landlen))
  ALLOCATE(rzexc_tmp(landlen))
  ALLOCATE(catdef_tmp(landlen))
  ALLOCATE(sfmc(landlen))
  ALLOCATE(rzmc(landlen))
  ALLOCATE(prmc(landlen))

  ! READ MODEL PARAMETERS
  WRITE(*,*)'Reading model parameters ',TRIM(LSC)
  CALL READ_LSC

  ! READ INITIAL CONDITIONS
  WRITE(*,*) 'Reading initial conditions from ',TRIM(INITIAL)
  CALL READ_INITIAL

  ! CHECK INITIAL CONDITIONS AND MODEL GEOMETRY
  WRITE(*,*) 'Checking initial conditions and model geometry'
  CALL CHECK_INITIAL

  ! BUILD CURRENT FILE SUFFIX
  WRITE(SUFFIX,'(".",i4.4,i2.2,".nc")')YEAR0,MONTH0

  ! OPEN ATMOSPHERIC FORCING DATA
  FORCFILE = TRIM(FORCING)//TRIM(SUFFIX)
  WRITE(*,*) 'Opening atmospheric data ',TRIM(FORCFILE)
  CALL OPEN_FORCING
  tsteplen_save = tsteplen

  ! OPEN OUTPUT FILES
  OUTFILES(1) = TRIM(RESULT)//"/wb"//TRIM(SUFFIX)
  WRITE(*,*) 'Opening output files:'
  CALL OPEN_OUTPUT
  WRITE(*,*) 'WB file: ',TRIM(OUTFILES(1))

  ! OPEN RESTART FILE
  RESTFILE = TRIM(RESTART)//TRIM(SUFFIX)
  WRITE(*,*) 'Opening restart file ',TRIM(RESTFILE)
  CALL OPEN_RESTART

  ! Initialize time counters
  CALL COMPUTE_NUM_FSTEPS(FORCING_DT,YEAR0,MONTH0,DAY0,YEAR_FINAL,MONTH_FINAL,DAY_FINAL,num_fsteps)
  MADTT = FORCING_DT/MODEL_DT
  IF (MADTT < 1) THEN
     WRITE(*,*)'ERROR: MODEL_DT',MODEL_DT,'IS GREATER THAN FORCING_DT',FORCING_DT
     STOP
  END IF
  OUTPUT_STEP_RATIO = OUTPUT_DT/MODEL_DT
  IF (OUTPUT_STEP_RATIO < 1) THEN
     WRITE(*,*)'ERROR: MODEL_DT',MODEL_DT,'IS GREATER THAN OUTPUT_DT',OUTPUT_DT
     STOP
  END IF
  IDTS = MODEL_DT
  ! This is from Yun's original LDAS code; I'm not sure I trust it
  !  IDT = INT(3600/MODEL_DT)
  ! This is my interpretation of how to handle IDT;
  ! since IDT is an integer, we can't handle sub-hourly time steps
  IDT = INT(MODEL_DT/3600)

  tsteps_per_day = 24*3600/FORCING_DT
  OUTPUT_STEP = 1
  MODEL_STEP_COUNT = 1
  year = YEAR0
  month = MONTH0
  day = DAY0
  day_of_year = JULDAY0
  hour = 0
  minute = 0
  sec = 0
  FORCING_STEP = 1
  FSTEP_COUNT = 1

  write(*,*)'tsteps_per_day ',tsteps_per_day
  write(*,*)'FORCING_DT ',FORCING_DT
  write(*,*)'MODEL_DT ',MODEL_DT
  write(*,*)'MADTT ',MADTT
  write(*,*)'OUTPUT_DT ',OUTPUT_DT
  write(*,*)'OUTPUT_STEP_RATIO ',OUTPUT_STEP_RATIO
  write(*,*)'num_fsteps ',num_fsteps

  ! INITIALIZE PREVIOUS TIMESTEP VALUES
  srfexc_tmp = CAT_PROGN%srfexc
  rzexc_tmp  = CAT_PROGN%rzexc
  catdef_tmp = CAT_PROGN%catdef
  CALL calc_soil_moist(landlen, CP%vegcls, CP%vgwmax, CP%cdcr1, CP%cdcr2, CP%wpwet, CP%poros, CP%psis, CP%bee, &
       CP%ars1, CP%ars2, CP%ars3, CP%ara1, CP%ara2, CP%ara3, CP%ara4, CP%arw1, CP%arw2, CP%arw3, CP%arw4, &
       srfexc_tmp, rzexc_tmp, catdef_tmp, sfmc, rzmc, prmc)
  DO I=1,landlen
  SoilMoistTotal(I) = (prmc(I)*(CP(I)%cdcr2/(1-CP(I)%wpwet)))/CP(I)%poros
  END DO
  SoilMoistTotal_old = SoilMoistTotal
  DO I = 1, landlen
     SWE(I) = 0
     SnowDepth(I) = 0
     DO J =1,N_snow
        SWE(I) = SWE(I) + CAT_PROGN(I)%wesn(J)
        SnowDepth(I) = SnowDepth(I) + CAT_PROGN(I)%sndz(J)
     END DO
  END DO
  SWE_old = SWE

  ! READ ATMOSPHERIC FORCING DATA - BEGINNING OF 1ST FORCING STEP
  IF (FORCING_STEP .LE. tsteplen) THEN
     CALL READ_FORCING(FORCING_STEP,SWdownnow,LWdownnow,Tairnow,Qairnow,Rainfnow,Snowfnow,PSurfnow,Windnow)
  ELSE
     WRITE(*,*)'ERROR: FORCING_STEP',FORCING_STEP,'not found in forcing file',FORCFILE
     STOP
  ENDIF

  ! RUN TIME AND SPATIAL LOOP
  DO

     ! READ ATMOSPHERIC FORCING DATA - BEGINNING OF NEXT FORCING STEP
     IF (FSTEP_COUNT .LT. num_fsteps) THEN
        ! At least one more forcing step remains in this simulation
        IF (FORCING_STEP .LT. tsteplen) THEN
           ! Next forcing step is contained in current forcing file
           CALL READ_FORCING(FORCING_STEP+1,SWdownplus,LWdownplus,Tairplus,Qairplus,Rainfplus,Snowfplus,PSurfplus,Windplus)
        ELSE
           ! We've reached the end of the current forcing file(s)
           tsteplen_save = tsteplen
           CALL CLOSE_FORCING
           ! Build new forcing file name(s)
           year_new = year
           month_new = month
           day_new = day
           hour_new = hour
           minute_new = minute
           sec_new = sec
           CALL INCREMENT_DATE(FORCING_DT,year_new,month_new,day_new,hour_new,minute_new,sec_new)
           WRITE(SUFFIX_NEW,'(".",i4.4,i2.2,".nc")') year_new,month_new
           FORCFILE = TRIM(FORCING)//TRIM(SUFFIX_NEW)
           ! Check whether next forcing files exist
           INQUIRE(FILE=FORCFILE,EXIST=forc_exist)
           IF (forc_exist) THEN
              ! Open new forcing file(s)
              WRITE(*,*) 'Opening atmospheric data ',TRIM(FORCFILE)
              CALL OPEN_FORCING
              CALL READ_FORCING(1,SWdownplus,LWdownplus,Tairplus,Qairplus,Rainfplus,Snowfplus,PSurfplus,Windplus)
           ELSE
              ! Forcings have run out before end of simulation
              WRITE(*,*) 'ERROR: Cannot find next atmospheric data file',TRIM(FORCFILE)
              STOP
           ENDIF
        ENDIF
     ELSE
        ! We are on the final forcing step of the simulation; set *plus vars equal to *now vars
        SWdownplus = SWdownnow
        LWdownplus = LWdownnow
        Tairplus   = Tairnow
        Qairplus   = Qairnow
        Rainfplus  = Rainfnow
        Snowfplus  = Snowfnow
        PSurfplus  = PSurfnow
        WINDplus   = WINDnow
        PotEvapplus   = PotEvapnow
     END IF

     FORCING_STEP_OF_DAY = mod((FORCING_STEP-1),tsteps_per_day) + 1

     IF (MODEL_STEP_COUNT == 1) THEN

        ! Initialize output variables that must be aggregated from model time step to output time step
        Rainf = 0.0
        Snowf = 0.0
        Evap = 0.0
        PotEvap = 0.0
        CanEvap = 0.0
        Transp = 0.0
        BareEvap = 0.0
        Qs   = 0.0
        Qsb  = 0.0
        Qsm  = 0.0
        Qst  = 0.0

     END IF

     ! For each forcing step, there are MADTT model timesteps
     DO MODEL_STEP = 1,MADTT

!!$        WRITE(*,*)'MODEL_STEP ', MODEL_STEP
!!$        WRITE(*,*)'OUTPUT_STEP ', OUTPUT_STEP
!!$        WRITE(*,*)'MODEL_STEP_COUNT ', MODEL_STEP_COUNT

        DO I = 1,landlen
           ! INTERPOLATE FORCING DATA TO MODEL TIME STEP
           ! This produces DT_* variables and CZMODEL, given *now and *plus variables and time/location
           CALL INTERP_FORCING(day_of_year,FORCING_STEP_OF_DAY,MODEL_STEP,I)

           ! Store interpolated forcings to Catchment forcing data structure
           MET_FORCE(I)%SWdown = SWdownnow(I)
           MET_FORCE(I)%LWdown = LWdownnow(I)
           MET_FORCE(I)%Psurf = PSurfnow(I)
           MET_FORCE(I)%Qair = Qairnow(I)
           MET_FORCE(I)%Tair = Tairnow(I)
           MET_FORCE(I)%Wind = Windnow(I)
           MET_FORCE(I)%Rainf = Rainfnow(I)
           MET_FORCE(I)%Snowf = Snowfnow(I)
           
            ! Interpolate Greenness and LAI
           DO K=2,13
              IF (day_of_year .le. akday(K) .and. day_of_year .gt. akday(K-1) )  khi=K
           END DO
           klo = khi - 1
           whi = (day_of_year - akday(klo)) / (akday(khi) - akday(klo))
           wlo = (akday(khi) - day_of_year) / (akday(khi) - akday(klo))
           IF (day_of_year .lt. akday(12)) THEN
              GREEN(I) = wlo * CGREEN(I,klo) + whi * CGREEN(I,khi)
              LAI(I) = wlo * CLAI(I,klo) + whi * CLAI(I,khi)
           ELSE ! in December interpolate with January's value
              GREEN(I) = wlo * CGREEN(I,klo) + whi * CGREEN(I,1)
              LAI(I) = wlo * CLAI(I,klo) + whi * CLAI(I,1)
           END IF
!!$           IF (day_of_year .ge. 334 .and. I .eq. 136) THEN
!!$              WRITE(*,*) day_of_year,LAI(I),GREEN(I),CLAI(I,klo),CLAI(I,klo-1),CLAI(I,khi),CLAI(I,khi-1),CGREEN(I,klo),CGREEN(I,khi),CGREEN(I,khi-1)
!!$           END IF
           
        END DO ! spatial loop

        ! CALL LAND-SURFACE PHYSICS
        CALL propagate_cat(landlen, MODEL_DT_REAL, day_of_year, LATL, LAND, CP, &
             MET_FORCE, GREEN, LAI, CZMODEL, MODIS_PARAM, CAT_PROGN, CAT_DIAGN)

        ! AGGREGATE OUTPUT VARIABLES TO OUTPUT TIME STEP
         DO I=1,landlen 
            ! fluxes from Catchment are in kg/m^2s units, therefore we accumulate them in kg/m^2
            ! and then divide by OUTPUT_DT to get back to kg/m^2s at the time of output
            Evap(I) = Evap(I) + CAT_DIAGN(I)%evap * MODEL_DT_REAL ! evaporation from Catchment is already in mm/s
            Qs(I) = Qs(I) + CAT_DIAGN(I)%runsrf * MODEL_DT_REAL ! surface runoff from Catchment is already in mm/s
            Qsb(I) = Qsb(I) + CAT_DIAGN(I)%bflow * MODEL_DT_REAL
            Rainf(I) = Rainf(I) + MET_FORCE(I)%Rainf * MODEL_DT_REAL
            Snowf(I) = Snowf(I) + MET_FORCE(I)%Snowf * MODEL_DT_REAL
            CanEvap(I) = CanEvap(I) + CAT_DIAGN(I)%eint / ALHE * MODEL_DT_REAL
            Transp(I) = Transp(I) + CAT_DIAGN(I)%eveg / ALHE * MODEL_DT_REAL
            BareEvap(I) = BareEvap(I) + CAT_DIAGN(I)%esoi / ALHE * MODEL_DT_REAL
        END DO

        IF (MODEL_STEP_COUNT == OUTPUT_STEP_RATIO) THEN

           ! rainfall and snowfall
           Rainf = Rainf / OUTPUT_DT_REAL
           Snowf = Snowf / OUTPUT_DT_REAL

           ! evaporation terms
           Evap = Evap / OUTPUT_DT_REAL
           CanEvap = CanEvap / OUTPUT_DT_REAL
           Transp = Transp / OUTPUT_DT_REAL
           BareEvap = BareEvap / OUTPUT_DT_REAL

           ! runoff
           Qs = Qs / OUTPUT_DT_REAL
           Qsb = Qsb / OUTPUT_DT_REAL

           ! canopy interception
           CanopInt = CAT_PROGN%capac

           ! average temperature
           AvgSurfT = CAT_DIAGN%tsurf

           ! soil moisture
          srfexc_tmp = CAT_PROGN%srfexc
          rzexc_tmp  = CAT_PROGN%rzexc
          catdef_tmp = CAT_PROGN%catdef
          CALL calc_soil_moist(landlen, CP%vegcls, CP%vgwmax, CP%cdcr1, CP%cdcr2, CP%wpwet, CP%poros, CP%psis, CP%bee, &
               CP%ars1, CP%ars2, CP%ars3, CP%ara1, CP%ara2, CP%ara3, CP%ara4, CP%arw1, CP%arw2, CP%arw3, CP%arw4, &
               srfexc_tmp, rzexc_tmp, catdef_tmp, sfmc, rzmc, prmc)
          SoilMoistTotal = 0.0
          DO I=1,landlen
          SoilMoistTotal(I) = (prmc(I)*(CP(I)%cdcr2/(1-CP(I)%wpwet)))/CP(I)%poros
	  END DO

          ! swe and snow depth
          DO I = 1, landlen
              SWE(I) = 0
              SnowDepth(I) = 0
              DO J =1,N_snow
                 SWE(I) = SWE(I) + CAT_PROGN(I)%wesn(J)
                 SnowDepth(I) = SnowDepth(I) + CAT_PROGN(I)%sndz(J)
              END DO
           END DO

           ! change in storage terms
           DelSoilMoist  = SoilMoistTotal - SoilMoistTotal_old
           DelSWE = SWE - SWE_old
           DelSurfStor   = 0.0
           DelIntercept  = 0.0

           ! Update "old" values
           SoilMoistTotal_old = SoilMoistTotal
           SWE_old       = SWE

           ! Write results for this output interval to output files
           CALL WRITE_OUTPUT(OUTPUT_STEP)

!!$           DO I=1,landlen
!!$              IF (AvgSurfT(I) /= AvgSurfT(I)) THEN
!!$                 WRITE(*,*) AvgSurfT(I),I,year,month,day
!!$              END IF
!!$           END DO
           WRITE(*,*) year,month,day

           ! re-initialize aggregated variables
           DO I = 1,landlen
              Qs(I)=0.0
              Qsb(I) = 0.0
              Evap(I)=0.0
              CanEvap(I) = 0.0
              Transp(I) = 0.0
              BareEvap(I) = 0.0
           END DO

           MODEL_STEP_COUNT = 1
           OUTPUT_STEP = OUTPUT_STEP + 1
           
        ELSE
           
           MODEL_STEP_COUNT = MODEL_STEP_COUNT + 1
           
        END IF

     END DO ! MADT loop

     ! Close output files if appropriate
     IF (FSTEP_COUNT == num_fsteps .OR. FORCING_STEP == tsteplen_save) THEN
        ! Close the output file(s)
        CALL CLOSE_OUTPUT
        ! Write restart file
        WRITE(*,*) 'Writing restart file ',RESTFILE
        CALL WRITE_RESTART
     ENDIF

     ! Check for end of simulation
     IF (FSTEP_COUNT == num_fsteps) THEN
        WRITE(*,*)'Reached end of simulation'
        EXIT
     ENDIF

     ! Save this forcing step's forcing values
     SWdownminus = SWdownnow
     LWdownminus = LWdownnow
     PSurfminus  = PSurfnow
     Tairminus   = Tairnow
     Qairminus   = Qairnow
     Windminus   = Windnow
     PotEvapminus= PotEvapnow

     SWdownnow = SWdownplus
     Rainfnow  = Rainfplus
     Snowfnow  = Snowfplus
     LWdownnow = LWdownplus
     PSurfnow  = PSurfplus
     Tairnow   = Tairplus
     Qairnow   = Qairplus
     Windnow   = Windplus
     PotEvapnow= PotEvapplus

     ! Increment forcing counters
     FORCING_STEP = FORCING_STEP + 1
     FSTEP_COUNT = FSTEP_COUNT +1 

     ! Calculate next forcing step's time counter values
     CALL INCREMENT_DATE(FORCING_DT,year,month,day,hour,minute,sec)
     CALL GREG2JUL(year,month,day,day_of_year)

     ! Open new output files if appropriate
     IF (FORCING_STEP .GT. tsteplen_save) THEN
        WRITE(SUFFIX,'(".",i4.4,i2.2,".nc")')year,month
        ! Check whether next forcing files exist
        FORCFILE = TRIM(FORCING)//TRIM(SUFFIX)
        INQUIRE(FILE=FORCFILE,EXIST=forc_exist)
        IF (forc_exist) THEN
           ! Open new restart file
           RESTFILE = TRIM(RESTART)//TRIM(SUFFIX)
           WRITE(*,*) 'Opening restart file ',TRIM(RESTFILE)
           CALL OPEN_RESTART
           ! Open new output file(s)
           OUTFILES(1) = TRIM(RESULT)//"/wb"//TRIM(SUFFIX)
           WRITE(*,*) 'Opening output files:'
           CALL BUILD_DATE_STRING(YEAR,MONTH,DAY,0,0,0,MODEL_START_TIME)
           CALL OPEN_OUTPUT
           WRITE(*,*) 'WB file: ',TRIM(OUTFILES(1))
           OUTPUT_STEP = 1
           tsteplen_save = tsteplen
           FORCING_STEP = 1
        ELSE
           IF (FSTEP_COUNT .LT. num_fsteps) THEN
              WRITE(*,*)'ERROR: Reached end of forcing files before end of simulation'
           ENDIF
           STOP
        ENDIF
     ENDIF

  END DO ! TIME LOOP


END PROGRAM catchment