import os
import cv2
import numpy as np
from pathlib import Path
from sklearn.linear_model import RANSACRegressor

# ---------------------------
# 读取 ROI 列表 txt
# ---------------------------
def load_rois_from_txt(txt_path):
    rois = []
    if not os.path.exists(txt_path):
        print(f"❌ ROI 文件不存在: {txt_path}")
        return rois
    with open(txt_path, 'r') as f:
        for line in f:
            s = line.strip()
            if s:
                try:
                    x, y, w, h = map(int, s.split(','))
                    rois.append((x, y, w, h))
                except Exception as e:
                    print(f"⚠️ 无法解析 ROI 行 '{s}': {e}")
    return rois

# ---------------------------
# PCA + RANSAC + 迭代去离群点拟合直线
# ---------------------------
def fit_line_best(points, distance_thresh=5, max_iter=5):
    if len(points) < 2:
        return None

    points = points.astype(np.float32)

    for _ in range(max_iter):
        mean = np.mean(points, axis=0)
        cov = np.cov(points.T)
        eigvals, eigvecs = np.linalg.eig(cov)
        idx = np.argmax(eigvals)
        direction = eigvecs[:, idx]
        vx, vy = direction
        x0, y0 = mean

        dists = np.abs(vy*(points[:,0]-x0) - vx*(points[:,1]-y0)) / np.hypot(vx, vy)
        inliers = points[dists <= distance_thresh]

        if len(inliers) == len(points) or len(inliers) < 2:
            break
        points = inliers

    if len(points) < 2:
        return None

    # RANSAC 拟合 y = kx + b
    X = points[:, 0].reshape(-1, 1)
    y = points[:, 1]
    try:
        ransac = RANSACRegressor(residual_threshold=distance_thresh)
        ransac.fit(X, y)
        k = ransac.estimator_.coef_[0]
        b = ransac.estimator_.intercept_

        vx = 1 / np.sqrt(1 + k**2)
        vy = k / np.sqrt(1 + k**2)
        x0 = np.mean(points[:,0])
        y0 = k*x0 + b
    except:
        mean = np.mean(points, axis=0)
        cov = np.cov(points.T)
        eigvals, eigvecs = np.linalg.eig(cov)
        idx = np.argmax(eigvals)
        direction = eigvecs[:, idx]
        vx, vy = direction
        x0, y0 = mean

    return vx, vy, x0, y0

# ---------------------------
# 封装函数：读取 ROI txt -> 拟合直线 -> 可视化
# ---------------------------
def fit_lines_from_image_txt(image_path, roi_txt_path, distance_thresh=5, draw_overlay=True):
    """
    输入：
        image_path: 原图路径
        roi_txt_path: ROI txt 文件路径，每行 x,y,w,h
        distance_thresh: 直线拟合残差阈值
        draw_overlay: 是否在原图上叠加拟合直线
    输出：
        lines: 每个 ROI 的拟合直线 [(vx, vy, x0, y0), ...]
        overlay_img: 可视化原图叠加拟合直线
    """
    img = cv2.imread(image_path)
    if img is None:
        print(f"❌ 无法读取图片: {image_path}")
        return [], None

    rois = load_rois_from_txt(roi_txt_path)
    if not rois:
        print("❌ 没有有效 ROI")
        return [], None

    overlay_img = img.copy() if draw_overlay else None
    lines = []

    for idx, (x, y, w, h) in enumerate(rois):
        roi = img[y:y+h, x:x+w]
        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 100, 200)
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

        longest_contour = max(contours, key=lambda c: cv2.arcLength(c, closed=False), default=None)

        if longest_contour is not None and len(longest_contour) >= 2:
            points = longest_contour.reshape(-1, 2)
            line = fit_line_best(points, distance_thresh=distance_thresh, max_iter=5)
            if line:
                vx, vy, x0, y0 = line
                lines.append(line)

                if draw_overlay:
                    cols = gray.shape[1]
                    lefty = int(y0 - vy/vx * x0)
                    righty = int(y0 + vy/vx * (cols - x0))

                    # 绘制在原图
                    cv2.line(overlay_img, (x, y + lefty), (x + cols - 1, y + righty), (0, 0, 255), 2)
                    cv2.drawContours(overlay_img, [longest_contour + np.array([x, y])], -1, (0, 255, 0), 1)
            else:
                lines.append(None)
        else:
            lines.append(None)

    return lines, overlay_img

# ---------------------------
# 使用示例
# ---------------------------
if __name__ == "__main__":
    image_path = "../test_image/1.jpg"
    roi_txt_path = "../roi_coordinates/1_rois1.txt"

    lines, overlay = fit_lines_from_image_txt(image_path, roi_txt_path, distance_thresh=5)

    for idx, line in enumerate(lines):
        print(f"ROI {idx} 拟合直线: {line}")

    if overlay is not None:
        cv2.imwrite("overlay_result.jpg", overlay)
        print("✅ 原图叠加拟合直线已保存: overlay_result.jpg")