* ==============================================================================
* ciquartile: Creating spike plots with interquarterile range
* Author: Zong Huang
*
* Based off of Alexander Staudt's cibar
* ==============================================================================

* over1 : over-variable
* over2 : additional over-variable
* over3 : additional over-variable
* bar : plot spikes as bar; default is scatter
* horizontal : plot spikes horizontally; default is vertically
* reverse : reverse order of spikes
* missing : omit missing spikes (buggy)
* gap2 : specify gap between spike groups (over2); default is 67% (of a spike)
* gap3 : specify gap between spike groups (over3); default is 100% (of a spike)
* spikeopts : specify the look of the spikes (global)
* spikecolor : specify the color of the spikes, delimited by spaces
* spikegap : specify gap within spike groups; default is 0% (of a spike)
* spike : specify the look for each spike, delimited by ||
* spikegroup : specify the look for each spike group (over2), delimited by ||
* ciopts : specify the look of the confidence intervals (global)
* cicolor : specify the color of the confidence intervals, delimited by spaces
* ci : specify the look for each confidence interval, delimited by ||
* cigroup : specify the look for each confidence interval group (over2), delimited by ||
* graphopts : specify additional graph options

program define ciquartile
    syntax anything(name = var) [if], ///
           over1(string) [over2(string) over3(string) ///
           bar horizontal missing reverse ///
           gap2(integer 67) gap3(integer 100) ///
           spikeopts(string asis) spikecolor(string asis) spikegap(integer 0) ///
           spike(string asis) spikegroup(string asis) ///
           ciopts(string asis) cicolor(string asis) ///
           ci(string asis) cigroup(string asis) ///
           graphopts(str asis)]

    preserve

    * ==========================================================================
    * Parse options
    * ==========================================================================

    * If option
    mark touse `if'
    qui keep if touse == 1

    * Direction option
    if "`horizontal'" == "horizontal" {
        local a1 "y"
        local a2 "x"
        local a1_angle "labgap(2) angle(0)"
        local a1_order "s median"
    }
    else {
        local a1 "x"
        local a2 "y"
        local a1_angle "labgap(2)"
        local a1_order "median s"
    }

    * Group options
    local bylist = "`over1' `over2' `over3'"
    local n_over = wordcount("`bylist'")
    local gap2 = `gap2' / 100 - 1
    local gap3 = `gap3' / 100

    * Spike options
    local spikeopts = "fintensity(inten100) `a1'axis(1 2 3)" + " " + `"`spikeopts'"'
    local sg = `spikegap' / 100

    * Confidence interval options
    local ciopts = "lcolor(gs8)" + " " + `"`ciopts'"'

    * ==========================================================================
    * Calculate mean and confidence interval
    * ==========================================================================

    * Calculate percentiles
    bysort `bylist': egen lb     = pctile(`var'), p(25)
    bysort `bylist': egen median = mean(`var')
    bysort `bylist': egen ub     = pctile(`var'), p(75)


    * Get percentiles
    keep `bylist' lb median ub
    duplicates drop

    foreach x in `bylist' {
        qui drop if `x' == . // Confidence intervals not calculated for missing groups
    }

    * ==========================================================================
    * Get group labels
    * ==========================================================================

    if `n_over' == 2 {
        capture confirm existence `over2'
        if _rc == 6 {
            local over2 = "`over3'"
            display as txt "You specified over3() but not over2(). The value of over3() is used in over2() instead."
        }
    }

    * Reindex group levels sequentially
    forvalues i = 1/`n_over' {
        qui tab `over`i''
        local labels`i' = `r(r)'
        qui levelsof(`over`i''), local(levels`i')
    }

    forvalues i = 1/`n_over' {
        qui gen over`i'_n = .
        tokenize "`levels`i''"
        forvalues j = 1/`labels`i'' {
            qui replace over`i'_n = `j' if `over`i'' == ``j''
        }
    }

    * Label sequential group levels
    forvalues i = 1/`n_over' {
        tokenize "`levels`i''"
        forvalues j = 1/`labels`i'' {
            local ltext : label (`over`i'') ``j''
            label define vlabel`i' `j' `"`ltext'"', add
            if `j' == `labels`i'' label value over`i'_n vlabel`i'
        }
    }

    forvalues i = 1/`n_over' {
        qui levelsof(over`i'_n), local(levels`i'_n)
    }

    * ==========================================================================
    * Generate graph labels (single group case)
    * ==========================================================================

    if `n_over' == 1 {
        * Adjust spacing for spikegap
        qui tab over1_n
        qui gen s = over1_n
        qui replace s = s * (1 + `sg')

        * Reverse spacing if applicable
        if "`reverse'" == "reverse" replace s = s * -1

        * Get spacing for spikes
        qui sum s
        local s_min = `r(min)' - .75
        local s_max = `r(max)' + .75

        * Get labels for spikes
        local axis = "`a1'label(`s_min'" + " " + `"" ""' + " " + "`s_max'" + " " + `"" ""' + ", noticks `a1_angle')"
    }

    * ==========================================================================
    * Generate graph labels (two group case)
    * ==========================================================================

    if `n_over' == 2 {
        * Add missing spikes
        forvalues i = 1/`labels1' {
            forvalues j = 1/`labels2' {
                qui tab median if over1_n == `i' & over2_n == `j'
                if `r(r)' == 0 {
                    local obsv = _N + 1
                    qui set obs `obsv'
                    qui replace over1_n = `i' in `obsv'
                    qui replace over2_n = `j' in `obsv'
                }
            }
        }
        sort over2_n over1_n

        * Adjust spacing for spikegap
        qui gen s = .
        tokenize "`levels1_n'"
        forvalues i = 1/`labels1' {
            qui replace s = over1_n + (``i'' - 1) * (`sg') if over1_n == ``i''
        }

        * Adjust spacing for gap2
        tokenize "`levels2_n'"
        forvalues i = 1/`labels2' {
            qui replace s = s + (`labels1' + 1 + `gap2' + `sg') * (`i' - 1) if over2_n == ``i''
        }

        * Get spacing for spikes
        forvalues i = 1/`labels2' {
            qui sum s if over2_n == `i'
            local s_`i' = `r(mean)'
            if `i' == `labels2' {
                local s_min = 0
                local s_max = `r(max)' + 1
            }
        }

        * Reverse spacing if applicable
        if "`reverse'" == "reverse" {
            replace s = s * -1
            forvalues i = 1/`labels2' {
                local s_`i' = `s_`i'' * -1
            }
            local s_max = `s_max' * -1
        }

        * Get labels for spikes
        local axis = "`a1'label(`s_min'" + " " + `"" ""'
        forvalues i = 1/`labels2' {
            local l_`i' : label (over2_n) `i'

            local axis = `"`axis'"' + " " + "`s_`i''" + " " + `"`"`l_`i''"'"'
        }
        local axis = `"`axis'"' + " " + "`s_max'" + " " + `"" ""' + " " + ", noticks `a1_angle')"
    }

    * ==========================================================================
    * Generate graph labels (three group case)
    * ==========================================================================

    if `n_over' == 3 {
        * Add missing spikes
        forvalues i = 1/`labels1' {
            forvalues j = 1/`labels2' {
                forvalues k = 1/`labels3' {
                    qui tab median if over1_n == `i' & over2_n == `j' & over3_n == `k'
                    if `r(r)' == 0 {
                        local obsv = _N+1
                        qui set obs `obsv'
                        qui replace over1_n = `i' in `obsv'
                        qui replace over2_n = `j' in `obsv'
                        qui replace over3_n = `k' in `obsv'
                    }
                }
            }
        }
        sort over3_n over2_n over1_n

        * Drop missing cross-groups (if applicable)
        qui egen over_cg = group(over3_n over2_n)

        if "`missing'" == "missing" {
            bysort over2_n over3_n: egen nonmissing = count(median)
            qui drop if nonmissing == 0
        }

        sort over_cg `over1'

        * Index cross-groupings
        qui tab over_cg
        local labels_cg = `r(r)'
        qui levelsof(over_cg), local(levels_cg)

        * Adjust spacing for spikegap
        qui gen index = _n
        qui gen over1_n2 = mod(index, `labels1')
        qui replace over1_n2 = `labels1' if over1_n2 == 0

        qui gen s = .
        tokenize "`levels1_n'"
        forvalues i = 1/`labels1' {
            qui replace s = over1_n2 + (``i'' - 1) * (`sg') if over1_n2 == ``i''
        }

        * Adjust spacing for gap2
        tokenize "`levels_cg'"
        forvalues i = 1/`labels_cg' {
            qui replace s = s + (`i' - 1) * (`labels1' + 1 + `gap2' + `sg') if over_cg == ``i''
        }

        * Adjust spacing for gap3
        forvalues i = 1/`labels3' {
            qui replace s = s + (-`gap2' + 1) * (`i' - 1) * `gap3' if over3_n == `i'
        }

        * Get spacing for spikes
        tokenize "`levels_cg'"

        forvalues i = 1/`labels_cg' {
            qui sum s if over_cg == ``i''
            local s_`i' = `r(mean)'
            if `i' == `labels2' {
                local s_min = 0
                local s_max = `r(max)' + 1
            }
        }

        * Reverse spacing if applicable
        if "`reverse'" == "reverse" {
            replace s = s * -1
            forvalues i = 1/`labels_cg' {
                local s_`i' = `s_`i'' * -1
            }
            local s_max = `s_max' * -1
        }

        * Get labels for cross-groupings
        tokenize "`levels3_n'"
        forvalues i = 1/`labels3' {
            tokenize "`levels2_n'"
            forvalues j = 1/`labels2' {

                local l`j' : label (over2_n) ``j''
                local y = `j' + (`i' - 1) * `labels2'
                local y_l = "`y'" + " " + `"`"`l`j''"'"'

                if `i' == 1 & `j' == 1 local vlabel = `"`y_l'"'
                else local vlabel = `"`vlabel'"' + " " + `"`y_l'"'
            }
        }

        qui label define vlabel `vlabel'
        qui label value over_cg vlabel

        * Get labels for spikes
        tokenize "`levels_cg'"

        local axis = "`a1'label(`s_min'" + " " + `"" ""'
        forvalues i = 1/`labels_cg' {
            local l_`i' : label (over_cg) ``i''

            local axis = `"`axis'"' + " " + "`s_`i''" + " " + `"`"`l_`i''"'"'
        }
        local axis = `"`axis'"' + " " + "`s_max'" + " " + `"" ""' + " " + ", noticks `a1_angle')"

        * Get labels for spike groupings
        forvalues i = 1/`labels3' {
            forvalues j = 1/`labels_cg' {
                qui sum s if over3_n == `i'
            }
            local ss_`i' = `r(mean)'
            local ll_`i' : label (over3_n) `i'

            if `i' == 1 local l3 = "`ss_`i''" + " " + `"`"`ll_`i''"'"'
            else local l3 = `"`l3'"' + " " + "`ss_`i''" + " " + `"`"`ll_`i''"'"'

            if `i' == `labels3' {
                local axis = `"`axis'"' + " " + "`a1'label(" + `"`l3'"' + ", noticks `a1_angle' axis(2))"
            }
        }
    }

    * ==========================================================================
    * Generate graph components (single group case)
    * ==========================================================================

    if `n_over' == 1 {
        * Generate graph components
        forvalues j = 1/`labels1' {
            if "`bar'" == "bar" {
                qui gen spike`j' = "bar median s if over1_n == `j', `spikeopts' `horizontal'"
            }
            else {
                qui gen spike`j' = "scatter `a1_order' if over1_n == `j', `spikeopts'"
            }

            qui gen rcap`j' = "rcap lb ub s if over1_n == `j', `ciopts' `horizontal'"
        }

        * Add colors to graph command
        if wordcount(`"`spikecolor'"') > 0 {
            qui gen scolors = `"`spikecolor'"'
            qui split scolors, gen(scolor)

            forvalues j = 1/`r(nvars)' {
                cap qui replace spike`j' = spike`j' + " color(" + scolor`j' + ")"
            }
        }

        if wordcount(`"`cicolor'"')  > 0 {
            qui gen rcolors = `"`cicolor'"'
            qui split rcolors, gen(rcolor)

            forvalues j = 1/`r(nvars)' {
                cap qui replace rcap`j' = rcap`j' + " lcolor(" + rcolor`j' + ")"
            }
        }

        * Add options to graph command
        if wordcount(`"`spike'"') > 0 {
            gen sopts = `"`spike'"'
            qui split sopts, gen(sopt) p("||")

            forvalues j = 1/`r(nvars)' {
                cap qui replace spike`j' = spike`j' + " " + sopt`j'
            }
        }

        if wordcount(`"`ci'"') > 0 {
            gen ropts = `"`ci'"'
            qui split ropts, gen(ropt) p("||")

            forvalues j = 1/`r(nvars)' {
                cap qui replace rcap`j' = rcap`j' + " " + ropt`j'
            }
        }
    }

    * ==========================================================================
    * Generate graph components (multiple group case)
    * ==========================================================================

    if `n_over' > 1 {
        * Generate graph components
        forvalues i = 1/`labels1' {
            forvalues j = 1/`labels2' {
                if "`bar'" == "bar" {
                    qui gen spike`i'`j' = "bar median s if over1_n == `i' & over2_n == `j', `spikeopts' `horizontal'"
                }
                else {
                    qui gen spike`i'`j' = "scatter `a1_order' if over1_n == `i' & over2_n == `j', `spikeopts'"
                }

                qui gen rcap`i'`j' = "rcap lb ub s if over1_n == `i' & over2_n == `j', `ciopts' `horizontal'"
            }
        }

        * Add colors to graph command
        if wordcount(`"`spikecolor'"') > 0 {
            qui gen scolors = `"`spikecolor'"'
            qui split scolors, gen(scolor)

            forvalues i = 1/`r(nvars)' {
                forvalues j = 1/`labels2' {
                    cap qui replace spike`i'`j' = spike`i'`j' + " color(" + scolor`i' + ")"
                }
            }
        }

        if wordcount(`"`cicolor'"')  > 0 {
            qui gen rcolors = `"`cicolor'"'
            qui split rcolors, gen(rcolor)

            forvalues i = 1/`r(nvars)' {
                forvalues j = 1/`labels2' {
                    cap qui replace rcap`i'`j' = rcap`i'`j' + " lcolor(" + rcolor`i' + ")"
                }
            }
        }

        * Add options to graph command
        if wordcount(`"`spike'"') > 0 {
            gen sopts = `"`spike'"'
            qui split sopts, gen(sopt) p("||")

            forvalues i = 1/`r(nvars)' {
                forvalues j = 1/`labels2' {
                    cap qui replace spike`i'`j' = spike`i'`j' + " " + sopt`i'
                }
            }
        }

        if wordcount(`"`ci'"') > 0 {
            gen ropts = `"`ci'"'
            qui split ropts, gen(ropt) p("||")

            forvalues i = 1/`r(nvars)' {
                forvalues j = 1/`labels2' {
                    cap qui replace rcap`i'`j' = rcap`i'`j' + " " + ropt`i'
                }
            }
        }

        if wordcount(`"`spikegroup'"') > 0 {
            gen sgopts = `"`spikegroup'"'
            qui split sgopts, gen(sgopt) p("||")

            forvalues i = 1/`labels1' {
                forvalues j = 1/`r(nvars)' {
                    cap qui replace spike`i'`j' = spike`i'`j' + " " + sgopt`j'
                }
            }
        }

        if wordcount(`"`cigroup'"') > 0 {
            gen rgopts = `"`cigroup'"'
            qui split rgopts, gen(rgopt) p("||")

            forvalues i = 1/`labels1' {
                forvalues j = 1/`r(nvars)' {
                    cap qui replace rcap`i'`j' = rcap`i'`j' + " " + rgopt`j'
                }
            }
        }
    }

    * ==========================================================================
    * Generate graph
    * ==========================================================================

    * Concatenate graph components
    if "`bar'" == "bar" {
        qui egen graph = concat(spike* rcap*), punct(") (")
    }
    else {
        qui egen graph = concat(rcap* spike*), punct(") (")
    }
    qui replace graph = "(" + graph + ")"
    local graph = graph[1]

    * ==========================================================================
    * Get legend (single group case)
    * ==========================================================================

    if `n_over' == 1 {
        forvalues j = 1/`labels1' {
            local l`j' : label (over1_n) `j'
            if "`bar'" == "bar" {
                local v`j' = `j'
            }
            else {
                local v`j' = `j' + `labels1'
            }

            if `j' == 1 {
                local legend1 = `"`v`j'' "`l`j''""'
            }
            else {
                local legend1 = `"`legend1'"' + " " + `"`v`j'' "`l`j''""'
            }
        }

        local legend = `"legend(order(`legend1') region(lcolor(white)))"'
    }

    * ==========================================================================
    * Get legend (multiple group case)
    * ==========================================================================

    if `n_over' > 1 {
        forvalues j = 1/`labels1' {
            local l`j' : label (over1_n) `j'
            if "`bar'" == "bar" {
                local v`j' = `j' * 2
            }
            else {
                local v`j' = (`j' + `labels1') * 2
            }

            if `j' == 1 {
                local legend1 = `"`v`j'' "`l`j''""'
            }
            else {
                local legend1 = `"`legend1'"' + " " + `"`v`j'' "`l`j''""'
            }
        }

        local legend = `"legend(order(`legend1') region(lcolor(white)))"'
    }

    * ==========================================================================
    * Plot graph
    * ==========================================================================
    local format = `"`legend' plotregion(margin(bargraph)) graphregion(color(white))"' + " " + ///
                   `"`a2'title("") `a1'title("") `a1'title("", axis(2)) `a1'title("", axis(3))"' + " " + ///
                   `"`a1'scale(axis(2) noline) `a1'scale(axis(3) noline)"' + " " + ///
                   `"`a1'label(none, axis(2) noticks) `a1'label(none, axis(3) noticks)"' + " " + ///
                   `"`a1'label(, nogrid) `a2'label(, grid format(%9.3g))"'

    dis `"`graph', `format' `axis' `graphopts'"'

    twoway `graph', `format' `axis' `graphopts'
end

exit