*! version 1.0.0  15nov2023
program wgtdistrim
    
    version 16.1
    
    syntax varname(numeric) [ if ] [ in ] ///
    , Generate(namelist max=2)            ///
      UPper(real)                         ///
    [                                     ///
        LOwer(real 0)                     ///
        ITERate(integer 10)               ///
        TOLerance(real 0)                 ///
        NORMalize                         ///
    ]
    
    marksample touse
    
    local wgtvar : copy local varlist
    
    typlist_and_varlist_of `generate'
    
    /*
        We now have local macros
        
            wgtvar  <varname> holding untrimmed weights
            typlist  type of new variable from generate(), e.g., float
            varlist  name of new variable from generate()
    */
    
    if ( (`upper'<0) | (`upper'>=1) ) ///
        option_invalid upper() 125
    
    if ( (`lower'<0) | (`lower'>=1) ) ///
        option_invalid lower() 125
    
    if (`iterate' < 1) ///
        option_invalid iterate() 125
    
    if (`tolerance' < 0) ///
        option_invalid tolerance() 125
    
    mata : wgtdistrim(                ///
        st_local("wgtvar"),           ///
        st_local("touse"),            ///
        (`lower', 1-`upper'),         ///
        `iterate',                    ///
        `tolerance',                  ///
        ("`normalize'"=="normalize"), ///
        st_local("typlist"),          ///
        st_local("varlist")           ///
        )
    
end


program typlist_and_varlist_of
    
    capture noisily syntax newvarname(numeric)
    if ( _rc ) ///
        option_invalid generate() _rc
    
    if ("`varlist'" == "") ///
        option_invalid generate() 102
    
    c_local typlist : copy local typlist
    c_local varlist : copy local varlist
    
end


program option_invalid
    
    args option rc
    
    display as err "option `option' invalid"
    exit `rc'
    
end


/*  _________________________________________________________________________
                                                                     Mata  */

version 16.1


mata :


mata set matastrict   on
mata set mataoptimize on


void wgtdistrim(
    
    string scalar    wgtvar,
    string scalar    touse,
    real   rowvector lower_upper,
    real   scalar    iter,
    real   scalar    tolerance,
    real   scalar    normalize,
    string scalar    typlist,
    string scalar    varlist
    
    )
{
    real colvector w_kt
    real scalar    n
    real scalar    i
    
    
    w_kt = st_data(., wgtvar, touse)
    n    = rows(w_kt)
    
    confirm_sampling_weights(w_kt, n)
    
    summarize_iteration(0, minmax(w_kt), .)
    
    for (i=1; i<=iter; i++) {
        
        if (mreldif_w_kt(w_kt,n,lower_upper,i) <= tolerance)
            break
        
    }
    
    if ( normalize ) {
        
        w_kt = w_kt:/mean(w_kt)
        summarize_iteration(min((i,iter)), minmax(w_kt), .)
        
    }
    
    st_store(., st_addvar(typlist,varlist), touse, w_kt)
}


real scalar mreldif_w_kt(
    
    real colvector w_kt,
    real scalar    n,
    real rowvector lower_upper,
    real scalar    iteration
    
    )
{
    real scalar    w_bar
    real scalar    s2
    real scalar    alpha
    real scalar    beta
    real rowvector w_op
    real scalar    mreldif
    
    
    w_bar = mean(w_kt)
    s2    = quadvariance(w_kt)
    /*
        (6) in Potter (1990, 227) uses the population variance
        method of moments estimator, dividing by n, not (n-1)
        
    s2    = quadcolsum((w_kt:-w_bar):^2 / n)
    */
    
    alpha = (w_bar*(n*w_bar-1) / (n*s2)) + 2
    beta  = (n*w_bar-1)*(alpha-1)
    
    w_op = 1:/(n*invibetatail(alpha,beta,lower_upper))
    
    mreldif = trim_weights(w_kt, w_op)
    
    summarize_iteration(iteration, minmax(w_kt), mreldif)
    
    return(mreldif)
}


real scalar trim_weights(
    
    real colvector w_kt,
    real rowvector w_op
    
    )
{
    real matrix    kappa
    real scalar    gamma
    real colvector w_was
    
    
    /*
        The notation below follows Chen et al. (2017, 232)
    */
    
    kappa = (w_kt:<=w_op[1], w_kt:>=w_op[2])
    
    if ( !any(kappa) )
        return(0)
    
    w_was = w_kt
    
    gamma = (
        (quadcolsum(w_kt)-quadsum(kappa:*w_op)) 
        / 
        quadcolsum((1:-rowsum(kappa)):*w_kt)
        )
    
    w_kt = gamma:*w_kt
    
    if ( any(kappa[,1]) )
        w_kt[selectindex(kappa[,1])] = J(colsum(kappa[,1]),1,w_op[1])
    
    if ( any(kappa[,2] ) )
        w_kt[selectindex(kappa[,2])] = J(colsum(kappa[,2]),1,w_op[2])
    
    return( mreldif(w_kt,w_was) )
}


void summarize_iteration(
    
    real scalar    iteration,
    real rowvector min_max,
    real scalar    mreldif
    
    )
{
    printf("{txt}Iteration %f:", iteration)
    printf("{col 20}{txt}min = {res}%9.0g", min_max[1])
    printf("{col 40}{txt}max = {res}%9.0g", min_max[2])
    printf("{col 60}{txt}rel. diff = {res}%9.0g", mreldif)
    printf("\n")
}


void confirm_sampling_weights(
    
    real colvector w_kt,
    real scalar    n
    
    )
{
    if (n < 2)
        exit(error(2000+n)) // insufficient observations
    
    /*
        Potter (1990, 227)
        
            f(w) = n (1/nw)^(a+1)*(1 - 1/nw)^(b-1) / B(a,b)
            
                for 1/n <= w <= .
    */
    
    if ( !all((1/n):<=w_kt) ) {
        
        errprintf("weights must be greater than %f\n", 1/n)
        exit(any(w_kt:<0) ? 402 : 459)
        
    }
    
    if ( any(w_kt:<1) )
        printf("{txt}note: weights less than 1 encountered\n\n")
}


end


exit


/*  _________________________________________________________________________
                                                              version history

1.0.0   15nov2023   release on public GitHub repository
0.5.0   14nov2023   option upper() required; no default
                    minor refactoring
0.4.0   14nov2023   defaults upper(.01) lower(0)
                    use (unbiased) sample variance 
                    no longer allow zero weights
                    more informative error message for generate()
0.3.0   14nov2023   specify separate trimming levels for upper and lower bound
0.2.0   11nov2023   trim large weights from upper tail only
                    warning for sampling weights < 1
0.1.0   10nov2023   upload to private GitHub repository