clear all
version 14
set more off
set varabbrev off
reghdfe, compile
set matsize 2000
cap mkdir temp

adopath + source/lib/stata/heterogeneous


set graphics off
program main
    import delimited source/lib/designs_main.csv, varn(1) clear
    levelsof design, loc(databases)

    local release_large ../datastore/analysis/release_large
    foreach database of local databases {
        local options release_large(`release_large') database(`database')
        load_data,    `options'

        estimate_bjs, mata(`database'_bjs)
        estimate_sa,  mata(`database'_sa)
        estimate_cdh, mata(`database'_cdh)

        export_runs, `options' mata(`database'_bjs) tag(bjs)
        export_runs, `options' mata(`database'_sa)  tag(sa)
        export_runs, `options' mata(`database'_cdh) tag(cdh)
    }
end


capture program drop load_data
program load_data
    syntax, release_large(str) database(str)

    local imopts varn(1) case(preserve) clear
    import delimited `"`release_large'/simulation/`database'.csv"', `imopts'
    if ( strpos("`:type eventtime'", "str") ) replace eventtime = "" if lower(trim(eventtime)) == "nan"
    destring eventtime, replace
	gen post = (eventtime >= 0)
	replace post = 0 if mi(eventtime)

    * NOTE: We compute post for use in cdh before manipulating
    * eventtime. Is this right?

    tempvar first last
    isid run index time
    gen event = time - eventtime
    egen `first' = min(time), by(index run)
    egen `last'  = max(time), by(index run)
    replace eventtime = . if (event < `first') | (event > `last')
    replace event = . if (event < `first') | (event > `last')
end


capture program drop estimate_bjs
program estimate_bjs
    syntax, MATAsave(str) [MAXiter(real 0)]

    qui levelsof run, loc(runs)
    mata `matasave' = J(0, 3, .)
    local maxiter = cond(`maxiter' == 0, ., `maxiter')
    local niter = 0
    timer clear 1

    foreach run of local runs {
        qui timer on 1
        if ( `++niter' < `maxiter' ) {
            disp `niter'
            preserve
                qui keep if run == `run'
                timer on 1
                qui did_imputation y index time event, pretrends(6) h(0/5) minn(0)
                mata post = st_matrix("r(table)")[1, 1::6]'
                mata pre  = st_matrix("r(table)")[1, 12::7]'
                mata post = post :- pre[length(pre)]
                mata pre  = pre  :- pre[length(pre)]
                mata coef = pre \ post
                mata ix   = -6::5
                timer off 1
                timer list
                mata `matasave' = `matasave' \ (J(rows(coef), 1, `run'), ix, coef)
            restore
        }
    }
end


capture program drop estimate_sa
program estimate_sa
    syntax, MATAsave(str) [MAXiter(real 0)]

    qui levelsof run, loc(runs)
    mata `matasave' = J(0, 3, .)
    local maxiter = cond(`maxiter' == 0, ., `maxiter')
    local niter = 0
    timer clear 1

    foreach run of local runs {
        qui timer on 1
        if ( `++niter' < `maxiter' ) {
            disp `niter'
            preserve
                qui keep if run == `run'
                timer on 1
                qui sun_abraham y, i(index) t(time) event(event) b(.) noshareadjust
                timer off 1
                timer list
                mata ix  = strtoreal(st_matrixrowstripe("r(v)")[., 2])
                mata sel = selectindex((ix :>= -6) :* (ix :<= 5))
                mata `matasave' = `matasave' \ (J(rows(sel), 1, `run'), ix[sel], st_matrix("r(v)")[sel])
            restore
        }
    }
end


capture program drop estimate_cdh
program estimate_cdh
    syntax, MATAsave(str) [MAXiter(real 0)]

    qui levelsof run, loc(runs)
    mata `matasave' = J(0, 3, .)
    local maxiter = cond(`maxiter' == 0, ., `maxiter')
    local niter = 0
    timer clear 1

    foreach run of local runs {
        qui timer on 1
        if ( `++niter' < `maxiter' ) {
            disp `niter'
            preserve
                qui keep if run == `run'
                timer on 1
                qui did_multiplegt y index time post, robust_dynamic longdiff_placebo  placebo(5) dynamic(5)
                local effect_0 = e(effect_0)
                forvalues x = 1/5 {
                    local effect_`x'  = e(effect_`x')
                    local placebo_`x' = e(placebo_`x')
                }
                timer off 1
                timer list
                mata coefs = J(12, 2, 0)
                forvalues x = 1/12{
                    mata coefs[`x', 1] = -7+`x'
                }
                forvalues x = 1/5{
                    mata coefs[`x', 2] = `placebo_`=6-`x'''
                }
                forvalues x = 7/12{
                    mata coefs[`x', 2] = `effect_`=`x'-7''
                }
                mata `matasave' = `matasave' \ (J(rows(coefs), 1, `run'), coefs)
            restore
        }
    }
end


capture program drop export_runs
program export_runs
    syntax, release_large(str) database(str) MATAsave(str) [tag(str)]

    clear
    mata stata(sprintf("set obs %g", rows(`matasave')))
    mata (void) st_addvar(("long", "long", "double"), ("run", "eventtime", "coef"))
    mata st_store(., ("run", "eventtime", "coef"), `matasave')
    save `"`release_large'/estimation/HeterogeneityRobust/`database'_`tag'_results.dta"', replace

    egen p2_5  = pctile(coef), by(eventtime) p(2.5)
    egen p50   = pctile(coef), by(eventtime) p(50)
    egen p97_5 = pctile(coef), by(eventtime) p(97.5)

    keep  eventtime p2_5 p50 p97_5
    order eventtime p2_5 p50 p97_5
    duplicates drop

    save `"`release_large'/estimation/HeterogeneityRobust/`database'_`tag'_pctiles.dta"', replace
    export delimited `"`release_large'/estimation/HeterogeneityRobust/`database'_`tag'_pctiles.csv"', replace
end


main

