import cv2
import numpy as np
import matplotlib.pyplot as plt

def ransac_line_detection(image_path, threshold=5, max_iterations=1000):
    """
    使用RANSAC算法检测图像中的直线
    :param image_path: 输入图像路径
    :param threshold: 距离阈值
    :param max_iterations: 最大迭代次数
    :return: 原图和检测到直线的图像
    """
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 边缘检测
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # 获取边缘点坐标
    edge_points = np.column_stack(np.where(edges > 0))
    edge_points = edge_points[:, ::-1]  # 交换x和y坐标

    if len(edge_points) < 2:
        return img, img.copy()

    best_line = None
    best_inliers = 0
    best_model = None

    for iteration in range(max_iterations):
        # 随机选择两个点
        sample_indices = np.random.choice(len(edge_points), size=2, replace=False)
        p1, p2 = edge_points[sample_indices]

        # 避免选择相同的点
        if np.array_equal(p1, p2):
            continue

        # 计算直线参数 (ax + by + c = 0)
        x1, y1 = p1
        x2, y2 = p2

        if x1 == x2:  # 垂直线
            continue

        # 计算直线方程 ax + by + c = 0
        a = y2 - y1
        b = x1 - x2
        c = x2 * y1 - x1 * y2

        # 计算所有点到直线的距离
        distances = np.abs(a * edge_points[:, 0] + b * edge_points[:, 1] + c) / np.sqrt(a**2 + b**2)

        # 统计内点
        inlier_mask = distances < threshold
        inlier_count = np.sum(inlier_mask)

        # 更新最佳模型
        if inlier_count > best_inliers:
            best_inliers = inlier_count
            best_line = (a, b, c)
            best_model = inlier_mask

    # 绘制最佳直线
    result_img = img.copy()
    if best_line is not None and best_inliers > 10:  # 至少需要10个内点才认为是有效直线
        a, b, c = best_line
        if b != 0:  # 非垂直线
            # 计算直线在图像边界上的交点
            x1, x2 = 0, img.shape[1] - 1
            y1 = int((-c - a * x1) / b)
            y2 = int((-c - a * x2) / b)
            cv2.line(result_img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        else:  # 垂直线
            x = int(-c / a)
            cv2.line(result_img, (x, 0), (x, img.shape[0] - 1), (0, 255, 0), 2)

    return img, result_img

def multiple_ransac_lines(image_path, threshold=5, max_iterations=1000, min_line_points=20):
    """
    使用RANSAC检测多条直线
    :param image_path: 输入图像路径
    :param threshold: 距离阈值
    :param max_iterations: 最大迭代次数
    :param min_line_points: 检测一条直线所需的最少点数
    :return: 原图和检测到多条直线的图像
    """
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 边缘检测
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # 获取边缘点坐标
    edge_points = np.column_stack(np.where(edges > 0))
    edge_points = edge_points[:, ::-1]  # 交换x和y坐标

    if len(edge_points) < 2:
        return img, img.copy()

    remaining_points = edge_points.copy()
    result_img = img.copy()
    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255)]

    line_count = 0
    max_lines = 6  # 最多检测6条直线

    while len(remaining_points) >= min_line_points and line_count < max_lines:
        best_line = None
        best_inliers = 0
        best_model = None

        for iteration in range(max_iterations):
            # 随机选择两个点
            sample_indices = np.random.choice(len(remaining_points), size=2, replace=False)
            p1, p2 = remaining_points[sample_indices]

            # 避免选择相同的点
            if np.array_equal(p1, p2):
                continue

            # 计算直线参数 (ax + by + c = 0)
            x1, y1 = p1
            x2, y2 = p2

            if x1 == x2:  # 垂直线
                continue

            # 计算直线方程 ax + by + c = 0
            a = y2 - y1
            b = x1 - x2
            c = x2 * y1 - x1 * y2

            # 计算所有剩余点到直线的距离
            distances = np.abs(a * remaining_points[:, 0] + b * remaining_points[:, 1] + c) / np.sqrt(a**2 + b**2)

            # 统计内点
            inlier_mask = distances < threshold
            inlier_count = np.sum(inlier_mask)

            # 更新最佳模型
            if inlier_count > best_inliers:
                best_inliers = inlier_count
                best_line = (a, b, c)
                best_model = inlier_mask

        # 如果找到了足够的内点，则绘制直线并移除这些点
        if best_line is not None and best_inliers >= min_line_points:
            a, b, c = best_line
            if b != 0:  # 非垂直线
                # 计算直线在图像边界上的交点
                x1, x2 = 0, img.shape[1] - 1
                y1 = int((-c - a * x1) / b)
                y2 = int((-c - a * x2) / b)
                color = colors[line_count % len(colors)]
                cv2.line(result_img, (x1, y1), (x2, y2), color, 2)
            else:  # 垂直线
                x = int(-c / a)
                color = colors[line_count % len(colors)]
                cv2.line(result_img, (x, 0), (x, img.shape[0] - 1), color, 2)

            # 移除已匹配的点
            inlier_mask = best_model
            remaining_points = remaining_points[~inlier_mask]
            line_count += 1
        else:
            break  # 没有找到更多直线

    return img, result_img

def ransac_circle_detection(image_path, threshold=5, max_iterations=1000):
    """
    使用RANSAC算法检测图像中的圆形
    :param image_path: 输入图像路径
    :param threshold: 距离阈值
    :param max_iterations: 最大迭代次数
    :return: 原图和检测到圆形的图像
    """
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 边缘检测
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # 获取边缘点坐标
    edge_points = np.column_stack(np.where(edges > 0))
    edge_points = edge_points[:, ::-1]  # 交换x和y坐标

    if len(edge_points) < 3:
        return img, img.copy()

    best_circle = None
    best_inliers = 0

    for iteration in range(max_iterations):
        # 随机选择三个点
        sample_indices = np.random.choice(len(edge_points), size=3, replace=False)
        p1, p2, p3 = edge_points[sample_indices]

        # 检查三点是否共线
        area = abs((p1[0] * (p2[1] - p3[1]) + p2[0] * (p3[1] - p1[1]) + p3[0] * (p1[1] - p2[1])) / 2)
        if area < 1e-6:  # 三点共线
            continue

        # 计算外接圆的圆心和半径
        # 使用三点确定圆的方法
        x1, y1 = p1
        x2, y2 = p2
        x3, y3 = p3

        # 计算圆心 (cx, cy) 和半径 r
        a = x1 * (y2 - y3) - y1 * (x2 - x3) + x2 * y3 - x3 * y2
        bx = (x1 * x1 + y1 * y1) * (y3 - y2) + (x2 * x2 + y2 * y2) * (y1 - y3) + (x3 * x3 + y3 * y3) * (y2 - y1)
        by = (x1 * x1 + y1 * y1) * (x2 - x3) + (x2 * x2 + y2 * y2) * (x3 - x1) + (x3 * x3 + y3 * y3) * (x1 - x2)

        if abs(a) < 1e-6:
            continue

        cx = -bx / (2 * a)
        cy = -by / (2 * a)
        r = np.sqrt((x1 - cx)**2 + (y1 - cy)**2)

        # 计算所有点到圆的距离
        distances = np.abs(np.sqrt((edge_points[:, 0] - cx)**2 + (edge_points[:, 1] - cy)**2) - r)

        # 统计内点
        inlier_mask = distances < threshold
        inlier_count = np.sum(inlier_mask)

        # 更新最佳模型
        if inlier_count > best_inliers:
            best_inliers = inlier_count
            best_circle = (int(cx), int(cy), int(r))

    # 绘制最佳圆形
    result_img = img.copy()
    if best_circle is not None and best_inliers > 10:  # 至少需要10个内点才认为是有效圆形
        cv2.circle(result_img, best_circle[:2], best_circle[2], (0, 255, 0), 2)

    return img, result_img

def compare_ransac_hough(image_path):
    """
    比较RANSAC和霍夫变换的直线检测效果
    :param image_path: 输入图像路径
    :return: 原图、RANSAC结果和霍夫变换结果
    """
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 边缘检测
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # RANSAC直线检测
    ransac_img = img.copy()
    edge_points = np.column_stack(np.where(edges > 0))
    edge_points = edge_points[:, ::-1]  # 交换x和y坐标

    if len(edge_points) >= 2:
        best_line = None
        best_inliers = 0

        for iteration in range(500):  # 减少迭代次数以提高速度
            sample_indices = np.random.choice(len(edge_points), size=2, replace=False)
            p1, p2 = edge_points[sample_indices]

            if np.array_equal(p1, p2):
                continue

            x1, y1 = p1
            x2, y2 = p2

            if x1 == x2:
                continue

            a = y2 - y1
            b = x1 - x2
            c = x2 * y1 - x1 * y2

            distances = np.abs(a * edge_points[:, 0] + b * edge_points[:, 1] + c) / np.sqrt(a**2 + b**2)
            inlier_mask = distances < 5
            inlier_count = np.sum(inlier_mask)

            if inlier_count > best_inliers:
                best_inliers = inlier_count
                best_line = (a, b, c)

        if best_line is not None and best_inliers > 10:
            a, b, c = best_line
            if b != 0:
                x1, x2 = 0, img.shape[1] - 1
                y1 = int((-c - a * x1) / b)
                y2 = int((-c - a * x2) / b)
                cv2.line(ransac_img, (x1, y1), (x2, y2), (0, 255, 0), 2)

    # 霍夫变换直线检测
    hough_img = img.copy()
    lines = cv2.HoughLines(edges, 1, np.pi/180, 100)

    if lines is not None:
        for rho, theta in lines[:, 0]:
            a = np.cos(theta)
            b = np.sin(theta)
            x0 = a * rho
            y0 = b * rho
            x1 = int(x0 + 1000 * (-b))
            y1 = int(y0 + 1000 * (a))
            x2 = int(x0 - 1000 * (-b))
            y2 = int(y0 - 1000 * (a))
            cv2.line(hough_img, (x1, y1), (x2, y2), (0, 0, 255), 2)

    return img, ransac_img, hough_img

def robust_ransac_detection(image_path):
    """
    鲁棒的RANSAC检测（结合多种策略）
    :param image_path: 输入图像路径
    :return: 原图和鲁棒RANSAC检测结果
    """
    # 读取图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 应用不同的预处理技术
    # 1. 高斯模糊
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    
    # 2. 边缘检测
    edges = cv2.Canny(blurred, 50, 150, apertureSize=3)

    # 获取边缘点
    edge_points = np.column_stack(np.where(edges > 0))
    edge_points = edge_points[:, ::-1]  # 交换x和y坐标

    if len(edge_points) < 2:
        return img, img.copy()

    # 多轮RANSAC检测
    result_img = img.copy()
    remaining_points = edge_points.copy()
    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)]
    line_count = 0
    max_lines = 4

    while len(remaining_points) >= 20 and line_count < max_lines:
        best_line = None
        best_inliers = 0
        best_model = None

        # 动态调整参数
        adaptive_threshold = 3 + line_count  # 随着检测的直线增多，适当放宽阈值

        for iteration in range(300):  # 适中的迭代次数
            sample_indices = np.random.choice(len(remaining_points), size=2, replace=False)
            p1, p2 = remaining_points[sample_indices]

            if np.array_equal(p1, p2):
                continue

            x1, y1 = p1
            x2, y2 = p2

            if x1 == x2:
                continue

            a = y2 - y1
            b = x1 - x2
            c = x2 * y1 - x1 * y2

            distances = np.abs(a * remaining_points[:, 0] + b * remaining_points[:, 1] + c) / np.sqrt(a**2 + b**2)
            inlier_mask = distances < adaptive_threshold
            inlier_count = np.sum(inlier_mask)

            if inlier_count > best_inliers:
                best_inliers = inlier_count
                best_line = (a, b, c)
                best_model = inlier_mask

        if best_line is not None and best_inliers >= 15:
            a, b, c = best_line
            if b != 0:
                x1, x2 = 0, img.shape[1] - 1
                y1 = int((-c - a * x1) / b)
                y2 = int((-c - a * x2) / b)
                color = colors[line_count % len(colors)]
                cv2.line(result_img, (x1, y1), (x2, y2), color, 2)
            
            # 移除已匹配的点
            inlier_mask = best_model
            remaining_points = remaining_points[~inlier_mask]
            line_count += 1
        else:
            break

    return img, result_img

# 使用示例
if __name__ == "__main__":
    # 注意：需要提供实际的图像路径
    # img, result = ransac_line_detection('image.jpg', 5, 1000)
    # multi_img, multi_result = multiple_ransac_lines('image.jpg', 5, 1000, 20)
    # circle_img, circle_result = ransac_circle_detection('image.jpg', 5, 1000)
    # comp_img, ransac_result, hough_result = compare_ransac_hough('image.jpg')
    # robust_img, robust_result = robust_ransac_detection('image.jpg')
    pass