capture program drop sun_abraham
program sun_abraham, rclass
    syntax varlist(numeric ts fv) /// dependent_var covariates
        [if] [in]                 /// subset
        /// [aw fw pw iw]         /// [weight type = exp]
        ,                         ///
        Individuals(varname)      ///
        Time(varname)             /// time variable (can be relative or absolute)
        Event(varname)            /// time when when individual receives treatment
                                  /// (event); must be relative or absolute to match
                                  /// time() variable.
    [                             ///
        endpoints(str)            /// cut off leads and lags at endpoints; examples:
                                  ///     endpoints(. 5)
                                  ///     endpoints(-6 .)
                                  ///     endpoints(-6 5)
                                  ///
                                  ///
        Basecohorts(str)          /// cohorts (by event time) to exclude from estimation
                                  /// (i.e. event time to use as comparison groups).
        noshareadjust             /// [do not] adjust the variance for the shares
    ]

    * Parsing
    * -------

    unab varlist: `varlist'
    marksample touse, strok
    markout `touse' `individuals' `cohorts'

    gettoken Y X: varlist

    * Check options
    * -------------

    foreach var in time event {
        cap noi assert ``var'' == round(``var'') if `touse'
        if _rc {
            disp as err "`var'() variable '`var'' must be an integer."
            exit _rc
        }
    }

    qui count if `event' == `time' & `touse'
    if ( `r(N)' == 0 ) {
        disp as err "No treatment found in event (`event') and time (`time') variables."
        exit 198
    }

    if ( "`endpoints'" != "" ) {
        if ( `:list sizeof endpoints' != 2 ) {
            disp as err "endpoints() incorrectly specified; include exactly two (left and right)"
            exit 198
        }
        gettoken left_endpoint right_endpoint: endpoints
    }

    * Run main specification
    * ----------------------

    tempname V C W R b breshape
    tempname v vse delta dse
    tempvar E t share

    qui {
        egen `E' = group(`event') if `touse'
        sum `E' if `touse', meanonly
        local Elevels = `r(max)' - `r(min)' + 1
        local emax    = `r(max)'

        gen `t' = `time' - `event'
        if ( "`endpoints'" != "" ) {
            if !mi(`left_endpoint') {
                replace `t' = `left_endpoint' if `t' < `left_endpoint'
            }
            if !mi(`right_endpoint') {
                replace `t' = `right_endpoint' if `t' > `right_endpoint'
            }
        }

        sum `t' if `touse', meanonly
        local tlevels = `r(max)' - `r(min)' + 1
        local tmin    = `r(min)'
        replace `t' = `t' - `tmin' + 1 if `touse'
        replace `t' = 0 if (`time' - `event') == -1 & `touse'

        * Never-treated (`event' is missing) are not given dummies, so
        * they are coded the same way as the base cohorts.

        qui count if mi(`event') & `touse'
        local nnever = `r(N)'
        replace `E' = 0 if mi(`event') & `touse'
        replace `t' = 0 if mi(`event') & `touse'
    }

    if ( "`shareadjust'" == "" ) {
        * Variance from the share coefficients [xx should this have the `X'?]
        qui gen `share' = .
        estimates clear
        local estimates
        forvalues el = 1 / `Elevels' {
            * disp `el'
            tempname est`el'
            qui replace `share' = (`E' == `el') if `touse'
            qui reg `share' i(1/`tlevels').`t' if `touse', noc
            qui estimates store `est`el''
            local estimates `estimates' `est`el''
        }
        qui suest `estimates', vce(robust)
        mata Vshare = st_matrix("e(V)")

        mata sel = colshape(rowshape(1::rows(Vshare), `Elevels')[., 1::`tlevels']', 1)
        mata Vshare = Vshare[sel, sel]
    }

    if ( (`"`basecohorts'"' == "") & (`nnever' == 0) ) {
        qui sum `event' if `touse', meanonly
        local basecohorts `r(max)'
    }

    foreach ell of local basecohorts {
        replace `E' = 0 if `event' == `ell' & `touse'
    }

    * Step 1: Compute delta_el relative to base cohorts
    * -------------------------------------------------

    // if ( `"`weight'"' == "" ) local wgt [`weight' `exp']

    qui reghdfe `Y' i(1/`tlevels').`t'#i(1/`Elevels').`E' `X' if `touse' `in' `wgt', ///
        absorb(`individuals' `time') cluster(`individuals') noconstant

    mata `V' = st_matrix("e(V)")
    mata `b' = st_matrix("e(b)")'

    matrix `dse'   = J(`tlevels', `Elevels', .)
    matrix `delta' = J(`tlevels', `Elevels', .)

    forvalues tl = 1 / `tlevels' {
        forvalues el = 1 / `Elevels' {
            matrix `delta'[`tl', `el'] = _b[`tl'.`t'#`el'.`E']
            matrix `dse'[`tl', `el']   = _se[`tl'.`t'#`el'.`E']
        }
    }

    qui replace `t' = `t' + `tmin' - 1 if `touse'
    qui replace `t' = -1 if (`time' - `event') == -1 & `touse'
    qui replace `E' = `emax' if `E' == 0 & `touse'

    qui replace `t' = . if mi(`event') & `touse'
    qui replace `E' = . if mi(`event') & `touse'

    qui levelsof `event' if `touse'
    matrix colnames `delta' = `r(levels)'
    matrix colnames `dse'   = `r(levels)'

    qui levelsof `t' if `touse'
    matrix rownames `delta' = `r(levels)'
    matrix rownames `dse'   = `r(levels)'

    qui replace `t' = `t' - `tmin' + 1 if `touse'

    * Step 2: Compute weights
    * -----------------------

    tempvar N
    markout `touse' `event'
    egen `N' = sum(`touse'), by(`t' `E')

    mata `C' = sort(uniqrows(st_data(., "`t' `E' `N'",  st_local("touse"))), (1, 2))
    mata `W' = J(`tlevels', `Elevels', .)
    mata for (i = 1; i <= rows(`C'); i++) `W'[`C'[i, 1], `C'[i, 2]] = `C'[i, 3]
    mata `W' = `W' :* (st_matrix(st_local("delta")) :!= 0)
    mata `W' = `W' :/ rowsum(`W')
    mata _editmissing(`W', 0)

    * Step 3: Combine!
    * ----------------

    mata `R' = J(`tlevels', rows(`V'), 0)
    mata for(i = 1; i <= `tlevels'; i++) `R'[i, ((i - 1) * `Elevels' + 1)::(i * `Elevels')] = `W'[i, .]

    tempname addShare
    if ( "`shareadjust'" == "" ) {
        mata `breshape' = J(`tlevels', rows(Vshare), 0)
        mata for(i = 1; i <= `tlevels'; i++) `breshape'[i, ((i - 1) * `Elevels' + 1)::(i * `Elevels')] = st_matrix(st_local("delta"))[i, .]

        * Note that the variance from the share coefficient should be 0 if
        * there is only one group, as the share is 1 in that case.

        mata `addShare' = (rowsum(st_matrix(st_local("delta")) :!= 0) :> 1) :* (`breshape' * Vshare * `breshape'')
        mata mata drop Vshare
    }
    else {
        mata `addShare' = 0
    }

    mata st_matrix(st_local("v"),   `R' * `b')
    mata st_matrix(st_local("vse"), sqrt(diagonal(`R' * `V' * `R'' :+ `addShare')))
    // mata `R' * `b', sqrt(diagonal(`R' * `V' * `R''))

    matrix colnames `v'   = v
    matrix colnames `vse' = vse
    matrix rownames `v'   = `:rownames `delta''
    matrix rownames `vse' = `:rownames `delta''

    matrix list `v'
    matrix list `vse'
    matrix list `delta'
    matrix list `dse'

    return matrix v     = `v'
    return matrix vse   = `vse'
    return matrix delta = `delta'
    return matrix dse   = `dse'
end
