-- Module:ChoroplethMap generates colored maps of countries and other entities
-- Documentation and master version: https://en.wikipedia.org/wiki/Module:ChoroplethMap
-- Authors: User:Sophivorus
-- License: CC-BY-SA-4.0

local ChoroplethMap = {}

local lang = mw.getContentLanguage()

-- Helper function to get an argument
-- Arguments from Lua calls have priority over parent arguments from templates
local function getArg( key, default )
	local frame = mw.getCurrentFrame()
	local parent = frame:getParent()
	for k, value in pairs( parent.args ) do
		if k == key and mw.text.trim( value ) ~= '' then
			return value
		end
	end
	for k, value in pairs( frame.args ) do
		if k == key and mw.text.trim( value ) ~= '' then
			return value
		end
	end
	return default
end

-- Get geographical data from Commons
-- and make an index to optimize search
local source = getArg( 'source', 'ChoroplethMap.map' )
local sourceFeatures = mw.ext.data.get( source ).data.features
local sourceIndex = {}
for i = 0, #sourceFeatures do -- don't use ipairs because it starts at 1
	local feature = sourceFeatures[ i ]
	if feature.properties.wikidata then
		sourceIndex[ feature.properties.wikidata ] = i
	end
	if feature.properties.iso_a2 then
		sourceIndex[ feature.properties.iso_a2 ] = i
	end
	if feature.properties.name then
		sourceIndex[ feature.properties.name ] = i
	end
end

-- Helper function to get the Wikidata id of an entity
local function getID( entity )
	if not entity then
		return
	end
	entity = mw.text.trim( entity )
	if entity == '' then
		return
	end
	if string.match( entity, '^Q%d+$' ) then
		return entity
	end
	local index = sourceIndex[ entity ]
	if index then
		local feature = sourceFeatures[ index ]
		return feature.properties.wikidata
	end
	return mw.wikibase.getEntityIdForTitle( entity )
end

function ChoroplethMap.main( frame )

	-- Get the arguments
	local entities = getArg( 'entities' )
	local color = getArg( 'color', '#f00' )
	local view = getArg( 'view', 'Q2' ) -- Q2 is the Wikidata item for the Earth
	local latitude = getArg( 'latitude' )
	local longitude = getArg( 'longitude' )
	local zoom = getArg( 'zoom' )
	local width = getArg( 'width', '250' )
	local height = getArg( 'height', '200' )
	local align = getArg( 'align' )
	local frameless = getArg( 'frameless' )
	local mapstyle = getArg( 'mapstyle' )
	local alt = getArg( 'alt' )
	local caption = getArg( 'caption' )
	local legends = getArg( 'legends' )
	local logarithmicScale = getArg( 'logarithmic-scale' )
	local minOpacity = getArg( 'min-opacity', '0' )

	-- Parse the entities into a clean data table
	local data = {}
	if entities then

		-- If any color directive is an absolute value
		-- we need to figure out what the top value is
		-- before we can color any of them
		local topValue
		for entity in mw.text.gsplit( entities, '[;\n]' ) do
			entity = mw.text.trim( entity )
			local value = entity:match( '^([0-9.,]-) *: *.+$' )
			if value then
				value = lang:parseFormattedNumber( value )
				if not topValue or value > topValue then
					topValue = value
				end
			end
		end

		local opacity = .7
		local value
		local min = lang:parseFormattedNumber( minOpacity )
		for entity in mw.text.gsplit( entities, '[;\n]' ) do
			entity = mw.text.trim( entity )
			if entity ~= '' then
				local colorPart, entityPart = entity:match( '^([#0-9A-Za-z .,%%]-) *: *(.+)$' )
				if colorPart and entityPart then
					if colorPart:match( '^([0-9.,]+) ?%%$' ) then
						value = colorPart
						local plainValue = lang:parseFormattedNumber( mw.text.trim( value, '%%' ) )
						if logarithmicScale then
							opacity = min + ( math.log( math.max( plainValue, 1 ) ) / math.log( 100 ) ) * ( 1 - min )
						else
							opacity = min + ( plainValue / 100 ) * ( 1 - min )
						end
					elseif colorPart:match( '^[0-9.,]+$' ) then
						value = colorPart
						local plainValue = lang:parseFormattedNumber( value )
						if logarithmicScale then
							opacity = min + ( math.log( math.max( plainValue, 1 ) ) / math.log( topValue ) ) * ( 1 - min )
						else
							opacity = min + ( plainValue / topValue ) * ( 1 - min )
						end
					else
						color = colorPart
						value = nil
						opacity = .7
					end
					entity = entityPart
				end
				local id = getID( entity )
				if id then
					local title = mw.wikibase.getSitelink( id )
					local label = mw.wikibase.getLabel( id )
					table.insert( data, { id = id, title = title, label = label, value = value, color = color, opacity = opacity } )
				end
			end
		end
	end

	-- Make the features
	local features = {}
	local done = {}

	-- Make the view feature
	if view ~= 'auto' then
		local viewID = getID( view )
		local viewIndex = sourceIndex[ viewID ]
		local viewFeature
		if viewIndex then
			viewFeature = sourceFeatures[ viewIndex ]
			viewFeature.properties[ 'fill-opacity' ] = 0
			viewFeature.properties[ 'stroke-opacity' ] = 0
		else
			viewFeature = {
				[ 'type' ] = 'ExternalData',
				[ 'service' ] = 'geoshape',
				[ 'ids' ] = viewID,
				[ 'properties' ] = {
					[ 'fill-opacity' ] = 0,
					[ 'stroke-opacity' ] = 0
				}
			}
		end
		table.insert( features, viewFeature )
		done[ viewID ] = true
	end

	-- Make the entity features
	for _, entity in ipairs( data ) do
		if not done[ entity.id ] then
			local entityFeature
			local entityWikitext = '[[' .. entity.title .. '|' .. entity.label .. ']]' .. ( entity.value and ( '<br>' .. entity.value ) or '' )
			local entityIndex = sourceIndex[ entity.id ]
			if entityIndex then
				entityFeature = sourceFeatures[ entityIndex ]
				entityFeature.properties[ 'title' ] = entityWikitext
				entityFeature.properties[ 'fill' ] = entity.color
				entityFeature.properties[ 'fill-opacity' ] = entity.opacity
				entityFeature.properties[ 'stroke-width' ] = 0.1
			else
				entityFeature = {
					[ 'type' ] = 'ExternalData',
					[ 'service' ] = 'geoshape',
					[ 'ids' ] = entity.id,
					[ 'properties' ] = {
						[ 'title' ] = entityWikitext,
						[ 'fill' ] = entity.color,
						[ 'fill-opacity' ] = entity.opacity,
						[ 'stroke-width' ] = 0.1
					}
				}
			end
			table.insert( features, entityFeature )
			done[ entity.id ] = true
		end
	end

	local json = mw.text.jsonEncode( features )

	-- Build the automatic legends
	if legends then

		-- First look for repeated colors
		local legendsData = {}
		for _, entity in ipairs( data ) do
			local repeated
			for _, legendData in ipairs( legendsData ) do
				if legendData.color == entity.color and legendData.opacity == entity.opacity then
					repeated = legendData
					break
				end
			end
			if repeated then
				repeated.wikitext = repeated.wikitext .. ', [[' .. entity.title .. '|' .. entity.label .. ']]'
			else
				local legendData = {
					[ 'wikitext' ] = '[[' .. entity.title .. '|' .. entity.label .. ']]',
					[ 'value' ] = entity.value,
					[ 'color' ] = entity.color,
					[ 'opacity' ] = entity.opacity
				}
				table.insert( legendsData, legendData )
			end
		end

		legends = mw.html.create( 'div' )
		legends:css( 'column-width', '250px' )
		for _, legendData in ipairs( legendsData ) do
			local colorBox = mw.html.create( 'span' )
			colorBox:css( 'background', legendData.color )
			colorBox:css( 'display', 'inline-block' )
			colorBox:css( 'margin-right', '.5em' )
			colorBox:css( 'opacity', legendData.opacity )
			colorBox:css( 'vertical-align', 'middle' )
			colorBox:css( 'width', '1em' )
			colorBox:css( 'height', '1em' )

			local wikitext = legendData.wikitext
			if legendData.value then
				local message = '$1: $2'
				if lang:getCode() == 'fr' then
					message = '$1 : $2'
				end
				wikitext = mw.message.newRawMessage( message, legendData.value, wikitext ):plain()
			end

			local legend = mw.html.create( 'div' )
			legend:node( colorBox )
			legend:wikitext( wikitext )
			legends:node( legend )
		end
		legends = tostring( legends )
		if caption then
			caption = caption .. legends
		else
			caption = legends
		end
	end

	local attributes = {
		[ 'latitude' ] = latitude,
		[ 'longitude' ] = longitude,
		[ 'zoom' ] = zoom,
		[ 'width' ] = width,
		[ 'height' ] = height,
		[ 'align' ] = align,
		[ 'text' ] = caption,
		[ 'alt' ] = alt,
		[ 'frameless' ] = frameless,
		[ 'mapstyle' ] = mapstyle
	}

	return frame:extensionTag( 'mapframe', json, attributes )
end

return ChoroplethMap