import math

from PIL import Image
from PIL import ImageEnhance

from sd_bmab import util
from sd_bmab.base.context import Context
from sd_bmab.base.processorbase import ProcessorBase


def calc_color_temperature(temp):
	white = (255.0, 254.11008387561782, 250.0419083427406)

	temperature = temp / 100

	if temperature <= 66:
		red = 255.0
	else:
		red = float(temperature - 60)
		red = 329.698727446 * math.pow(red, -0.1332047592)
		if red < 0:
			red = 0
		if red > 255:
			red = 255

	if temperature <= 66:
		green = temperature
		green = 99.4708025861 * math.log(green) - 161.1195681661
	else:
		green = float(temperature - 60)
		green = 288.1221695283 * math.pow(green, -0.0755148492)
	if green < 0:
		green = 0
	if green > 255:
		green = 255

	if temperature >= 66:
		blue = 255.0
	else:
		if temperature <= 19:
			blue = 0.0
		else:
			blue = float(temperature - 10)
			blue = 138.5177312231 * math.log(blue) - 305.0447927307
			if blue < 0:
				blue = 0
			if blue > 255:
				blue = 255

	return red / white[0], green / white[1], blue / white[2]


class FinalProcessorBasic(ProcessorBase):
	def __init__(self) -> None:
		super().__init__()
		self.noise_alpha_final = 0
		self.contrast = 1
		self.brightness = 1
		self.sharpeness = 1
		self.color_saturation = 1
		self.color_temperature = 0

	def preprocess(self, context: Context, image: Image):
		self.contrast = context.args['contrast']
		self.brightness = context.args['brightness']
		self.sharpeness = context.args['sharpeness']
		self.color_saturation = context.args['color_saturation']
		self.color_temperature = context.args['color_temperature']
		self.noise_alpha_final = context.args['noise_alpha_final']
		return True

	def process(self, context: Context, image: Image):

		if self.noise_alpha_final != 0:
			context.add_generation_param('BMAB noise alpha final', self.noise_alpha_final)
			img_noise = util.generate_noise(context.sdprocessing.seed, image.size[0], image.size[1])
			image = Image.blend(image, img_noise, alpha=self.noise_alpha_final)

		if self.contrast != 1:
			context.add_generation_param('BMAB contrast', self.contrast)
			enhancer = ImageEnhance.Contrast(image)
			image = enhancer.enhance(self.contrast)

		if self.brightness != 1:
			context.add_generation_param('BMAB brightness', self.brightness)
			enhancer = ImageEnhance.Brightness(image)
			image = enhancer.enhance(self.brightness)

		if self.sharpeness != 1:
			context.add_generation_param('BMAB sharpeness', self.sharpeness)
			enhancer = ImageEnhance.Sharpness(image)
			image = enhancer.enhance(self.sharpeness)

		if self.color_saturation != 1:
			context.add_generation_param('BMAB color', self.color_saturation)
			enhancer = ImageEnhance.Color(image)
			image = enhancer.enhance(self.color_saturation)

		if self.color_temperature != 0:
			context.add_generation_param('BMAB color temperature', self.color_temperature)
			temp = calc_color_temperature(6500 + self.color_temperature)
			az = []
			for d in image.getdata():
				az.append((int(d[0] * temp[0]), int(d[1] * temp[1]), int(d[2] * temp[2])))
			image = Image.new('RGB', image.size)
			image.putdata(az)

		return image

	def postprocess(self, context: Context, image: Image):
		pass