module runoff_mapper_mod

    use config_mod

    implicit none

    private

    public setup_runoff_maps
    public map_runoff

    integer               :: GatherMap(GridNX,GridNY)
    integer               :: ScatterMap(GridNX,GridNY)
    integer               :: CalvingMap(GridNX,GridNY)

    real(RNP)             :: CellArea(GridNX,GridNY)

    real(RNP),allocatable :: BasinRunoff(:)
    real(RNP),allocatable :: BasinCalving(:)
    real(RNP),allocatable :: InvDischargeArea(:)
    real(RNP),allocatable :: InvCalvingArea(:)

    real(RNP),pointer     :: RunoffATM(:,:)
    real(RNP),pointer     :: CalvingATM(:,:)
    real(RNP),pointer     :: RunoffOCE(:,:)
    real(RNP),pointer     :: CalvingOCE(:,:)

contains

    ! *=========================================================================
    ! * SUBROUTINE setup_runoff_maps
    ! *-------------------------------------------------------------------------
    ! * Read GatherMap and ScatterMap (see comment to map_runoff below for
    ! * details) as well as CellAreas from NetCDF files and initialise the rest
    ! * of the module variables.
    ! *=========================================================================
    subroutine setup_runoff_maps

        use NetCDF
        use CPLNG

        ! Locals
        integer :: ncid,varid,NumBasins
        integer :: i,j

        call nf90_chkerr(nf90_open(RunoffMapFileName,NF90_CLOBBER,ncid))

        call nf90_chkerr(nf90_inq_varid(ncid,'drainage_basin_id',varid))
        call nf90_chkerr(nf90_get_var(ncid,varid,GatherMap))

        call nf90_chkerr(nf90_inq_varid(ncid,'arrival_point_id',varid))
        call nf90_chkerr(nf90_get_var(ncid,varid,ScatterMap))

        call nf90_chkerr(nf90_inq_varid(ncid,'calving_point_id',varid))
        call nf90_chkerr(nf90_get_var(ncid,varid,CalvingMap))

        call nf90_chkerr(nf90_close(ncid))

        ! The variables 'drainage_basin' and 'arrival_point' from the
        ! runoff maps file are coded as follows:
        ! drainage_basin: >0 over land
        !                 -2 over sea
        ! arrival_point:  >0 discharge point (sea near coast)
        !                 -1 over land, lake
        !                 -2 over sea, no discharge point
        ! For the runoff mapper, we just need all values >0 (basin number) or 0
        ! for points that don't contribute to runoff to discharge areas. Hence,
        ! we cut off all values <0
        GatherMap  = MAX(0,GatherMap)
        ScatterMap = MAX(0,ScatterMap)
        CalvingMap = MAX(0,CalvingMap)

        ! The maximum values in GatherMap and ScatterMap indicate the number of
        ! basins, assuming that the basin numbers are continuous. Hence, both
        ! should be the same.
        NumBasins = MAXVAL(GatherMap)

        if (MAXVAL(ScatterMap)/=NumBasins .or. MAXVAL(CalvingMap)/=NumBasins) then
            call ERROR('Inconsistent number of basins in gather and scatter/calving maps')
        endif

        ! Allocate basin arrays. Note basin number 0, which is a virtual basin
        ! to take care of all points that do not contribute to a discharge area.
        allocate(BasinRunoff(0:NumBasins))
        allocate(InvDischargeArea(0:NumBasins))
        allocate(BasinCalving(0:NumBasins))
        allocate(InvCalvingArea(0:NumBasins))

        ! Read the cell areas from NetCDF file
        call nf90_chkerr(nf90_open(CellAreaFileName,NF90_CLOBBER,ncid))
        call nf90_chkerr(nf90_inq_varid(ncid,CellAreaVarName,varid))
        call nf90_chkerr(nf90_get_var(ncid,varid,CellArea))
        call nf90_chkerr(nf90_close(ncid))

        ! Compute the inverse of the discharge/calving areas for all basins
        InvDischargeArea = 0.0_RNP
        InvCalvingArea   = 0.0_RNP

        do j=1,GridNY
            do i=1,GridNX
                InvDischargeArea(ScatterMap(i,j)) = InvDischargeArea(ScatterMap(i,j)) + CellArea(i,j)
                InvCalvingArea(CalvingMap(i,j))   = InvCalvingArea(CalvingMap(i,j)) + CellArea(i,j)
            enddo
        enddo

        InvDischargeArea = 1.0_RNP / InvDischargeArea
        InvCalvingArea   = 1.0_RNP / InvCalvingArea

        ! Let the local runoff arrays point to the coupling fields
        ! handled by CPLNG
        RunoffATM  => CPLNG_FLD(CPLNG_IDX('R_Runoff_atm'))%D
        CalvingATM => CPLNG_FLD(CPLNG_IDX('R_Calving_atm'))%D
        RunoffOCE  => CPLNG_FLD(CPLNG_IDX('R_Runoff_oce'))%D
        CalvingOCE => CPLNG_FLD(CPLNG_IDX('R_Calving_oce'))%D

    contains

        subroutine nf90_chkerr(iret)
            integer,intent(in) :: iret
            if (iret /= NF90_NOERR) call ERROR(trim(nf90_strerror(iret)))
        end subroutine nf90_chkerr

    end subroutine setup_runoff_maps

    ! *=========================================================================
    ! * SUBROUTINE map_runoff
    ! *-------------------------------------------------------------------------
    ! * Distribute runoff from land over the ocean using river basins. This is
    ! * basically a gather-scatter operation with conversions between flux and
    ! * mass in between.
    ! * The runoff is received from the atmosphere (over land points, but that's
    ! * transparent to the runoff-mapper) after it has been interpolated by
    ! * OASIS to the intermediate grid used here. Access to the atmosphere
    ! * runoff is via RunoffATM that points to the coupling field handled by
    ! * CPLNG.
    ! * The runoff is then gathered by basins in the BasinRunoff array. The
    ! * gather operation uses GatherMap for the grid-point-to-basin mapping. The
    ! * runoff is converted from flux [kg/m^2] to mass [kg] in that step by
    ! * mupliplying with the CellArea. This is needed to distribute the overall
    ! * mass to the basins.
    ! * Finally, BasinRunoff is distributed (scattered) to the grid points by
    ! * using ScatterMap, which maps every basin to a number of discharge
    ! * points. The runoff is converted back from mass [kg] to flux [kg/m^2] by
    ! * dividing by the respective discharge area. The result is written to
    ! * RunoffOCE (pointer to coupling field) before it is sent to the ocean by
    ! * OASIS via CPLNG.
    ! *=========================================================================
    subroutine map_runoff
        ! Locals
        integer :: i,j,basin

        ! Gather runoff from the source grid (RunoffATM) to basins
        BasinRunoff  = 0.0_RNP
        BasinCalving = 0.0_RNP
        do j=1,GridNY
            do i=1,GridNX
                BasinRunoff(GatherMap(i,j)) = BasinRunoff(GatherMap(i,j)) + RunoffATM(i,j)*CellArea(i,j)
                BasinCalving(GatherMap(i,j)) = BasinCalving(GatherMap(i,j)) + CalvingATM(i,j)*CellArea(i,j)
            enddo
        enddo

        ! Scatter runoff from basins to target grid (RunoffOCE, CalvingOce)
        RunoffOCE  = 0.0_RNP
        CalvingOCE = 0.0_RNP
        do j=1,GridNY
            do i=1,GridNX
                basin = ScatterMap(i,j)
                if (basin>0) then
                    RunoffOCE(i,j)  = BasinRunoff(basin)*InvDischargeArea(basin)
                endif
                basin = CalvingMap(i,j)
                if (basin>0) then
                    CalvingOCE(i,j) = BasinCalving(basin)*InvCalvingArea(basin)
                endif
            enddo
        enddo

    end subroutine map_runoff

end module runoff_mapper_mod
