;+
; NAME:
;    ukkola_drought
;
; PURPOSE:
;    This function identifies the occurrence of meteorological drought using 
;    the Ukkula and others (2020) method.
;
; CATEGORY:
;    Time Series Analysis
;
; CALLING SEQUENCE:
;    result = ukkola_drought( pr_data )
;
; INPUTS:
;    PR_DATA:  A float vector containing a time series of monthly total 
;        precipitation.  Of length N_TIME.
;    ID_THRESHOLD, STATS, THRESHOLD
;
; KEYWORD PARAMETERS:
;    DAILY:  If input then the data is interpreted as being daily rather than 
;        monthly.  This function then applies the Ukkula algorithm but with 
;        running 91-day averages rather than the 3-month averages.  A string 
;        specifying the calendar.  Supported values are:  '360_day', '365_day'.
;    DURATION:  If STATS is set, then this returns an N_DROUGHT vector 
;        reporting the duration of each of the N_DROUGHT identified drought 
;        events, in units of months.
;    ID_THRESHOLD:  An optional integer two-element vector specifying the 
;        indices of the first and last elements in PR_DATA to use for 
;        estimating the seasonally-varying threshold.  The default is to use 
;        all of PR_DATA.
;    INTENSITY_MAX:  If STATS is set, then this returns an N_DROUGHT vector 
;        reporting the maximum monthly intensity each of the N_DROUGHT 
;        identified drought events, with intensity defined as the difference 
;        between the smoothed precipitation and annual cycle of the drought 
;        threshold.
;    INTENSITY_MEAN:  If STATS is set, then this returns an N_DROUGHT vector 
;        reporting the average intensity each of the N_DROUGHT identified 
;        drought events, with intensity defined as the difference between the 
;        smoothed precipitation and annual cycle of the drought threshold.
;    START_ID:  If STATS is set, then this returns an N_DROUGHT vector 
;        reporting the index of the first month in each of the N_DROUGHT 
;        identified drought events.
;    STATS:  If set then statistics concerning the drought are returned in 
;        DURATION, INTENSITY_MAX, INTENSITY_MEAN, and START_ID.  The default is 
;        for these outputs to not be returned.
;    THRESHOLD:  An optional input vector of length 12 months defining the 
;        monthly precipitation threshold for defining drought.  If not input, 
;        this keyword parameter returns the estimated 12 monthly threshold 
;        values.
;
; OUTPUTS:
;    RESULT:  Returns an integer vector of length N_TIME with values of 1 
;        marking the occurrence of drought and of 0 zero marking the absence of 
;        drought.
;    DURATION, INTENSITY_MAX, INTENSITY_MEAN, START_ID, THRESHOLD
;
; USES:
;    filter.pro
;    quantile_threshold.pro
;
; PROCEDURE:
;    This function used the method of Ukkola and others (2020) to identify 
;    meteorological drought events.  Drought is defined as occurring if the 
;    3-month running average of precipitation falls below the 15th percentile 
;    for that month.  A drought event is any continuous sequence of drought 
;    months.
;
; EXAMPLE:
;    ; Define a seasonally-invariant threshold of 50
;    threshold = 50. + fltarr( 12 )
;    ; Define a two-year precipitation vector with four values below the 
;    ; threshold
;    pr_data = 51. + fltarr( 24 )
;    pr_data[11:14] = [ 49., 35., 45., 47. ]
;    ; Identify the drought
;    result = ukkola_drought( pr_data, threshold=threshold, stats=1, $
;        duration=duration, intensity_max=intensity_max, $
;        intensity_mean=intensity_mean, start_id=start_id )
;    ; RESULT will contain zeroes except for five consecutive ones in 
;    ; RESULT[11:15].  The drought designation in RESULT[15] comes from the 
;    ; 3-month smoothing operation, as we transition from drought to no-drought 
;    ; conditions.
;    ; The START_ID will be 11, DURATION will be 5 months, INTENSITY_MAX will 
;    ; be 7.67, and INTENSITY_MEAN will be 4.47.  The non-integer intensity 
;    ; values arise from the 3-month smoothing.
;
; REFERENCES:
;    Ukkola, A. M., M. G. De Kauwe, M. L. Roderick, G. Abramowitz, and A. J. 
;        Pitman.  2020.  Robust future changes in meteorological drought in 
;        CMIP6 projections despite uncertainty in precipitation.  Geophysical 
;        Research Letters, 47, e2020GL087820, 
;        https://doi.org/10.1029/2020GL087820
;
; MODIFICATION HISTORY:
;    Written by:  Daithi A. Stone (dastone@runbox.com), 2026-06-11
;    Modified:  DAS, 2026-07-02 (Added DAILY keyword input with support for 
;        daily data)
;-

;***********************************************************************

FUNCTION UKKOLA_DROUGHT, $
    PR_DATA, $
    DAILY=daily, $
    ID_THRESHOLD=id_threshold, $
    THRESHOLD=pr_thresh_abs, $
    DURATION=drought_duration, $
    INTENSITY_MAX=drought_intensity_max, $
      INTENSITY_MEAN=drought_intensity_mean, $
    START_ID=drought_start, $
    STATS=stats_opt

;***********************************************************************
; Constants and Options

; The option to work with daily data
daily_opt = keyword_set( daily )

; If we are working with daily data
if daily_opt eq 1 then begin
  ; The number of days (time steps) in a year
  if daily eq '360_day' then begin
    n_tina = 360
  endif else if daily eq '365_day' then begin
    n_tina = 365
  endif else begin
    ; Not yet supported
    stop
  endelse
; If we are working with monthly data
endif else begin
  ; The number of months (time steps) in a year
  n_tina = 12
endelse

; The number of months in the input data
n_time = n_elements( pr_data )
; The number of complete years
n_year = n_time / n_tina
if n_time lt 1 then stop

; The smoothing filter width (in months or days)
if daily_opt eq 0 then begin
  ; Three months
  filter_len = 3
endif else begin
  ; Ninety days
  filter_len = 91
endelse

; The quantile threshold for drought
pr_thresh_quant = 0.15

; Ensure proper threshold input
if max( n_elements( pr_thresh_abs ) eq [ 0, n_tina ] ) ne 1 then stop

;***********************************************************************
; Identify drought

; Take 3-month moving average
pr_smooth = filter( pr_data, filter_len, 'boxcar' )

; Calculate absolute thresholds if not input
if not( keyword_set( pr_thresh_abs ) ) then begin
  ; Initialise vector of absolute thresholds for calendar months
  pr_thresh_abs = fltarr( n_tina )
  ; Initialise vector of start of years within reference period
  index = lindgen( n_year ) * n_tina
  if keyword_set( id_threshold ) then begin
    id = where( ( index ge id_threshold[0] ) and ( index le id_threshold[1] ), $
        n_id )
    if n_id le 3 then stop
    index = index[id]
  endif
  ; Iterate through calendar months or days
  for i_month = 0, n_tina - 1 do begin
    ; Extract all instances of this month or calendar day
    temp = pr_smooth[index+i_month]
    ; Estimate the value for the drought quantile
    pr_thresh_abs[i_month] = quantile_threshold( temp, pr_thresh_quant )
  endfor
endif

; Initialise output drought flag vector
drought_flag = intarr( n_time )
; Iterate through time
for i_time = 0l, n_time - 1l do begin
  ; Flag if this value is less than the threshold for the calendar month
  drought_flag[i_time] = pr_smooth[i_time] le pr_thresh_abs[i_time mod n_tina]
endfor

;***********************************************************************
; Estimate drought statistics

; If statistics are requested
if keyword_set( stats_opt ) then begin
  ; Determine the start of the first drought
  id = ( where( drought_flag eq 1, n_id ) )[0]
  ; If we did not find any droughts
  if n_id eq 0 then begin
    ; Return no-value flags
    drought_start = -1
    drought_duration = -1
    drought_intensity_max = -1
    drought_intensity_mean = -1
    n_drought = 0
  ; If we have droughts
  endif else begin
    ; Record the start and duration of the first drought
    drought_start = [ id[0] ]
    drought_duration = [ 1l ]
    n_drought = 1l
    ; Initialise in-drought flag
    in_flag = 1l
    ; Interate through time from then next month
    for i_time = id[0] + 1l, n_time - 1l do begin
      ; If this is a drought month
      if drought_flag[i_time] eq 1 then begin
        ; If this is a new drought
        if in_flag eq 0 then begin
          ; Start recording this drought
          drought_start = [ drought_start, i_time ]
          drought_duration = [ drought_duration, 1 ]
          n_drought = n_drought + 1
          in_flag = 1
        ; If we are already in a drought
        endif else begin
          ; Add to the duration
          drought_duration[n_drought-1] = drought_duration[n_drought-1] + 1
        endelse
      ; If this is not a drought month
      endif else begin
        ; Flag that we are no longer in drought
        in_flag = 0
      endelse
    endfor
    ; Restrict to droughts at least 30 days long
    if daily_opt eq 1 then begin
      ; Find droughts at least 30 days long
      id = where( drought_duration ge 30, n_drought )
      ; If there are none
      if n_drought eq 0 then begin
        ; Return no-value flags
        drought_start = -1
        drought_duration = -1
        drought_intensity_max = -1
        drought_intensity_mean = -1
      ; If there are some droughts
      endif else begin
        ; Restrict to those droughts
        drought_start = drought_start[id]
        drought_duration = drought_duration[id]
      endelse
    endif
    ; Estimate drought intensity
    if n_drought gt 0 then begin
      drought_intensity_mean = fltarr( n_drought )
      drought_intensity_max = fltarr( n_drought )
      for i_drought = 0, n_drought - 1 do begin
        id = drought_start[i_drought] + indgen( drought_duration[i_drought] )
        temp = pr_thresh_abs[id mod n_tina] - pr_smooth[id]
        drought_intensity_mean[i_drought] = mean( temp )
        drought_intensity_max[i_drought] = max( temp )
      endfor
    endif
  endelse
endif

;***********************************************************************
; The End

return, drought_flag
END