import torch


def compute_start_stop_exposures(reference, tone_mapper, tmax, tmin):
	"""
	Computes start and stop exposure for HDR-FLIP based on given tone mapper and reference image.
	Refer to the Visualizing Errors in Rendered High Dynamic Range Images
	paper for details about the formulas

	:param reference: float tensor (with NxCxHxW layout) containing reference images (nonnegative values)
	:param tone_mapper: string describing which tone mapper should be assumed
	:param tmax: float describing the t value used to find the start exposure
	:param tmin: float describing the t value used to find the stop exposure
	:return: two float tensors (with Nx1x1x1 layout) containing start and stop exposures, respectively, to use for HDR-FLIP
	"""

	device = reference.device

	if tone_mapper == "reinhard":
		k0 = 0
		k1 = 1
		k2 = 0
		k3 = 0
		k4 = 1
		k5 = 1

		x_max = tmax * k5 / (k1 - tmax * k4)
		x_min = tmin * k5 / (k1 - tmin * k4)
	elif 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

		c0 = (k1 - k4 * tmax) / (k0 - k3 * tmax)
		c1 = (k2 - k5 * tmax) / (k0 - k3 * tmax)
		x_max = - 0.5 * c0 + torch.sqrt(((torch.tensor([0.5]).to(device) * c0) ** 2) - c1)

		c0 = (k1 - k4 * tmin) / (k0 - k3 * tmin)
		c1 = (k2 - k5 * tmin) / (k0 - k3 * tmin)
		x_min = - 0.5 * c0 + torch.sqrt(((torch.tensor([0.5]).to(device) * c0) ** 2) - c1)
	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

		c0 = (k1 - k4 * tmax) / (k0 - k3 * tmax)
		c1 = (k2 - k5 * tmax) / (k0 - k3 * tmax)
		x_max = - 0.5 * c0 + torch.sqrt(((torch.tensor([0.5]).to(device) * c0) ** 2) - c1)

		c0 = (k1 - k4 * tmin) / (k0 - k3 * tmin)
		c1 = (k2 - k5 * tmin) / (k0 - k3 * tmin)
		x_min = - 0.5 * c0 + torch.sqrt(((torch.tensor([0.5]).to(device) * c0) ** 2) - c1)

	# Convert reference to luminance
	lum_coeff_r = 0.2126
	lum_coeff_g = 0.7152
	lum_coeff_b = 0.0722
	Y_reference = reference[:, 0:1, :, :] * lum_coeff_r + reference[:, 1:2, :, :] * lum_coeff_g + reference[:, 2:3, :, :] * lum_coeff_b

	# Compute start exposure
	Y_hi = torch.amax(Y_reference, dim=(2, 3), keepdim=True)
	start_exposure = torch.log2(torch.maximum(1e-8 * torch.ones_like(x_max), x_max / Y_hi))

	# Compute stop exposure
	dim = Y_reference.size()
	Y_ref = Y_reference.view(dim[0], dim[1], dim[2]*dim[3])
	Y_lo = torch.median(Y_ref, dim=2).values.unsqueeze(2).unsqueeze(3)
	stop_exposure = torch.log2(torch.maximum(1e-8 * torch.ones_like(x_min), x_min / Y_lo))

	return start_exposure, stop_exposure