! Test runoff spreading for NEMO_ORCA1 coupled to AGCM
!
! NCS, 04/2017
!
! Added Case 4. Global spreading of runoff in subpolar North Atlantic
! D. Yang, Oct/2017
!
! Added Case 5. Global spreading of runoff in western subpolar North Atlantic
! D. Yang, Nov/2017
!
! Added Case 6. Global spreading of runoff in Hudson Bay and adjacent area
! D. Yang, Nov/2017
!
! Added Case 7. Global spreading of runoff north of ~40N
! D. Yang, 10/Nov/2017
! 
! Added Case 8. Global spreading of runoff south of ~40N
! D. Yang, 10/Nov/2017
!
! Added Case 9. Global spreading of runoff in subpolar North Pacific
! D. Yang, 22/Nov/2017
!
! Added Case 10. Global spreading of runoff in most of Arctic Ocean
! D. Yang, 22/Nov/2017
!
! Added Case 11. Global spreading of runoff in part of Arctic Ocean, Black and part of Caspian Seas.
! D. Yang, 22/Nov/2017
!
! Added spread_fraction and Case 12. Global spreading of one-half of the runoff north of ~40N 
! D. Yang, 15/Dec/2017


MODULE spreading
    ! Spreads CanAM river runoff to prevent negative salinities in CanNEMO.

    private
    public :: runoff_spreading

    integer, parameter :: dp=8

CONTAINS

  subroutine runoff_spreading(runoff_in, tmask, ocean_areas, runoff_modified)
      !======================================================================
      ! Spreads CanAM river runoff to prevent negative salinities in CanNEMO.
      !
      ! Brute force manual spreading with hard coded cell indices. Global
      ! integral runoff is preserved. 
      !
      ! ** NOTE ** Only valid for ORCA1.
      !
      ! Neil Swart, 04/2017
      !
      !=======================================================================
      IMPLICIT NONE
      REAL(dp), intent (in), DIMENSION(:,:)  :: runoff_in, tmask, ocean_areas
      REAL(dp), intent (out), DIMENSION(:,:) :: runoff_modified
      REAL(dp) glob_runoff_in, glob_runoff_out                                ! Globally integrated input and output runoff
      REAL(dp) :: spr_frac = 1.0_dp                                           ! Spreading fraction

      ! Verify this is ORCA1 sized grid.
      IF ((SIZE(runoff_in,1) /= 362) .OR. (SIZE(runoff_in,2) /=292)) THEN
          write(6,*)"runoff_spreading: O_Runoff NOT ORCA1 shaped. Exiting." 
          STOP 1
      ENDIF
      
      runoff_modified = runoff_in

      ! Trouble grid-points for runoff
      ! Case 1. ORCA1 Arctic: i=305:313, j=291 - > i=300:304, j=286:291
      !call weighted_spread(runoff_modified, ocean_areas*tmask, (/305,313,291,291/), (/300,304,287,291/), spr_frac, runoff_modified)

      ! Case 2. ORCA1 Arctic (north fold mirror): i=50:58,, j=291 -> i=59:63, j=287:291
      !call weighted_spread(runoff_modified, ocean_areas*tmask, (/50,58,291,291/), (/59,63,287,291/), spr_frac, runoff_modified)

      ! Case 3. Global spreading
      !call weighted_spread(runoff_modified, ocean_areas*tmask, (/1,362,1,292/), (/1,362,1,292/), runoff_modified)

      ! Case 4. Global spreading of runoff in subpolar North Atlantic (north of 39.9278N and 100.5W~20.5E)
      !call weighted_spread(runoff_modified, ocean_areas*tmask, (/188,309,207,291/), (/1,362,1,292/), runoff_modified)

      ! Case 5. Global spreading of runoff in western subpolar North Atlantic (north of 39.9278N and 100.5W~14.5W)
      !call weighted_spread(runoff_modified, ocean_areas*tmask, (/188,274,207,291/), (/1,362,1,292/), runoff_modified)

      ! Case 6. Global spreading of runoff in Hudson Bay and adjacent area
      !call weighted_spread(runoff_modified, ocean_areas*tmask, (/190,235,213,266/), (/1,362,1,292/), runoff_modified)

      ! Case 7. Global spreading of runoff north of ~40N
      !call weighted_spread(runoff_modified, ocean_areas*tmask, (/1,362,207,291/), (/1,362,1,292/), runoff_modified)

      ! Case 8. Global spreading of runoff south of ~40N
      !all weighted_spread(runoff_modified, ocean_areas*tmask, (/1,362,1,206/), (/1,362,1,292/), runoff_modified)

      ! Case 9. Global spreading of runoff in subpolar North Pacific
      !call weighted_spread(runoff_modified, ocean_areas*tmask, (/33,172,207,245/), (/1,362,1,292/), runoff_modified)

      ! Case 10. Global spreading of runoff in most of Arctic Ocean
      !call weighted_spread(runoff_modified, ocean_areas*tmask, (/33,172,246,291/), (/1,362,1,292/), runoff_modified)

      ! Case 11. Global spreading of runoff in part of Arctic Ocean, Black and part of Caspian Seas
      !call weighted_spread(runoff_modified, ocean_areas*tmask, (/310,362,207,291/), (/1,362,1,292/), runoff_modified)

      ! Case 12. Global spreading of one-half of the runoff north of ~40N
      !pr_frac = 0.25_dp
      !all weighted_spread(runoff_modified, ocean_areas*tmask, (/1,362,207,291/), (/1,362,1,292/), spr_frac, runoff_modified)

      ! Compute and print the global runoff integrals before and after spreading.
      glob_runoff_in = SUM(runoff_in * ocean_areas * tmask)
      glob_runoff_out = SUM(runoff_modified * ocean_areas * tmask)
      write(6,*)'runoff_spreading: Input runoff (1e9 kg/s):', glob_runoff_in/1e9_dp
      write(6,*)'runoff_spreading: Out runoff (1e9 kg/s)  :', glob_runoff_out/1e9_dp
  end subroutine runoff_spreading

  subroutine weighted_spread(field, weights, sidx, tidx, spr_frac, field_out)
      !
      ! Do a weighted spread, or redistribution, of a field from a source to target 
      ! region.
      ! 
      ! Weights are equivalent to grid-cell areas (that have been masked as appropriate)
      !
      IMPLICIT NONE
      REAL(dp), INTENT(IN), DIMENSION(:,:)   :: field, weights
      REAL(dp), INTENT(OUT), DIMENSION(:,:)  :: field_out
      INTEGER, DIMENSION(4)                  :: sidx, tidx  ! Format is:
                                                            ! (imin,imax,jmin,jmax)
      REAL(dp) :: source_integral, spreading_area, target_additive
      REAL(dp) :: spr_frac

      source_integral = spr_frac * SUM(field(sidx(1):sidx(2), sidx(3):sidx(4))*weights(sidx(1):sidx(2), sidx(3):sidx(4)))
      spreading_area = SUM(weights(tidx(1):tidx(2), tidx(3):tidx(4)))
      target_additive = source_integral / spreading_area
      field_out = field
      ! Remove from the source region
      field_out(sidx(1):sidx(2), sidx(3):sidx(4)) = field_out(sidx(1):sidx(2), sidx(3):sidx(4))  &
      &                                           - spr_frac * field(sidx(1):sidx(2), sidx(3):sidx(4))
      ! Add to the target region
      field_out(tidx(1):tidx(2), tidx(3):tidx(4)) = field_out(tidx(1):tidx(2), tidx(3):tidx(4)) + target_additive
  end subroutine

END MODULE spreading
