import torch


def tone_map(img, tone_mapper, exposure):
	"""
	Applies exposure compensation and tone mapping.
	Refer to the Visualizing Errors in Rendered High Dynamic Range Images
	paper for details about the formulas.

	:param img: float tensor (with NxCxHxW layout) containing nonnegative values
	:param tone_mapper: string describing the tone mapper to apply
	:param exposure: float tensor (with Nx1x1x1 layout) describing the exposure compensation factor
	"""

	device = img.device

	# Exposure compensation
	x = (2 ** exposure) * img

	# Set tone mapping coefficients depending on tone_mapper
	if tone_mapper == "reinhard":
		lum_coeff_r = 0.2126
		lum_coeff_g = 0.7152
		lum_coeff_b = 0.0722

		Y = x[:, 0:1, :, :] * lum_coeff_r + x[:, 1:2, :, :] * lum_coeff_g + x[:, 2:3, :, :] * lum_coeff_b
		return torch.clamp(torch.div(x, 1 + Y), 0.0, 1.0)

	if tone_mapper == "hable":
		# Source: https://64.github.io/tonemapping/
		A = 0.15
		B = 0.50
		C = 0.10
		D = 0.20
		E = 0.02
		F = 0.30
		k0 = A * F - A * E
		k1 = C * B * F - B * E
		k2 = 0
		k3 = A * F
		k4 = B * F
		k5 = D * F * F

		W = 11.2
		nom = k0 * torch.pow(W, torch.tensor([2.0]).to(device)) + k1 * W + k2
		denom = k3 * torch.pow(W, torch.tensor([2.0]).to(device)) + k4 * W + k5
		white_scale = torch.div(denom, nom)  # = 1 / (nom / denom)

		# Include white scale and exposure bias in rational polynomial coefficients
		k0 = 4 * k0 * white_scale
		k1 = 2 * k1 * white_scale
		k2 = k2 * white_scale
		k3 = 4 * k3
		k4 = 2 * k4
		# k5 = k5 # k5 is not changed
	else:
		# Source:  ACES approximation: https://knarkowicz.wordpress.com/2016/01/06/aces-filmic-tone-mapping-curve/
		# Include pre-exposure cancelation in constants
		k0 = 0.6 * 0.6 * 2.51
		k1 = 0.6 * 0.03
		k2 = 0
		k3 = 0.6 * 0.6 * 2.43
		k4 = 0.6 * 0.59
		k5 = 0.14

	x2 = torch.pow(x, 2)
	nom = k0 * x2 + k1 * x + k2
	denom = k3 * x2 + k4 * x + k5
	denom = torch.where(torch.isinf(denom), torch.Tensor([1.0]).to(device), denom)  # if denom is inf, then so is nom => nan. Pixel is very bright. It becomes inf here, but 1 after clamp below
	y = torch.div(nom, denom)
	return torch.clamp(y, 0.0, 1.0)