import cv2
import numpy as np

from PIL import Image
from PIL import ImageOps
from sd_bmab.base.context import Context
from sd_bmab.base.processorbase import ProcessorBase


class EdgeEnhancement(ProcessorBase):
	def __init__(self) -> None:
		super().__init__()

		self.edge_low_threadhold = 50
		self.edge_high_threadhold = 200
		self.edge_strength = 0.5

	def preprocess(self, context: Context, image: Image):
		if context.args['edge_flavor_enabled']:
			self.edge_low_threadhold = context.args['edge_low_threadhold']
			self.edge_high_threadhold = context.args['edge_high_threadhold']
			self.edge_strength = context.args['edge_strength']
		return context.args['edge_flavor_enabled']

	def process(self, context: Context, image: Image):
		context.add_generation_param('BMAB edge flavor low threadhold', self.edge_low_threadhold)
		context.add_generation_param('BMAB edge flavor high threadhold', self.edge_high_threadhold)
		context.add_generation_param('BMAB edge flavor strength', self.edge_strength)

		numpy_image = np.array(image)
		base = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
		arcanny = cv2.Canny(base, self.edge_low_threadhold, self.edge_high_threadhold)
		canny = Image.fromarray(arcanny)
		canny = ImageOps.invert(canny)

		img = image.convert('RGB')
		newdata = [(0, 0, 0) if mdata == 0 else ndata for mdata, ndata in zip(canny.getdata(), img.getdata())]
		newbase = Image.new('RGB', img.size)
		newbase.putdata(newdata)
		return Image.blend(img, newbase, alpha=self.edge_strength).convert("RGB")