import glob
import os

import cv2
import gradio as gr


def center_crop_mod64(img):
	"""Crop image to largest area that's divisible by 64"""
	h, w	= img.shape[:2]
	shortside, longside	= min(w, h), max(w, h)
	longside_crop	= (longside // 64) * 64

	left	= (w - longside_crop) // 2	if w > h else 0
	top	= (h - longside_crop) // 2	if h > w else 0
	right	= left + longside_crop	if w > h else shortside
	bottom	= top  + longside_crop	if h > w else shortside

	return img[top:bottom, left:right]

def load_example_images():
	"""Load all images from images/ directory"""
	image_dir = os.path.join(os.path.dirname(__file__), "images")
	if not os.path.exists(image_dir):
		return []

	extensions = ['*.png', '*.jpg', '*.jpeg', '*.webp']
	examples = []
	for ext in extensions:
		examples.extend(glob.glob(os.path.join(image_dir, ext)))
		examples.extend(glob.glob(os.path.join(image_dir, ext.upper())))

	return sorted(examples)

def apply_edge_detection(image_rgb, algorithm, apply_crop, convert2grayscale, params):
	"""Apply selected edge detection algorithm with specified parameters"""

	if image_rgb is None:
		return None

	# Apply cropping if requested
	if apply_crop:
		img_processed = center_crop_mod64(image_rgb)
	else:
		img_processed = image_rgb.copy()

	# Convert to grayscale if requested
	if convert2grayscale and len(img_processed.shape) == 3:
		img_input = cv2.cvtColor(cv2.cvtColor(img_processed, cv2.COLOR_RGB2GRAY), cv2.COLOR_GRAY2RGB)
	else:
		img_input = img_processed
	ret = img_input.copy()

	# set parameters


	# Detect edges
	ed = cv2.ximgproc.createEdgeDrawing()
	ed.setParams(params)
	ed.detectEdges(img_input)

	if algorithm == "Edges":
		ret = ed.getEdgeImage()
		return ret

	elif algorithm == "Lines":
		lines = ed.detectLines()

		# draw lines
		if lines is not None:
			for line in lines:
				x1, y1, x2, y2 = map(int, line[0])
				cv2.line(ret, (x1, y1), (x2, y2), (0, 255, 0), 1)

		return ret

	elif algorithm == "Ellipses":
		ellipsess = ed.detectEllipses()

		# draw ellipses
		if ellipsess is not None:
			for ellipses in ellipsess:
				ellipse = ellipses[0]
				center = (int(ellipse[0]), int(ellipse[1]))
				if ellipse[2] == 0: # Ellipse
					axes = (int(ellipse[3]), int(ellipse[4]))
					angle = ellipse[5]
					cv2.ellipse(ret, center, axes, angle, 0, 360, (0, 255, 0), 1)
				else: # Circle
					radius = int(ellipse[2])
					cv2.circle(ret, center, radius, (0, 255, 0), 1)

		return ret

	return img_input

# Create Gradio interface
with gr.Blocks(title="Edge Drawing") as app:
	examples = load_example_images()
	gr.Markdown("# Enhanced Edge Detection using [Edge Drawing](https://github.com/CihanTopal/ED_Lib) with [opencv-contrib-python](https://docs.opencv.org/4.x/d1/d1c/classcv_1_1ximgproc_1_1EdgeDrawing.html)")

	# Image row - Input and Output side by side
	with gr.Row():
		input_image = gr.Image(value=examples[0] if examples else None , label="Input Image" , type="numpy")
		output_image = gr.Image(value=None , label="Detection Result" , type="numpy")

	# Controls row - Processing Options on left, Parameters on right
	with gr.Row():
		# Processing Options
		with gr.Column():
			with gr.Group():
				gr.Markdown("### Processing Options")
				algorithm_radio	= gr.Radio(label="Detection mode", value="Edges", choices=["Edges", "Lines", "Ellipses"])

				paramterfree_checkbox	= gr.Checkbox(label	="Parameter Free Mode"	, value=True	, info="Auto-determine optimal parameters")
				convert2grayscale_checkbox	= gr.Checkbox(label	="Convert to Grayscale"	, value=False	, info="Force conversion to grayscale prior to edge detection (note that Edge Drawing has native support for color images)")
				crop_checkbox	= gr.Checkbox(label	="Apply center crop mod-64"	, value=False	, info="Crop to largest area divisible by 64 (used for Stable Diffusion)")

		# EdgeDrawing Parameters
		with gr.Column():
			with gr.Group():
				gr.Markdown("### Edge Drawing Parameters")

				# Parameter controls group (disabled by default)
				with gr.Group(visible=False) as param_group:
					with gr.Column():
						with gr.Row():
							gradient_operator = gr.Radio(choices=["PREWITT", "SOBEL", "SCHARR", "LSD"], value="PREWITT", label="Gradient Operator", info="indicates the operator used for gradient calculation. Default value is PREWITT")
							gradient_threshold	= gr.Slider(minimum=	1	, maximum=	100	, value=	20	, step=	1	, label="Gradient Threshold"	, scale=  1	, info="Threshold value of gradiential difference between pixels. Used to create gradient image. Default value is 20"	)

						with gr.Row():
							anchor_threshold	= gr.Slider(minimum=	0	, maximum=	255	, value=	0	, step=	1	, label="Anchor Threshold"	, scale=  1	, info="Threshold value used to select anchor points. Default value is 0"	)
							min_path_length	= gr.Slider(minimum=	1	, maximum=	50	, value=	10	, step=	1	, label="Min Path Length"	, scale=  1	, info="Minimum connected pixels length processed to create an edge segment. Default value is 10"	)

						with gr.Row():
							min_line_length	= gr.Slider(minimum=	-1	, maximum=	100	, value=	-1	, step=	1	, label="Min Line Length"	, scale=  1	, info="Minimum line length to detect. Default value is -1"	)
							line_fit_error_threshold	= gr.Slider(minimum=	0.1	, maximum=	10.0	, value=	1.0	, step=	0.1	, label="Line Fit Error Threshold"	, scale=  1	, info="Default value is 1.0"	)
							max_distance_between_two_lines	= gr.Slider(minimum=	1.0	, maximum=	20.0	, value=	6.0	, step=	0.1	, label="Max Distance Between Two Lines"	, scale=  1	, info="Default value is 6.0"	)

						with gr.Row():
							max_error_threshold	= gr.Slider(minimum=	0.1	, maximum=	5.0	, value=	1.3	, step=	0.1	, label="Max Error Threshold"	, scale=  1	, info="Default value is 1.3"	)
							scan_interval	= gr.Slider(minimum=	1	, maximum=	10	, value=	1	, step=	1	, label="Scan Interval"	, scale=  1	, info="Default value is 1"	)

						with gr.Row():
							sigma	= gr.Slider(minimum=	0.1	, maximum=	5.0	, value=	1.0	, step=	0.1	, label="Sigma"	, scale=  1	, info="Sigma value for internal GaussianBlur() function. Default value is 1.0"	)
							nfa_validation	= gr.Checkbox(label="NFA Validation"	, value=True	, info="Indicates if NFA (Number of False Alarms) algorithm will be used for line and ellipse validation. Default value is true")
							sum_flag	= gr.Checkbox(label="Sum Flag"	, value=True	, info="Default value is true")

	# Update function
	def update_output(
			image_rgb, algorithm, apply_crop, conv2grayscale, pf_mode,
			gradient_operator, anchor_threshold, gradient_threshold, min_path_length, min_line_length, line_fit_error_threshold, max_distance_between_two_lines, max_error_threshold, nfa_validation, scan_interval, sigma, sum_flag
		):
		operator_map = {
			"PREWITT"	: cv2.ximgproc.EdgeDrawing_PREWITT,
			"SOBEL"	: cv2.ximgproc.EdgeDrawing_SOBEL,
			"SCHARR"	: cv2.ximgproc.EdgeDrawing_SCHARR,
			"LSD"	: cv2.ximgproc.EdgeDrawing_LSD,
		}

		params = cv2.ximgproc.EdgeDrawing.Params()
		params.PFmode = pf_mode
		params.EdgeDetectionOperator = operator_map.get(gradient_operator, cv2.ximgproc.EdgeDrawing_PREWITT)
		if not pf_mode:
			params.AnchorThresholdValue	= anchor_threshold
			params.GradientThresholdValue	= gradient_threshold
			params.MinPathLength	= min_path_length
			params.MinLineLength	= min_line_length
			params.LineFitErrorThreshold	= line_fit_error_threshold
			params.MaxDistanceBetweenTwoLines	= max_distance_between_two_lines
			params.MaxErrorThreshold	= max_error_threshold
			params.NFAValidation	= nfa_validation
			params.ScanInterval	= scan_interval
			params.Sigma	= sigma
			params.SumFlag	= sum_flag

		ret = apply_edge_detection(image_rgb, algorithm, apply_crop, conv2grayscale, params)
		return ret

	# Function to toggle parameter group visibility
	def toggle_params(pf_mode):
		return gr.Group(visible=not pf_mode)

	# Set up event handlers
	inputs = [
		input_image,
		algorithm_radio,
		crop_checkbox,
		convert2grayscale_checkbox,
		paramterfree_checkbox,
		gradient_operator,
		anchor_threshold,
		gradient_threshold,
		min_path_length,
		min_line_length,
		line_fit_error_threshold,
		max_distance_between_two_lines,
		max_error_threshold,
		nfa_validation,
		scan_interval,
		sigma,
		sum_flag,
	]

	# Update output when any input changes
	for inp in inputs:
		inp.change(fn=update_output, inputs=inputs, outputs=output_image)

	# Toggle parameter group when parameter free mode changes
	paramterfree_checkbox.change(fn=toggle_params, inputs=paramterfree_checkbox, outputs=param_group)

	# Apply edge detection immediately on startup
	app.load(fn=update_output, inputs=inputs, outputs=output_image)

	# Examples
	if examples:
		gr.Examples(examples=examples, inputs=input_image, label="Example Images" )

if __name__ == "__main__":
	app.launch()