*! waffle v1.11 (05 May 2024)
*! Asjad Naqvi and Jared Colston

*v1.11 (05 May 2024): Bug fixes to how data was collapsed under different conditions. normvar now needs to be already the sum value.
*v1.1  (04 Apr 2024): Re-release. Allows wide and long form data.
*v1.0  (01 Mar 2022): First release by Jared Colston

* Code is based on the Waffle guide on Medium: https://medium.com/the-stata-guide/stata-graphs-waffle-charts-32afc7d6f6dd



capture program drop waffle 

program waffle, sortpreserve
version 15

syntax varlist(numeric min=1) [if] [in], ///
	[	///
		by(varname) over(varname) normvar(varname numeric) percent showpct format(string) palette(string)  ///
		ROWDots(real 20) COLDots(real 20) MSYMbol(string) MLWIDth(string) MSize(string)	 ///
		NDSYMbol(string)  NDSize(string) NDColor(string)	///   // No Data = ND
		COLs(real 4) LEGPOSition(string) LEGCOLumns(real 4) LEGSize(string) NOLEGend margin(string) ///
		aspect(numlist max=1 >0) note(passthru) title(passthru) subtitle(passthru) * 	]

	
	// check dependencies
	cap findfile colorpalette.ado
	if _rc != 0 {
		display as error "The palettes package is missing. Install the {stata ssc install palettes, replace:palettes} and {stata ssc install colrspace, replace:colrspace} packages."
		exit
	}	
	
	marksample touse, strok
	
preserve																			
quietly {
	
	keep if `touse'
	tokenize `varlist'		
	
	
	// preamble
	
	if "`over'" == ""	{
		gen _over = 1 
		local over _over
		local ovswitch  = 1
	}
	

	keep `varlist' `normvar' `over' `by'

	local length : word count `varlist'
	
	if `length' > 1 {
		foreach x of local varlist {
			local `x'_lab : var label `x'
		}
		
		collapse (sum) `varlist' `normvar', by(`over')

		// flatten to long
		foreach x of local varlist  {
			ren `x' y_`x'
		}	
		
		gen _i = _n
		
		reshape long y_, i(_i `over') j(_cats) string
		
		foreach x of  local varlist {
			replace _cats = "``x'_lab'" if _cats=="`x'"
		}		
		
	}
	if `length' == 1 {	
		
		if "`by'" == ""	{
			gen _by = 1 
			local by _by
		}		
		
		cap ren `by' _cats
		
		if "`normvar'"=="" {
			collapse (sum) `varlist', by(_cats `over')
		}
		else {
			collapse (sum) `varlist' (mean) `normvar', by(_cats `over')
		}

		cap ren `varlist' y_
	}	
	
	sort _cats `over'
	ren y_ _val

	
	egen _grp = group(_cats)
	
	
	local max = 0
	
	if `length' > 1 {
		if "`normvar'" == "" {
			levelsof _grp, local(lvls)
			
			foreach x of local lvls {
				summ _val if _grp==`x', meanonly
				if r(sum) > `max' local max = r(sum)
			}
		}
		else {
			levelsof _grp, local(lvls)
			
			foreach x of local lvls {
				summ `normvar' if _grp==`x', meanonly
				if r(sum) > `max' local max = r(sum)
			}
			
		}
	}
	if `length'==1 {
		summ `normvar', meanonly
		local max = r(max)
	}
	

	cap drop `normvar'
	
	gen double _share = .

	if "`percent'" == "" {
		replace _share = _val / `max' 
	}
	else {
		levelsof _grp, local(lvls)
		foreach x of local lvls {
			summ _val if _grp==`x', meanonly
						
			replace _share = _val / r(sum) if _grp==`x' 
		}
	}
	

	egen _tag = tag(_grp)
	
	bysort _cats: egen double _share_tot = sum(_share)
	
	
	local obsv = `rowdots' * `coldots'  // calculate the total number of rows per group
	
	expand `obsv' if _tag==1, gen(_control)	
		
	sort _grp  _cats _control `over'
	
	
	bysort _grp:	egen _y = seq(), b(`rowdots')  
	bysort _grp:	egen _x = seq(), t(`rowdots') 	

	by _grp: gen _id = _n				

	drop if _id > `obsv'
	drop _tag
	
	gen _dot = 0 
	
	egen _tag2 = tag(_grp `over')
	egen _tag3 = group(`over')
	
	levelsof _grp , local(lvls)
	levelsof _tag3, local(ovrs)
	
	local items = r(r)
	

	foreach x of local lvls {
		
		local start = 1
		local counter = 1
		
		foreach y of local ovrs {
					
			summ _share if _grp==`x' & _tag3==`y' & _tag2==1, meanonly
			local share = r(mean)
			
			if `r(N)' > 0 {
				summ _id  if _grp==`x', meanonly

				local gap = int(`share' * `r(max)')			
				local end = `start' + `gap'
				
				qui replace _dot = `counter' if _grp==`x' & inrange(_id, `start', `end') 
				local start = `end' + 1
				
			}
			local ++counter
		}
	}	
	

	cap confirm numeric var _cats
	if !_rc {
		decode _cats, gen(_temp)
	}
	else {
		gen _temp = _cats
	}
	
	
	if "`format'"  == "" {
		if "`showpct'"  == "" {
			local format %15.0fc
		}
		else {
			local format %6.2f
		}
	}		
	
	
	gen _label = ""

	if "`showpct'" == "" {
		levelsof _grp, local(lvls)
	
		foreach x of local lvls {
			summ _val if _grp==`x' & _tag2==1, meanonly
			replace _label = _temp + " (" + string(r(sum), "`format'") + ")"	if _grp==`x'
		}
	}
	else {
		replace _label = _temp + " (" + string(_share_tot * 100, "`format'") + "%)"	
	}
	

	capture drop _i 
	capture drop _control 
	capture drop _temp
	
	if "`msize'"   		== "" 	local msize		0.85	
	if "`msymbol'" 		== "" 	local msymbol	square	
	if "`mlwidth'" 		== "" 	local mlwidth	0.05	
	if "`ndsymbol'"		== "" 	local ndsymbol	square		
	if "`ndsize'"   	== "" 	local ndsize	0.5
	if "`legposition'"  == "" 	local legposition	6
	
	if "`ndcolor'" 	== "" 	{
		if "`normvar'" == "" {
			local ndcolor none
		}
		else {
			local ndcolor gs14
		}
	}
	
	if "`palette'" == "" {
		local palette tableau
	}
	else {
		tokenize "`palette'", p(",")
		local palette  `1'
		local poptions `3'
	}
	
	
	if "`legsize'"   == "" 	local legsize 2.2	
	
	// dots

	levelsof `over', local(ovlvls)
	
	local mlen : word count `msymbol' 
	local slen : word count `msize'
	local wlen : word count `mlwidth'
	

	levelsof _dot if _dot > 0, local(lvls)
	foreach x of local lvls {
		
		local b = min(`x', `mlen')
		local a : word `b' of `msymbol'
		
		local c = min(`x', `slen')
		local f : word `c' of `msize'
		
		local k = min(`x', `wlen')
		local p : word `k' of `mlwidth'		
		
		colorpalette `palette', nograph	n(`items') `poptions'
		local dots `dots' (scatter _y _x if _dot==`x',  msymbol(`a') msize(`f') mcolor("`r(p`x')'") mlwidth(`p') ) ///
	
		
		// legend
		
		capture confirm numeric variable `over'
		if !_rc {	
			local varn : label `over' `x'
			local entries `" `entries' `x'  "`varn'"  "'
		}
		else {
			local varn : word `x' of `ovlvls'
			local entries `" `entries' `x'  "`varn'"  "'
		}
			
			local mylegend legend(order("`entries'") position(`legposition') size(`legsize') col(`legcolumns')) 
		 
		
	}
	
	
	if "`ovswitch'" != "1" local myleg2 legend(position(`legposition'))
	

	if "`nolegend'" != ""  {
		local legswitch legend(off)
		local mylegend legend(off)		
		local myleg2 
	}	
	
	if "`ovswitch'" == "1"  {
		local legswitch legend(off)
		local mylegend legend(off)
	}
	
	
	if "`aspect'" == "" local aspect = `coldots' / `rowdots' // a good approximation
	if "`subtitle'" == "" local subtitle subtitle( , pos(6) size(2.5) nobox)
	
	
	// draw the graph
	
	twoway ///
		`dots' ///
		(scatter _y _x if _dot==0, msize(`ndsize') msymbol(`ndsymbol') mcolor(`ndcolor')) ///
			, ///
			ytitle("") yscale(noline) ylabel(, nogrid) ///
			xtitle("") xscale(noline) xlabel(, nogrid) ///
			by(, noiyaxes noixaxes noiytick noixtick noiylabel noixlabel `myleg2'	 ) ///
			by(_label, `title' `note' rows(`rows') cols(`cols') imargin(`margin') `legswitch' ) ///
			`subtitle' ///
			`mylegend'	///
				aspect(`aspect') `options'	
						
	
	
	*/
	}
restore
end




*************************
**** END OF PROGRAM *****
*************************