-- Complex numbers with metatable-based operators

local intpow = require("math.intpow")

local complex = {}

local metatable = { __index = complex }

function complex.new(
	real, --[[number, real part]]
	imaginary --[[number, imaginary part, defaults to 0]]
)
	return setmetatable({ real = real, imaginary = imaginary or 0 }, metatable)
end

-- Imaginary unit as a complex number
complex.i = complex.new(0, 1)

function complex.from_polar_coordinates(
	angle, --[[number, angle in radians]]
	abs --[[number, radius / norm / length]]
)
	return complex.new(abs * math.cos(angle), abs * math.sin(angle))
end

local function bin_op(name, operator)
	metatable["__" .. name] = function(self, other)
		-- Treat numbers as complex numbers with an imaginary part of 0
		if type(self) == "number" then
			self = complex.new(self)
		elseif type(other) == "number" then
			other = complex.new(other)
		end
		return operator(self, other)
	end
end

-- Arithmetic binary operators

metatable.__pow = intpow

bin_op("add", function(self, other)
	return complex.new(self.real + other.real, self.imaginary + other.imaginary)
end)

bin_op("sub", function(self, other)
	return complex.new(self.real - other.real, self.imaginary - other.imaginary)
end)

bin_op("mul", function(self, other)
	return complex.new(
		self.real * other.real - self.imaginary * other.imaginary,
		self.real * other.imaginary + self.imaginary * other.real
	)
end)

bin_op("div", function(self, other)
	local numerator = self * other:conjugate()
	local denominator = other.real ^ 2 + other.imaginary ^ 2
	return complex.new(numerator.real / denominator, numerator.imaginary / denominator)
end)

-- Unary operators

function metatable:__unm()
	return complex.new(-self.real, -self.imaginary)
end

function complex:conjugate()
	return complex.new(self.real, -self.imaginary)
end

function complex:abs()
	return (self.real ^ 2 + self.imaginary ^ 2) ^ 0.5
end

-- Comparison operators; only equality is reasonable here

bin_op("eq", function(self, other)
	return self.real == other.real and self.imaginary == other.imaginary
end)

-- Conversions

function complex:to_polar_coordinates()
	local angle, abs = math.atan2(self.imaginary, self.real), self:abs()
	return angle --[[number, angle in radians]],
		abs --[[number, radius / norm / length]]
end

function metatable:__tostring()
	return self.real .. "+" .. self.imaginary .. "i"
end

return complex