import numpy as np
from sklearn.linear_model import RANSACRegressor, LinearRegression
import matplotlib.pyplot as plt

def filter_middle_percent(data, percent):
    """
    保留数组中间80%的数据（删除首尾各10%）。

    参数:
        data (np.ndarray): 输入数组（可以是一维或多维，但会先展平）。

    返回:
        np.ndarray: 中间80%的数据。
    """
    # 展平数组（确保处理的是所有数据点）
    flattened_data = data.flatten()

    # # 计算10%和90%分位数
    # lower_bound = np.percentile(flattened_data, 15)
    # upper_bound = np.percentile(flattened_data, 75)
    #
    # # 筛选中间80%的数据
    # mask = (data >= lower_bound) & (data <= upper_bound)
    # 计算最大值、最小值和总范围
    data_min = np.min(flattened_data)
    data_max = np.max(flattened_data)
    data_range = data_max - data_min

    # 计算中间80%的上下界
    lower_bound = data_min + percent * data_range
    upper_bound = data_max - percent * data_range

    # 筛选数据
    mask = (flattened_data >= lower_bound) & (flattened_data <= upper_bound)

    return mask


def load_data(cameraModel, txt_name):
    """从用户输入的文件路径加载二维XY数据"""
    while True:
        filepath = txt_name.strip()
        if filepath.lower() == 'q':
            return None, None

        try:
            data = np.loadtxt(filepath)
            if data.shape[1] != 2:
                print("错误：文件必须包含两列数据（X和Y）")
                continue
            x = data[:, 0].reshape(-1, 1)
            y = data[:, 1].reshape(-1, 1)
            y = cameraModel.camera_height - y
            mask = filter_middle_percent(x, cameraModel.filt_percent)
            x_clean = x[mask]
            y_clean = y[mask]
            return x_clean.reshape(-1, 1), y_clean.reshape(-1, 1)
        except Exception as e:
            print(f"加载文件出错: {e}")


def ransac_fit(x, y, residual_threshold=2.5):
    """执行RANSAC拟合并返回模型和内点/外点"""
    ransac = RANSACRegressor(
        LinearRegression(),
        residual_threshold=residual_threshold,
        random_state=42
    )
    ransac.fit(x, y)

    inlier_mask = ransac.inlier_mask_
    outlier_mask = ~inlier_mask

    return ransac.estimator_, inlier_mask, outlier_mask


def get_data(cameraModel, txt_name):
    # 加载数据
    x, y = load_data(cameraModel, txt_name)
    if x is None:
        return 0, None, None, None, None, None, None, None, None

    # 第一次RANSAC拟合
    model1, inlier_mask1, outlier_mask1 = ransac_fit(x, y, residual_threshold=cameraModel.ransac_residual_threshold)

    x_inliers1 = x[inlier_mask1]
    y_inliers1 = y[inlier_mask1]

    # 获取第一次拟合的外点
    x_outliers1 = x[outlier_mask1]
    y_outliers1 = y[outlier_mask1]

    # 第二次RANSAC拟合（在外点上）
    model2, inlier_mask2, outlier_mask2 = None, None, None
    if len(x_outliers1) > 10:  # 确保有足够的外点进行第二次拟合
        model2, inlier_mask2, outlier_mask2 = ransac_fit(x_outliers1, y_outliers1,
                                                         residual_threshold=cameraModel.ransac_residual_threshold)

    x_inliers2 = x_outliers1[inlier_mask2]
    y_inliers2 = y_outliers1[inlier_mask2]

    # 获取第二次拟合的外点
    x_outliers2 = x_outliers1[outlier_mask2]
    y_outliers2 = y_outliers1[outlier_mask2]

    m1 = model1.predict(np.array([600]).reshape(-1, 1))
    m2 = model2.predict(np.array([600]).reshape(-1, 1))
    # 判断上下沿
    if m1 > m2:
        model_top, model_bot = model1, model2
        x_top, x_bot = x_inliers1, x_inliers2
        y_top, y_bot = y_inliers1, y_inliers2
    else:
        model_top, model_bot = model2, model1
        x_top, x_bot = x_inliers2, x_inliers1
        y_top, y_bot = y_inliers2, y_inliers1

    # 统一提取斜率和截距
    slope_top = model_top.coef_[0][0]
    intercept_top = model_top.intercept_[0]
    slope_bot = model_bot.coef_[0][0]
    intercept_bot = model_bot.intercept_[0]
    # plt.figure(figsize=(14, 7))
    #
    # # 绘制原始内点
    # plt.scatter(x_bot, y_bot,
    #             color='limegreen', marker='o', s=30, alpha=0.7,
    #             label='bot')
    # plt.scatter(x_top, y_top,
    #             color='red', marker='*', s=100, edgecolor='black',
    #             label='top')
    # plt.xlabel('X', fontsize=12)
    # plt.ylabel('Y', fontsize=12)
    # plt.title(f'{txt_name}', fontsize=14)
    # plt.legend(fontsize=10, loc='best')
    # plt.grid(True, alpha=0.3)
    # # plt.tight_layout()
    # plt.show()
    return 1, x_bot, y_bot, slope_bot, intercept_bot, x_top, y_top, slope_top, intercept_top