Editing Module:USN fleet totals/data

require('strict');

local patterns_tags = {
	'<nowiki>.-</nowiki>',
	'<!%-%-.-%-%->',
	'<pre>.-</pre>',
	'<syntaxhighlight.->.-</syntaxhighlight>',
	'<source.->.-</source>',													-- deprecated alias of syntaxhighlight tag
	}

local Article_content;
local grand_total = 0;
local planned_total = 0;

local wikitables_t = {
	['Commissioned_t'] = {},													-- k/v table where k is ship type and v is the number of that type
	['Non-commissioned_t'] = {},
	['Support_t'] = {},
	['Ready Reserve Force ships_t'] = {},
	['Reserve fleet_t'] = {},
	['Under construction_t'] = {},
	['On order_t'] = {},
	}

local wikitable_names_t = {'Commissioned_t', 'Non-commissioned_t', 'Support_t', 'Ready Reserve Force ships_t', 'Reserve fleet_t', 'Under construction_t', 'On order_t'};
local fleet_wikitables_t = {'Commissioned_t', 'Non-commissioned_t', 'Support_t', 'Ready Reserve Force ships_t', 'Reserve fleet_t'};	-- only these for n_of_m
local headers_t = {
	'Commissioned (USS)',
	'Non-commissioned (USNS)',
	'Support (MV, RV – \'\'<small>or no prefix</small>\'\')',
	'Ready Reserve Force ships (MV, SS, GTS)',
	'Reserve Fleet ships (USS, USNS)',
	'Under construction',
	'On order'
	}

local totals_t = {																-- table to hold total number of ships in these categories
	['Commissioned_t'] = 0,
	['Non-commissioned_t'] = 0,
	['Support_t'] = 0,
	['Ready Reserve Force ships_t'] = 0,
	['Reserve fleet_t'] = 0,
	['Under construction_t'] = 0,
	['On order_t'] = 0,
	}

local n_of_m_t = {}																-- table of total counts (m in 'n of m')


--[[--------------------------< A R T I C L E _ C O N T E N T _ G E T >----------------------------------------

get article content, remove certain html-like tags and their content so that this code doesn't include any citation
templates inside the tags as valid tagets; they are not.

]]

local function article_content_get ()
	if not Article_content then
		Article_content = mw.title.new ('List of current ships of the United States Navy'):getContent();
		for _, tag in ipairs (patterns_tags) do
			Article_content = Article_content:gsub (tag, '');					-- remove certain html-like tags and their content
		end
		while Article_content:match ('([\r\n]+|%-) *[\r\n]+|%-') do				-- are there multiple row markers without intervening column data?
			Article_content = Article_content:gsub ('([\r\n]+|%-) *[\r\n]+|%-', '%1');	-- remove duplicate row markers
		end
	end
end


--[[--------------------------< T Y P E _ L A B E L _ I N D E X _ G E T >--------------------------------------

get the labels from the tops of the current wikitable; the one we need is '!Type'.  This function assumes that
label markup for each table is one-line-per-label; label markup all-on-one-line (! <label> !! <label> !! ...) is
not supported.

]]

local function type_label_index_get (wikitable)
	local i = 0;
	for label in wikitable:gmatch ('[\r\n](![^\r\n]+)') do
		i = i + 1;
		if label:match ('! *Type') then
			return i;
		end
	end
end


--[[--------------------------< S H I P _ T Y P E _ C L E A N >------------------------------------------------

remove templates from <ship_type>

]]

local function ship_type_clean (ship_type)
	while ship_type:match ('%b{}') do											-- are there templates?
		ship_type = ship_type:gsub ('%b{}', '');								-- remove the template
	end
	
	return ship_type;
end


--[[--------------------------< S H I P _ T Y P E S _ G E T >--------------------------------------------------

find first row markup (|-) and fetch the column number that holds ship_types.  Then find each row markup, count
lines until ship_type line is located, extract the contents (ship type) and add entry to the appropriate wikitable
in wikitables_t; ship type already present, bump the count.

]]

local function ship_types_get (wikitable, wikitable_index)
	local find_pattern = '|%-';
	local type_label_index;
	
	local tstart, tend = wikitable:find (find_pattern);							-- find the table row marker in <wikitable>; should be column headers
	local i=0;

	if tstart then																-- if we found row marker (|-) for headers
		type_label_index = type_label_index_get (wikitable);					-- try to find index of type label (!Type)
		tstart, tend = wikitable:find (find_pattern, tend);						-- look for the next row
	else
		return nil;																-- TODO: error message?
	end

	while tstart do
		local i = 0;
		local ship_type;
		local rstart = tend;													-- row start and end return values from string.find()
		local rend;
		local pattern = '[\r\n]+|([^\r\n]*)';									-- newline, pipe, everything before the next newline

		while rstart do															-- nil if not found
			rstart, rend, ship_type = wikitable:find (pattern, rstart);

			if rstart then
				i = i + 1;
				ship_type = ship_type_clean (ship_type);
				
				if type_label_index == i then									-- is this the column that has the ship type entry?
					if wikitables_t[wikitable_names_t[wikitable_index]][ship_type] then		-- if we found a ship type that we already know about:
						wikitables_t[wikitable_names_t[wikitable_index]][ship_type] = wikitables_t[wikitable_names_t[wikitable_index]][ship_type] + 1;	-- bump the count
					else
						wikitables_t[wikitable_names_t[wikitable_index]][ship_type] = 1;	-- add ship type to the table with a count of 1
					end
					
					totals_t[wikitable_names_t[wikitable_index]] = totals_t[wikitable_names_t[wikitable_index]] + 1;	-- tally
					break;														-- and break out of the while
				end
			rstart = rend;														-- not yet on the correct column; reset the starting index to the end of the last find()
			end
		end

	tstart, tend = wikitable:find (find_pattern, tend);							-- search for another row marker (|-) in <wikitable>; begin at end of last row marker search
	end
end


--[[--------------------------< R E N D E R _ O U T P U T >----------------------------------------------------

extract data from various tables and append to <output_string>

]]

local function render_output (index, output_string)
	local out_t = {};															-- here we compose this section of <output_string>
	local temp_t = {};															-- a sequence of ship types and counts taken from wikitables_t; separate table for sorting
	
	local function sort (a, b)
		a = a:gsub ('%[%[(.+)%]%]', '%1');										-- remove outer wikilink markup
		b = b:gsub ('%[%[(.+)%]%]', '%1');
		a = a:gsub ('%* ', '');													-- remove unordered list markup
		b = b:gsub ('%* ', '');
		a = a:gsub ('^[^|]+|', '');												-- remove all but the display text from complex wikilinks
		b = b:gsub ('^[^|]+|', '');
		a = a:gsub (' *– *%d+', '');											-- remove the count
		b = b:gsub (' *– *%d+', '');
		return a < b;
	end
	
	table.insert (out_t, table.concat ({'<b>', headers_t[index], '</b> – ', totals_t[wikitable_names_t[index]]}));		-- header
	table.insert (out_t, '{{Div col|colwidth=22em}}');							-- start of columnar data
	for ship_type, count in pairs (wikitables_t[wikitable_names_t[index]]) do	-- make an unordered, unsorted, list
		local ship_type_total = n_of_m_t[ship_type] and table.concat ({' (<i>of ', n_of_m_t[ship_type], '</i>)'}) or '';
		if wikitable_names_t[index] == fleet_wikitables_t[index] then
			table.insert (temp_t, table.concat ({'* ', ship_type, ' – ', count, ship_type_total}));	-- make and add an unordered list item
		else
			table.insert (temp_t, table.concat ({'* ', ship_type, ' – ', count}));	-- make and add an unordered list item
		end
	end
	table.sort (temp_t, sort);													-- ascending sort; TODO: is there a better way to get a sorted list of ship types?
	table.insert (out_t, table.concat (temp_t, '\n'));							-- make a big string and add it to out_t
	table.insert (out_t, '{{div col end}}');									-- end of columnar data
	table.insert (out_t, '<hr />\n');											-- horizontal rule
	
	output_string = table.concat ({output_string, table.concat (out_t, '\n')});	-- make a big string and add it to output_string
	return output_string;														-- and done
end


--[[--------------------------< R E N D E R _ T O T A L S >----------------------------------------------------

render a wikitable list of totals; exclude Under construction and on-order ships

]]

local function render_totals (output_string)
	local out_t = {'<b>Totals</b>'};
	local tally = 0;

	table.insert (out_t, '{| style="margin-left:1.5em"');						-- start a wikitable indented 1.5em
	for i, wikitable_name in ipairs (wikitable_names_t) do						-- spin through the list of table names
		if 6 > i then															-- not under construction of on order
			local wikitable = wikitable_name:gsub ('_t', '');					-- remove suffix
			table.insert (out_t, '|-');											-- add row markup
			table.insert (out_t, table.concat ({'| ', wikitable, ': || style="text-align:right" | ', totals_t[wikitable_names_t[i]]}));	-- make a row
			tally = tally + totals_t[wikitable_names_t[i]];						-- and bump the grand total
		else
			break;
		end
	end
	grand_total = tally;														-- this will be rounded to the nearest integer evenly divisible by 5
	table.insert (out_t, '|-');														-- add row markup
	table.insert (out_t, table.concat ({'| <b>Grand total</b>: || style="text-align:right" | <b>', tally, '</b>'}));	-- add the grand total row
	table.insert (out_t, '|}');													-- close the wikitable

	output_string = table.concat ({output_string, table.concat (out_t, '\n')});	-- make a big string and add it to output_string
	return output_string;														-- and done
end


--[[--------------------------< N _ O F _ M _ G E T >----------------------------------------------------------

For ship types in commissioned, non-commissioned, support, ready researve, and reserve tables, count the number
of same-type ships so that we can render '(''of M'')' annotation for those ship types.

First gather a list of ship types that are listed in more than one table.  Compare commissioned to non-commissioned;
commissioned to support; etc.  Then non-commissioned to support; non-commissioned to ready researve; etc.  Continue
until we compare ready researve to reserve.  Do not compare reserve to itself.

Second, scroll through the list of duplicates and accumulate tallies of each ship type.

]]

local function n_of_m_get ()
	for i, wikitable_name in ipairs (fleet_wikitables_t) do						-- for each wikitable
		if #fleet_wikitables_t == i then										-- when we get to the last wikitable, don't compare it to itself
			break;
		end
		for ship_type, _ in pairs (wikitables_t[wikitable_name]) do
			for j=i+1, #fleet_wikitables_t do									-- index 1 looks in indexes 2, 3, ...
				if wikitables_t[fleet_wikitables_t[j]][ship_type] then
					n_of_m_t[ship_type] = 0;									-- this ship type appears in more than one wikitable
				end
			end
		end
	end

	for ship_type, _ in pairs (n_of_m_t) do										-- for each ship type in n_of_m table
		for _, wikitable_name in ipairs (fleet_wikitables_t) do					-- for each wikitable
			if wikitables_t[wikitable_name][ship_type] then						-- if the ship type is found in this wikitable
				n_of_m_t[ship_type] = n_of_m_t[ship_type] + wikitables_t[wikitable_name][ship_type];	-- add to the tally
			end
		end
	end
end


--[[--------------------------< U S N _ S H I P _ C O U N T E R >----------------------------------------------

tables in 'List of current ships of the United States Navy' are (in this order):
	Commissioned
	Non-commissioned
	Support
	Ready Reserve Force ships
	Reserve fleet
	Under construction
	On order

count the ship types in the Type columns of these tables, and then render pretty sorted lists of ship types with
their counts.

]]

local function fleet_totals ()
	article_content_get ();														-- attempt to get this article's content

	local wikitable_index = 0;
	local find_pattern = '{|';
	local tstart, tend = Article_content:find (find_pattern);					-- find the first wikitable

	while tstart do
		local wikitable = Article_content:match ('%b{}', tstart);				-- get the whole wikitable

		if not wikitable then
			break;																-- wikitable is nil for some reason (missing closing |} for example) so declare ourselves done
		end

		wikitable_index = wikitable_index + 1;
		ship_types_get (wikitable, wikitable_index);

		tstart, tend = Article_content:find (find_pattern, tend);				-- search for another template; begin at end of last wikitable search
	end

	n_of_m_get ();

	local output_string = '';
	for i, _ in ipairs (wikitable_names_t) do
		output_string = render_output (i, output_string);						-- make a pretty unordered list of ship types and counts for each wikitable
	end

	output_string = render_totals (output_string);								-- make a prety unordered list of 'classification' total numbers

	local frame = mw.getCurrentFrame();											-- do this to get access to frame:preprocess()
	return frame:preprocess (output_string);
end


--[[--------------------------< E X P O R T S >----------------------------------------------------------------
]]

return {
	fleet_totals_str = fleet_totals(),
	grand_total = grand_total,
	planned_total = totals_t['Under construction_t'] + totals_t['On order_t'],
	}