import math

import cv2
import numpy as np

from .algorithms import (
    ProcessingError,
    bit_plane,
    box_filter,
    contrast_stretch,
    gamma,
    gray_slice,
    gray_to_rgb,
    histogram_equalization,
    logarithmic,
    negative,
    normalize_to_uint8,
    to_gray,
    weighted_average,
    median_filter,
    high_boost,
    ensure_uint8,
)


CH_BASIC = "Basic"
CH3 = "Image Enhancement in the Spatial Domain"
CH4 = "Image Enhancement in the Frequency Domain"
CH6 = "Color Image Processing"


def with_meta(schema, *, label=None, description=None, show_when=None):
    if label:
        schema["label"] = label
    if description:
        schema["description"] = description
    if show_when:
        schema["show_when"] = show_when
    return schema


def odd_param(default=3, max_value=35, **meta):
    return with_meta({"type": "int", "default": default, "min": 3, "max": max_value, "step": 2, "odd": True}, **meta)


def float_param(default, min_value, max_value, step=0.1, **meta):
    return with_meta({"type": "float", "default": default, "min": min_value, "max": max_value, "step": step}, **meta)


def int_param(default, min_value, max_value, step=1, **meta):
    return with_meta({"type": "int", "default": default, "min": min_value, "max": max_value, "step": step}, **meta)


def select_param(default, choices, **meta):
    return with_meta({"type": "select", "default": default, "choices": choices}, **meta)


def bool_param(default=False, **meta):
    return with_meta({"type": "bool", "default": default}, **meta)


def kernel_preview(title, matrix, scale=None):
    return {"title": title, "matrix": matrix, "scale": scale}


def kernel_pair_preview(title, gx, gy):
    return {"title": title, "kernels": [{"label": "Gx", "matrix": gx}, {"label": "Gy", "matrix": gy}]}


def mask_param(default=3, max_value=35, label="Mask size"):
    schema = odd_param(default, max_value)
    schema["label"] = label
    return schema


def histeq(image, params):
    if image.ndim == 2:
        return histogram_equalization(image, params)
    channels = [histogram_equalization(image[:, :, idx], params)[:, :, 0] for idx in range(3)]
    return np.stack(channels, axis=2).astype(np.uint8)


def rgb_to_gray_matlab(image, params):
    red_weight = float(params.get("red_weight", 0.299))
    green_weight = float(params.get("green_weight", 0.587))
    blue_weight = float(params.get("blue_weight", 0.114))
    total = red_weight + green_weight + blue_weight
    if not np.isfinite(total) or math.isclose(total, 0.0):
        raise ProcessingError("Grayscale weights must have a non-zero finite sum.")
    weights = np.array([red_weight, green_weight, blue_weight], dtype=np.float32) / total
    gray = np.tensordot(image.astype(np.float32), weights, axes=([2], [0]))
    return gray_to_rgb(ensure_uint8(np.round(gray)))


def gaussian_noise(image, params):
    mean = float(params.get("mean", 0))
    variance = float(params.get("variance", 0.01))
    sigma = math.sqrt(max(variance, 0.0)) * 255.0
    noise = np.random.default_rng().normal(mean * 255.0, sigma, size=image.shape)
    return ensure_uint8(np.round(image.astype(np.float32) + noise))


def salt_pepper_noise(image, params):
    amount = float(params.get("amount", 0.03))
    salt_ratio = float(params.get("salt_ratio", 0.5))
    output = image.copy()
    rng = np.random.default_rng()
    mask = rng.random(image.shape[:2])
    salt = mask < amount * salt_ratio
    pepper = (mask >= amount * salt_ratio) & (mask < amount)
    output[salt] = 255
    output[pepper] = 0
    return output


def noise_filter(image, params):
    kind = params.get("kind", "gaussian")
    if kind == "salt_pepper":
        return salt_pepper_noise(image, params)
    return gaussian_noise(image, params)


def gaussian_filter(image, params):
    size = int(params.get("K", params.get("size", 3)))
    variance = float(params.get("Q", params.get("variance", 1.0)))
    if size < 3 or size % 2 == 0:
        raise ProcessingError("K must be an odd integer >= 3.")
    sigma = math.sqrt(max(variance, 1e-8))
    return cv2.GaussianBlur(image, (size, size), sigmaX=sigma, sigmaY=sigma, borderType=cv2.BORDER_REFLECT)


def max_filter(image, params):
    size = int(params.get("mask_size", params.get("N", params.get("size", 3))))
    if size < 3 or size % 2 == 0:
        raise ProcessingError("Mask size must be an odd integer >= 3.")
    return cv2.dilate(image, np.ones((size, size), np.uint8))


def min_filter(image, params):
    size = int(params.get("mask_size", params.get("N", params.get("size", 3))))
    if size < 3 or size % 2 == 0:
        raise ProcessingError("Mask size must be an odd integer >= 3.")
    return cv2.erode(image, np.ones((size, size), np.uint8))


def box_denoise(image, params):
    return box_filter(image, {"size": params.get("K", params.get("mask_size", params.get("size", 3)))})


def weighted_denoise(image, params):
    return weighted_average(image, {"size": 3})


def median_denoise(image, params):
    return median_filter(image, {"size": params.get("mask_size", params.get("N", params.get("size", 3)))})


def gaussian_denoise(image, params):
    size = params.get("K", params.get("mask_size", params.get("size", 3)))
    variance = params.get("Q", params.get("variance", 1.0))
    return gaussian_filter(image, {"K": size, "Q": variance})


def high_boost_slide(image, params):
    return high_boost(image, {"amplification": params.get("A", params.get("amplification", 1.5)), "size": params.get("K", params.get("size", 3))})


def laplacian_slide(image, params):
    mask_name = params.get("mask", "cross")
    kernels = {
        "cross": np.array([[0, 1, 0], [1, -5, 1], [0, 1, 0]], dtype=np.float32),
        "diagonal": np.array([[1, 1, 1], [1, -9, 1], [1, 1, 1]], dtype=np.float32),
    }
    kernel = kernels.get(mask_name)
    if kernel is None:
        raise ProcessingError("Unknown Laplacian mask.")
    channels = [cv2.filter2D(image[:, :, idx], cv2.CV_32F, kernel, borderType=cv2.BORDER_REFLECT) for idx in range(image.shape[2])]
    result = np.stack(channels, axis=2)
    return ensure_uint8(result)


def gradient_abs_sum(image, params):
    operator = params.get("operator", "sobel")
    gray = to_gray(image).astype(np.float32)
    if operator == "roberts":
        gx = np.array([[-1, 0], [0, 1]], dtype=np.float32)
        gy = np.array([[0, -1], [1, 0]], dtype=np.float32)
    else:
        gx = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=np.float32)
        gy = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=np.float32)
    fx = cv2.filter2D(gray, cv2.CV_32F, gx, borderType=cv2.BORDER_REFLECT)
    fy = cv2.filter2D(gray, cv2.CV_32F, gy, borderType=cv2.BORDER_REFLECT)
    return gray_to_rgb(normalize_to_uint8(np.abs(fx) + np.abs(fy)))


def rgb_channel(image, params):
    channel = params.get("channel", "r")
    index = {"r": 0, "g": 1, "b": 2}.get(channel, 0)
    return gray_to_rgb(image[:, :, index])


def fft_spectrum(image, params):
    gray = to_gray(image).astype(np.float32)
    spectrum = np.fft.fftshift(np.fft.fft2(gray))
    mode = params.get("mode", "log_magnitude")
    if mode == "phase":
        return gray_to_rgb(normalize_to_uint8(np.angle(spectrum)))
    magnitude = np.abs(spectrum)
    if mode == "log_magnitude":
        magnitude = np.log1p(magnitude)
    return gray_to_rgb(normalize_to_uint8(magnitude))


def operation(id, label, chapter, slide_group, func, params=None, supports="both", matrices=None, formula="", repeatable=True):
    return {
        "id": id,
        "label": label,
        "chapter": chapter,
        "slide_group": slide_group,
        "params": params or {},
        "supports": supports,
        "matrices": matrices or [],
        "formula": formula,
        "repeatable": repeatable,
        "func": func,
    }


OPERATIONS = [
    operation("histeq", "Histogram Equalization", CH_BASIC, "Histogram", histeq, formula="Default histeq: map gray levels by the cumulative histogram CDF.", repeatable=False),
    operation("negative", "Negative", CH3, "Point Processing", negative, formula="s = 255 - r", repeatable=False),
    operation("log", "Log", CH3, "Point Processing", logarithmic, {"c": float_param(1.44, 0.1, 5, 0.05, description="Scale factor in s = c log(1 + r).")}, formula="s = c log(1 + r)", repeatable=False),
    operation("gamma", "Power-Law / Gamma", CH3, "Point Processing", gamma, {"gamma": float_param(1.0, 0.1, 5, 0.05, description="Exponent gamma in s = c r^gamma."), "c": float_param(1.0, 0.1, 3, 0.05, description="Scale factor c in s = c r^gamma.")}, formula="s = c r^gamma", repeatable=False),
    operation("contrast_stretch", "Gray-Level Dynamic Range", CH3, "Piecewise Linear", contrast_stretch, {"low": int_param(30, 0, 254, description="Input gray level mapped toward 0."), "high": int_param(220, 1, 255, description="Input gray level mapped toward 255.")}, formula="Stretch [low, high] to [0, 255].", repeatable=False),
    operation("gray_slice", "Gray-Level Slicing", CH3, "Piecewise Linear", gray_slice, {"start": int_param(96, 0, 255, description="Lower bound A of highlighted range [A, B]."), "end": int_param(160, 0, 255, description="Upper bound B of highlighted range [A, B]."), "preserve_background": bool_param(True, description="Keep original background outside [A, B].")}, formula="Highlight gray range A <= r <= B.", repeatable=False),
    operation("bit_plane", "Bit-Plane Slicing", CH3, "Piecewise Linear", bit_plane, {"bit": int_param(7, 0, 7, description="Bit plane index, 0 least significant through 7 most significant.")}, formula="Output bit k of each gray level.", repeatable=False),
    operation("noise_filter", "Noise Filter", CH3, "Noise and Denoising", noise_filter, {
        "kind": select_param("gaussian", ["gaussian", "salt_pepper"], description="Select the noise model."),
        "mean": float_param(0, -1, 1, 0.01, description="Gaussian mean in normalized intensity units.", show_when={"param": "kind", "value": "gaussian"}),
        "variance": float_param(0.01, 0, 0.2, 0.005, description="Gaussian variance; sigma = sqrt(variance).", show_when={"param": "kind", "value": "gaussian"}),
        "amount": float_param(0.03, 0, 0.5, 0.01, description="Salt-and-pepper probability per pixel.", show_when={"param": "kind", "value": "salt_pepper"}),
        "salt_ratio": float_param(0.5, 0, 1, 0.05, description="Fraction of impulse noise assigned to salt.", show_when={"param": "kind", "value": "salt_pepper"}),
    }, formula="Gaussian: g=f+n. Salt-pepper: pixels become 0 or 255."),
    operation("box_filter", "Average / Box Filter", CH3, "Linear Low-Pass Filters", box_denoise, {"K": odd_param(3, 35, description="Odd mask dimension K for the K x K average mask.")}, matrices=[kernel_preview("1 / K^2 box mask", [["1", "1", "1"], ["1", "1", "1"], ["1", "1", "1"]], "1 / K^2")], formula="g = imfilter(f, ones(K,K)/K^2)"),
    operation("weighted_average", "Weighted Average Filter", CH3, "Linear Low-Pass Filters", weighted_denoise, matrices=[kernel_preview("Weighted average mask", [[1, 2, 1], [2, 4, 2], [1, 2, 1]], "1 / 16")], formula="g = imfilter(f, weighted mask)"),
    operation("gaussian_filter", "Gaussian Filter", CH3, "Linear Low-Pass Filters", gaussian_denoise, {"K": odd_param(3, 35, description="Odd Gaussian mask dimension K."), "Q": float_param(1.0, 0.01, 25, 0.1, description="Variance Q of the Gaussian mask.")}, formula="Gaussian mask controlled by K and variance Q."),
    operation("median_filter", "Median Filter", CH3, "Order-Statistics Filters", median_denoise, {"mask_size": mask_param(3, 25, "Window size")}, formula="g(x,y) = median of neighborhood."),
    operation("max_filter", "Max Filter", CH3, "Order-Statistics Filters", max_filter, {"mask_size": mask_param(3, 25, "Window size")}, formula="g(x,y) = max of neighborhood."),
    operation("min_filter", "Min Filter", CH3, "Order-Statistics Filters", min_filter, {"mask_size": mask_param(3, 25, "Window size")}, formula="g(x,y) = min of neighborhood."),
    operation("laplacian_slide", "Laplacian Sharpening Masks", CH3, "Sharpening Spatial Filters", laplacian_slide, {"mask": select_param("cross", ["cross", "diagonal"], description="Choose one of the taught sharpening masks.")}, matrices=[
        kernel_preview("Sharpening cross mask", [[0, 1, 0], [1, -5, 1], [0, 1, 0]]),
        kernel_preview("Sharpening diagonal mask", [[1, 1, 1], [1, -9, 1], [1, 1, 1]]),
    ], formula="Sharpen with selected Laplacian mask."),
    operation("gradient_abs_sum", "Gradient Operators", CH3, "Gradient Operator", gradient_abs_sum, {"operator": select_param("sobel", ["sobel", "roberts"], description="Choose Gx/Gy pair.")}, matrices=[
        kernel_pair_preview("Roberts Cross-Gradient", [[-1, 0], [0, 1]], [[0, -1], [1, 0]]),
        kernel_pair_preview("Sobel", [[-1, -2, -1], [0, 0, 0], [1, 2, 1]], [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]),
    ], formula="Gradient image = abs(imfilter(f,Gx)) + abs(imfilter(f,Gy))."),
    operation("high_boost", "High-Boost / Edge Emphasis", CH3, "High-Boost Filtering", high_boost_slide, {"A": float_param(1.5, 1, 6, 0.1, description="Boost factor A, where A >= 1."), "K": odd_param(3, 35, description="Odd averaging mask size used for the blurred image.")}, formula="f_hb = A f - blurred(f)."),
    operation("fft_spectrum", "FFT/DFT Spectrum View", CH4, "DFT and FFT", fft_spectrum, {"mode": select_param("log_magnitude", ["magnitude", "log_magnitude", "phase"], description="Choose magnitude, log magnitude, or phase display.")}, formula="F(u,v) = DFT{f(x,y)}", repeatable=False),
    operation("rgb_to_gray", "Convert to Grayscale", CH6, "Color Conversion", rgb_to_gray_matlab, {"red_weight": float_param(0.299, 0, 1, 0.001, description="R coefficient in gray = aR + bG + cB."), "green_weight": float_param(0.587, 0, 1, 0.001, description="G coefficient in gray = aR + bG + cB."), "blue_weight": float_param(0.114, 0, 1, 0.001, description="B coefficient in gray = aR + bG + cB.")}, formula="gray = 0.299R + 0.587G + 0.114B by default.", repeatable=False),
    operation("rgb_channel", "RGB Channel View", CH6, "RGB color model", rgb_channel, {"channel": select_param("r", ["r", "g", "b"], description="Select the RGB channel to view.")}, formula="Show one RGB channel as grayscale.", repeatable=False),
]

OPERATION_MAP = {item["id"]: item for item in OPERATIONS}


def operation_metadata():
    return [{key: value for key, value in item.items() if key != "func"} for item in OPERATIONS]


def apply_registered_operation(image, operation_id, params=None):
    item = OPERATION_MAP.get(operation_id)
    if item is None:
        raise ProcessingError(f"Unsupported operation '{operation_id}'.")
    operation_params = dict(params or {})
    repeat_count = int(operation_params.pop("_repeat", 1))
    if repeat_count < 1 or repeat_count > 20:
        raise ProcessingError("N must be between 1 and 20.")
    result = ensure_uint8(image)
    for _ in range(repeat_count):
        result = ensure_uint8(item["func"](result, operation_params))
    return result