Module:Sandbox/Erutuon/Climate/stats

local count = require "Module:fun".count

local function errorf(level, ...)
	if type(level) == "number" then
		return error(string.format(...), level + 1)
	else -- level is actually the format string.
		return error(string.format(level, ...), 2)
	end
end

local varargs_or_array_mt = {
	__index = {
		arr_func = function (self, arr, i, j)
			i = i or 1
			j = j or #arr
			assert(i > 0 and j > 0)
			
			if i == j then
				return arr[i]
			elseif i < j then -- |---i+++j---|
				return self.varargs_func(unpack(arr, i, j))
			else -- |+++j---i+++|
				return self.varargs_func(self.varargs_func(unpack(arr, 1, j)),
					self.varargs_func(unpack(arr, i)))
			end
		end,
	},
	__call = function (self, ...)
		if type(...) == 'table' then -- if first argument is table
			if select(2, ...) then -- if second argument
				return self:arr_func(...)
			else
				return self.varargs_func(unpack((...)))
			end
		else
			return self.varargs_func(...)
		end
	end,
}

-- Create a function that accepts two argument formats: a list of numbers, or
-- an array containing numbers followed by optional start and end indices
-- (similar to the arguments to "unpack", though here the end index can be
-- greater than the start index).
local function make_arr_and_varargs_func(varargs_func)
	return setmetatable({ varargs_func = varargs_func }, varargs_or_array_mt)
end

local function sum_varargs (...)
	local result = 0
	for i = 1, select('#', ...) do
		result = result + select(i, ...)
	end
	return result
end

local ops = {
	sum = make_arr_and_varargs_func(sum_varargs),
	mean = make_arr_and_varargs_func(
		function (...)
			return sum_varargs(...) / select('#', ...)
		end),
	min = make_arr_and_varargs_func(
		function (...)
			local min = math.huge
			local min_index
			for i = 1, select('#', ...) do
				local val = select(i, ...)
				if type(val) == "table" then
					val, i = val.value, val.index
				end
				if val < min then
					min = val
					min_index = i
				end
			end
			return { value = min, index = min_index }
		end),
	max = make_arr_and_varargs_func(
		function (...)
			local max = -math.huge
			local max_index
			for i = 1, select('#', ...) do
				local val = select(i, ...)
				if type(val) == "table" then
					val, i = val.value, val.index
				end
				if val > max then
					max = val
					max_index = i
				end
			end
			return { value = max, index = max_index }
		end),
}

local stats_mt = {}
stats_mt.ops = ops
-- Interprets a field name and then calculates, saves, and returns the desired
-- value. For instance, "Stats_object.mean" returns the mean of the values in
-- "Stats_object", "Stats_object.above_10" the number of values greater than 10.
function stats_mt:__index(key)
	if type(key) ~= "string" then
		return nil
	end
	
	-- Parse field names such as "above_10" or "below_60".
	local op, amount = key:match("(%a+)_(%d+)")
	if amount then
		amount = tonumber(amount)
		local count =
			op == "above" and count(function(val) return val > amount end, self)
			or op == "below" and count(function(val) return val < amount end, self)
			or errorf("Unrecognized operation %s", op)
		self[key] = count
		return count
	end
	
	-- Parse field names such as "mean" or "summer_min".
	local season
	season, op = key:match("^(%a-)_?(%a+)$")
	if not (season == "" or season == "summer" or season == "winter") then
		errorf("Unrecognized season %s", season)
	elseif not ops[op] then
		errorf("Unrecognized operation %s", op)
	end
	
	local result
	if season == "" then
		result = ops[op](self)
	else
		result = ops[op](self, unpack(self[season .. "_months"]))
	end
	self[key] = result
	return result
end

local seasons = { summer = { south = { 10, 3 }, north = { 4, 9 } } }
seasons.winter = {}
seasons.winter.south = { seasons.summer.south[2] + 1, seasons.summer.south[1] - 1 }
seasons.winter.north = { seasons.summer.north[2] + 1, seasons.summer.north[1] - 1 }

-- Allows "stats_mt" to be called as a function, returning a Stats object:
setmetatable(stats_mt, {
	__call = function (self, val, Southern_Hemisphere) -- constructor function
		val = val or {}
		local hemisphere = Southern_Hemisphere and "south" or "north"
		val.summer_months, val.winter_months =
			seasons.summer[hemisphere], seasons.winter[hemisphere]
		return setmetatable(val, self)
	end
})

return stats_mt