!
MODULE surf_land_bucket_hetero_mod
!
! 2025/04 A. Maison (adapted from surf_land_bucket_mod)
! Surface land bucket module with heterogeneous continental sub-surfaces
! This module is used when no external land model is choosen and iflag_hetero_surf = 1 or 2.
!
  IMPLICIT NONE

CONTAINS

  SUBROUTINE surf_land_bucket_hetero(itime, jour, knon, knindex, debut, dtime,&
       tsurf, p1lay, tq_cdrag, precip_rain, precip_snow, temp_air, &
       spechum, petAcoef, peqAcoef, petBcoef, peqBcoef, pref, plev, &
       u1, v1, gustiness, rugoro, swnet, lwnet, &
       snow, qsol, agesno, tsoil, &
       qsurf, z0m, z0h, alb1_new, alb2_new, evap, &
       fluxsens, fluxlat, tsurf_new, dflux_s, dflux_l, &
       tsurf_tersrf, tsoil_tersrf, qsurf_tersrf, tsurf_new_tersrf, &
       cdragm_tersrf, cdragh_tersrf, &
       swnet_tersrf, lwnet_tersrf, fluxsens_tersrf, fluxlat_tersrf)

    USE clesphys_mod_h
    USE dimsoil_mod_h, ONLY: nsoilmx
    USE yomcst_mod_h, ONLY: RD, RG, RCPD, RSIGMA
    USE compbl_mod_h
    USE dimpft_mod_h
    USE limit_read_mod
    USE surface_data
    USE simplehydrol_mod
    USE calcul_fluxs_mod
    USE cpl_mod
    USE dimphy
    USE geometry_mod, ONLY: longitude,latitude 
    USE mod_grid_phy_lmdz
    USE mod_phys_lmdz_para
    USE indice_sol_mod
    USE phys_state_var_mod, ONLY: frac_tersrf, ratio_z0m_z0h_tersrf, z0m_tersrf, &
                                  albedo_tersrf, beta_tersrf, inertie_tersrf, &
                                  hcond_tersrf
    USE surf_param_mod, ONLY: eff_surf_param, average_surf_var
    USE cdrag_mod
    USE albsno_mod, ONLY :  albsno
    USE calbeta_mod, ONLY :  calbeta
    USE calbeta_clim_mod, ONLY :  calbeta_clim

#ifdef ISO
    use infotrac_phy, ONLY: niso
#endif

!****************************************************************************************
! Bucket calculations for surface. 
!****************************************************************************************
!
! Input variables  
!****************************************************************************************
    INTEGER, INTENT(IN)                     :: itime, jour, knon
    INTEGER, DIMENSION(knon), INTENT(IN)    :: knindex
    LOGICAL, INTENT(IN)                     :: debut
    REAL, INTENT(IN)                        :: dtime
    REAL, DIMENSION(knon), INTENT(IN)       :: tsurf
    REAL, DIMENSION(knon), INTENT(IN)       :: p1lay
    REAL, DIMENSION(knon), INTENT(IN)       :: tq_cdrag
    REAL, DIMENSION(knon), INTENT(IN)       :: precip_rain, precip_snow
    REAL, DIMENSION(knon), INTENT(IN)       :: temp_air, spechum
    REAL, DIMENSION(knon), INTENT(IN)       :: petAcoef, peqAcoef
    REAL, DIMENSION(knon), INTENT(IN)       :: petBcoef, peqBcoef
    REAL, DIMENSION(knon), INTENT(IN)       :: pref
    REAL, DIMENSION(knon), INTENT(IN)       :: u1, v1, gustiness
    REAL, DIMENSION(knon), INTENT(IN)       :: rugoro
    REAL, DIMENSION(knon), INTENT(IN)       :: swnet, lwnet
    REAL, DIMENSION(klon, nbtersrf), INTENT(IN) :: tsurf_tersrf


! In/Output variables
!****************************************************************************************
    REAL, DIMENSION(knon), INTENT(INOUT)          :: snow, qsol
    REAL, DIMENSION(knon), INTENT(INOUT)          :: agesno
    REAL, DIMENSION(knon, nsoilmx), INTENT(INOUT) :: tsoil
    REAL, DIMENSION(knon), INTENT(INOUT)          :: plev
    REAL, DIMENSION(knon, nsoilmx, nbtersrf), INTENT(INOUT) :: tsoil_tersrf


! Output variables
!****************************************************************************************
    REAL, DIMENSION(knon), INTENT(OUT)       :: qsurf
    REAL, DIMENSION(knon), INTENT(OUT)       :: z0m, z0h
    REAL, DIMENSION(knon), INTENT(OUT)       :: alb1_new, alb2_new
    REAL, DIMENSION(knon), INTENT(OUT)       :: evap, fluxsens, fluxlat
    REAL, DIMENSION(knon), INTENT(OUT)       :: tsurf_new
    REAL, DIMENSION(knon), INTENT(OUT)       :: dflux_s, dflux_l
    !
    REAL, DIMENSION(klon, nbtersrf), INTENT(OUT) :: tsurf_new_tersrf
    REAL, DIMENSION(klon, nbtersrf), INTENT(OUT) :: qsurf_tersrf
    REAL, DIMENSION(klon, nbtersrf), INTENT(OUT) :: cdragm_tersrf, cdragh_tersrf
    REAL, DIMENSION(klon, nbtersrf), INTENT(OUT) :: swnet_tersrf, lwnet_tersrf
    REAL, DIMENSION(klon, nbtersrf), INTENT(OUT) :: fluxsens_tersrf, fluxlat_tersrf

#ifdef ISO
!    REAL, DIMENSION(ntiso,klon), INTENT(OUT) :: xtevap
!    REAL, DIMENSION(klon),       INTENT(OUT) :: h1
!    REAL, DIMENSION(niso,klon),  INTENT(OUT) :: xtrunoff_diag
    REAL, DIMENSION(knon)  :: runoff_diag
!    REAL, DIMENSION(niso,klon),  INTENT(IN)  :: Rland_ice
#endif
! Local variables
!****************************************************************************************
    REAL, DIMENSION(knon) :: soilcap, soilflux
    REAL, DIMENSION(knon) :: cal, beta, dif_grnd, capsol
    REAL, DIMENSION(knon) :: alb_neig, alb_lim, icesub, icemelt
    REAL, DIMENSION(knon) :: zfra
    REAL, DIMENSION(knon) :: radsol
    REAL, DIMENSION(knon) :: u0, v0, u1_lay, v1_lay
    REAL, DIMENSION(knon) :: dummy_riverflow,dummy_coastalflow 
    INTEGER               :: i, j, k
#ifdef ISO
    INTEGER               :: ixt
    REAL, PARAMETER       :: t_coup = 273.15
    REAL, DIMENSION(knon) :: fq_fonte_diag
    REAL, DIMENSION(knon) :: fqfonte_diag
    REAL, DIMENSION(knon) :: snow_evap_diag 
    REAL, DIMENSION(knon) :: fqcalving_diag 
    REAL                  :: max_eau_sol_diag  
    REAL, DIMENSION(knon) :: run_off_lic_diag 
    REAL :: coeff_rel_diag
#endif
    !
    REAL, DIMENSION(knon)           :: zlev, geop1, speed, pblh, ri_in, sst
    REAL, DIMENSION(knon)           :: beta_eff, inertie_eff, conv_ratio_eff
    REAL, DIMENSION(knon)           :: meansqT
    REAL, DIMENSION(klon, nbtersrf) :: z0h_tersrf, emis_tersrf, conv_ratio_tersrf
    REAL, DIMENSION(klon, nbtersrf) :: evap_tersrf
    REAL, DIMENSION(klon, nbtersrf) :: dflux_s_tersrf, dflux_l_tersrf
    REAL, DIMENSION(klon, nbtersrf) :: radsol_tersrf
    REAL, DIMENSION(klon, nbtersrf) :: zri_tersrf
    INTEGER, PARAMETER              :: klon_1D = 1

!
!****************************************************************************************

    ! *** Calculations common to the two flag values ***

    ! average albedo
    CALL eff_surf_param(klon, nbtersrf, albedo_tersrf, frac_tersrf, 'ARI', alb_lim)

    ! suppose zero surface speed
    u0(:)=0.0
    v0(:)=0.0
    u1_lay(:) = u1(:) - u0(:)
    v1_lay(:) = v1(:) - v0(:)
    speed(:) = (u1_lay(:)**2 + v1_lay(:)**2)**0.5
    !
    geop1(1:knon) = RD * temp_air(1:knon) / (0.5*(pref(1:knon)+p1lay(1:knon))) &
            * (pref(1:knon)-p1lay(1:knon))
    zlev(1:knon) = (plev(1:knon)*RD*temp_air(1:knon)/(pref(1:knon)*RG))/2.
    !
    ! compute roughness lengths
    DO i=1, knon
      DO j=1, nbtersrf
        IF (ratio_z0m_z0h_tersrf(i,j) .NE. 0.) THEN
          z0h_tersrf(i,j) = z0m_tersrf(i,j) / ratio_z0m_z0h_tersrf(i,j)
        ELSE
          z0h_tersrf(i,j) = 1.E-12
        ENDIF
      ENDDO
    ENDDO

    CALL eff_surf_param(klon, nbtersrf, z0m_tersrf, frac_tersrf, 'CDN', z0m, zlev)
    CALL eff_surf_param(klon, nbtersrf, z0h_tersrf, frac_tersrf, 'CDN', z0h, zlev)

    DO i=1, knon
       z0m(i) = MAX(1.5e-05,SQRT(z0m(i)**2 + rugoro(i)**2))
    END DO

    ! compute the ratio to convert and print soil depths in meters (conv_ratio = (cond/cap)^0.5 and cap = I^2/cond)
    DO j=1, nbtersrf
       conv_ratio_tersrf(:,j) = hcond_tersrf(:,j)/inertie_tersrf(:,j)
    ENDDO

    !
    ! *** Surface parameter aggregation ***
    !
    IF (iflag_hetero_surf == 1) THEN
      !* Calcultaion of fluxes 
      
      ! calculate total absorbed radiance at surface
      radsol(:) = 0.0
      radsol(1:knon) = swnet(1:knon) + lwnet(1:knon)

      ! calculate constants (needeed for capsol and dif_grnd)
      CALL calbeta(dtime, is_ter, knon, snow, qsol, beta, capsol, dif_grnd)
      IF (type_veget=='betaclim') THEN
        CALL calbeta_clim(knon,jour,latitude(knindex(1:knon)),beta)
      ENDIF

      ! mean evapotranspiration coefficient
      CALL eff_surf_param(klon, nbtersrf, beta_tersrf, frac_tersrf, 'ARI', beta_eff)
      beta = beta_eff

      ! calculate temperature, heat capacity and conduction flux in soil
      IF (soil_model) THEN
         CALL eff_surf_param(klon, nbtersrf, inertie_tersrf, frac_tersrf, 'ARI', inertie_eff)
         CALL eff_surf_param(klon, nbtersrf, conv_ratio_tersrf, frac_tersrf, 'ARI', conv_ratio_eff)
         !
         CALL soil_hetero(dtime, is_ter, knon, snow, tsurf, qsol,  & 
           & longitude(knindex(1:knon)), latitude(knindex(1:knon)), tsoil, soilcap, soilflux, &
           & inertie_eff, conv_ratio_eff)
         !
         DO i=1, knon
            cal(i) = RCPD / soilcap(i)
            radsol(i) = radsol(i)  + soilflux(i)
         END DO
      ELSE 
         cal(:) = RCPD * capsol(:)
         IF (klon_glo .EQ. 1) THEN
           cal(:) = 0.
         ENDIF
      ENDIF

      ! calculate fluxes
      CALL calcul_fluxs(knon, is_ter, dtime, &
           tsurf, p1lay, cal, beta, tq_cdrag, tq_cdrag, pref, &
           precip_rain, precip_snow, snow, qsurf,  &
           radsol, dif_grnd, temp_air, spechum, u1_lay, v1_lay, gustiness, &
           1.,petAcoef, peqAcoef, petBcoef, peqBcoef, &
           tsurf_new, evap, fluxlat, fluxsens, dflux_s, dflux_l)
    
      
      !* Calculate snow height, run_off, age of snow
#ifdef ISO
    DO i=1,knon
      ! initialisation:
      fqfonte_diag(i)  =0.0
      fq_fonte_diag(i) =0.0
      snow_evap_diag(i)=0.0
    ENDDO !DO i=1,knon 
#endif    

      CALL simplehydrol( knon, is_ter, knindex, dtime, &
           tsurf, precip_rain, precip_snow, &
           snow, qsol, tsurf_new, evap, icesub, icemelt &
#ifdef ISO    
     & ,fq_fonte_diag,fqfonte_diag,snow_evap_diag,fqcalving_diag   &
     & ,max_eau_sol_diag,runoff_diag,run_off_lic_diag,coeff_rel_diag   &
#endif
     &  )
    
      ! calculate the age of snow
      CALL albsno(klon,knon,dtime,agesno(:),alb_neig(:), precip_snow(:))  
    
      WHERE (snow(1 : knon) .LT. 0.0001) agesno(1 : knon) = 0.
    
      DO i=1, knon
         zfra(i) = MAX(0.0,MIN(1.0, snow(i)/(snow(i)+10.0)))
         alb_lim(i)  = alb_neig(i) *zfra(i) + alb_lim(i)*(1.0-zfra(i))
      END DO

      
      !* Return albedo : 
      !    alb1_new and alb2_new are here given the same values
      
      alb1_new(:) = 0.0
      alb2_new(:) = 0.0
      alb1_new(1:knon) = alb_lim(1:knon)
      alb2_new(1:knon) = alb_lim(1:knon)
       
      !* Send to coupler
      !  The run-off from river and coast are not calculated in the bucket modele.
      !  For testing purpose of the coupled modele we put the run-off to zero.
      IF (type_ocean=='couple') THEN
         dummy_riverflow(:)   = 0.0
         dummy_coastalflow(:) = 0.0
         CALL cpl_send_land_fields(itime, knon, knindex, &
              dummy_riverflow, dummy_coastalflow)
      ENDIF

    !
    ! ***  Flux aggregation ***
    !
    ELSE IF (iflag_hetero_surf == 2) THEN
      ! initialize output tables
      evap_tersrf(:,:)      = 0.
      fluxsens_tersrf(:,:)  = 0.
      fluxlat_tersrf(:,:)   = 0.
      tsurf_new_tersrf(:,:) = 0.
      dflux_s_tersrf(:,:)   = 0.
      dflux_l_tersrf(:,:)   = 0.
      radsol_tersrf(:,:)    = 0.
      swnet_tersrf(:,:)     = 0.
      lwnet_tersrf(:,:)     = 0.
      ! hyp: surface emissivity = 1
      emis_tersrf(:,:)      = 1.
 
      ! * calculate total absorbed radiance at surface
      DO j=1, nbtersrf
        ! SW
        swnet_tersrf(klon_1D,j) = (1. - albedo_tersrf(klon_1D,j)) / (1. - alb_lim(klon_1D)) * swnet(klon_1D)
        ! LW
        ! first order
        lwnet_tersrf(klon_1D,j) = lwnet(klon_1D) + 4. * emis_tersrf(klon_1D,j) * RSIGMA * tsurf(klon_1D)**3 * &
                                  (tsurf(klon_1D) - tsurf_tersrf(klon_1D,j))
      ENDDO
      ! LW second order corrections
      !- net = dwn -up; up=sig( T4 + 4sum%T3T' + 6sum%T2T'2 +...)
      IF (iflag_order2_sollw == 1) THEN
        meansqT(:) = 0. ! as working buffer
        !
        DO j=1, nbtersrf
          meansqT(klon_1D) = meansqT(klon_1D) + (tsurf_tersrf(klon_1D,j) - tsurf(klon_1D))**2 * frac_tersrf(klon_1D,j)
        ENDDO
        DO j=1, nbtersrf
          lwnet_tersrf(klon_1D,j) = lwnet_tersrf(klon_1D,j) + 6. * RSIGMA * tsurf(klon_1D)**2 * (meansqT(klon_1D) - &
                              (tsurf(klon_1D) - tsurf_tersrf(klon_1D,j))**2)
        ENDDO
      ENDIF
      ! net radiation
      radsol_tersrf(:,:) = swnet_tersrf(:,:) + lwnet_tersrf(:,:)
 
      ! * compute evapotranspiration coefficient
      capsol(:) = 1.0/(2.5578E+06*0.15)
      dif_grnd(:) = 0.

      ! unused variables in cdrag routine
      pblh(:) = 0.
      ri_in(:) = 0.
      sst(:) = 0.

      ! Loop on sub-surfaces
      DO j=1, nbtersrf
        ! * drag coefficients
        CALL cdrag(knon, is_ter, speed, temp_air, spechum, geop1, pref, pblh, &
          tsurf_tersrf(:,j), qsurf_tersrf(:,j), z0m_tersrf(:,j), z0h_tersrf(:,j), ri_in, 0, &
          cdragm_tersrf(:,j), cdragh_tersrf(:,j), zri_tersrf(:,j), plev, precip_rain, sst, p1lay)

        ! * calculate temperature, heat capacity and conduction flux in soil
        IF (soil_model) THEN
          !
          CALL soil_hetero(dtime, is_ter, knon, snow, tsurf_tersrf(:,j), qsol,  &
            longitude(knindex(1:knon)), latitude(knindex(1:knon)), tsoil_tersrf(:,:,j), soilcap, soilflux, &
            inertie_tersrf(:,j), conv_ratio_tersrf(:,j))
          !
          cal(:) = RCPD / soilcap(:)
          radsol_tersrf(:,j) = radsol_tersrf(:,j)  + soilflux(:)
          !
        ELSE
          cal = RCPD * capsol
          IF (klon_glo .EQ. 1) THEN
            cal = 0.
          ENDIF
        ENDIF

        ! * calcultaion of fluxes 
        CALL calcul_fluxs(knon, is_ter, dtime, &
             tsurf_tersrf(klon_1D,j), p1lay, cal, beta_tersrf(klon_1D,j), cdragh_tersrf(klon_1D,j), cdragh_tersrf(klon_1D,j), pref, &
             precip_rain, precip_snow, snow, qsurf_tersrf(klon_1D,j),  &
             radsol_tersrf(klon_1D,j), dif_grnd, temp_air, spechum, u1_lay, v1_lay, gustiness, &
             1.,petAcoef, peqAcoef, petBcoef, peqBcoef, &
             tsurf_new_tersrf(klon_1D,j), evap_tersrf(klon_1D,j), fluxlat_tersrf(klon_1D,j), fluxsens_tersrf(klon_1D,j), &
             dflux_s_tersrf(klon_1D,j), dflux_l_tersrf(klon_1D,j))

        ! if snow > 0
        ! calculate snow height, run_off, age of snow
        CALL simplehydrol( knon, is_ter, knindex, dtime, &
             tsurf_tersrf(:,j), precip_rain, precip_snow, &
             snow, qsol, tsurf_new_tersrf(:,j), evap_tersrf(:,j), icesub, icemelt &
#ifdef ISO    
     & ,fq_fonte_diag,fqfonte_diag,snow_evap_diag,fqcalving_diag   &
     & ,max_eau_sol_diag,runoff_diag,run_off_lic_diag,coeff_rel_diag   &
#endif
     &  )

      ENDDO ! loop on sub-surfaces

      ! calculate the age of snow
      CALL albsno(klon,knon,dtime,agesno(:),alb_neig(:), precip_snow(:))  
      WHERE (snow(1 : knon) .LT. 0.0001) agesno(1 : knon) = 0.
    
      DO i=1, knon
        zfra(i) = MAX(0.0,MIN(1.0, snow(i)/(snow(i)+10.0)))
        alb_lim(i)  = alb_neig(i) *zfra(i) + alb_lim(i)*(1.0-zfra(i))
      END DO

      ! return albedo : 
      !    alb1_new and alb2_new are here given the same values
      alb1_new(:) = 0.0
      alb2_new(:) = 0.0
      alb1_new(1:knon) = alb_lim(1:knon)
      alb2_new(1:knon) = alb_lim(1:knon)
       
      ! send to coupler
      !  the run-off from river and coast are not calculated in the bucket modele.
      !  for testing purpose of the coupled modele we put the run-off to zero.
      IF (type_ocean=='couple') THEN
         dummy_riverflow(:)   = 0.0
         dummy_coastalflow(:) = 0.0
         CALL cpl_send_land_fields(itime, knon, knindex, &
              dummy_riverflow, dummy_coastalflow)
      ENDIF

      ! * average of fluxes and surface variables
      qsurf = average_surf_var(klon, nbtersrf, qsurf_tersrf, frac_tersrf, 'ARI')
      tsurf_new = average_surf_var(klon, nbtersrf, tsurf_new_tersrf, frac_tersrf, 'ARI')
      evap = average_surf_var(klon, nbtersrf, evap_tersrf, frac_tersrf, 'ARI')
      fluxlat = average_surf_var(klon, nbtersrf, fluxlat_tersrf, frac_tersrf, 'ARI')
      fluxsens = average_surf_var(klon, nbtersrf, fluxsens_tersrf, frac_tersrf, 'ARI')
      dflux_l = average_surf_var(klon, nbtersrf, dflux_l_tersrf, frac_tersrf, 'ARI')
      dflux_s = average_surf_var(klon, nbtersrf, dflux_s_tersrf, frac_tersrf, 'ARI')
      DO k=1, nsoilmx
        tsoil(:,k) = average_surf_var(klon, nbtersrf, tsoil_tersrf(:,k,:), frac_tersrf, 'ARI')
      ENDDO

      ! order 2 correction to tsurf_new, for radiation computations (main atm effect of Ts)
      IF (iflag_order2_sollw == 1) THEN
        meansqT(:) = 0. ! as working buffer
        DO j=1, nbtersrf
          meansqT(klon_1D) = meansqT(klon_1D)+(tsurf_tersrf(klon_1D,j)-tsurf_new(klon_1D))**2 *frac_tersrf(klon_1D,j)
        ENDDO
        tsurf_new(:) = tsurf_new(:)+1.5*meansqT(:)/tsurf_new(:)
      ENDIF   ! iflag_order2_sollw == 1

    ENDIF ! iflag_hetero_surf
!
!* End
!
  END SUBROUTINE surf_land_bucket_hetero
!
!****************************************************************************************
!
END MODULE surf_land_bucket_hetero_mod
