import os
import numpy as np
import matplotlib.pyplot as plt
import rasterio
from rasterio.windows import from_bounds
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)


def extract_mask_values(image_path, mask_src, band_index=0):
    """
    从图像中提取与掩膜重叠区域内掩膜值为1的像素值，按行主序返回。

    参数
    ----------
    image_path : str
        图像文件路径（需包含 .hdr 头文件）
    mask_src : rasterio.DatasetReader
        已打开的掩膜文件（用于读取掩膜数组和变换）
    band_index : int
        波段索引（0‑based）

    返回
    -------
    list of float
        按行主序排列的像素值
    """
    with rasterio.open(image_path) as img_src:
        # 获取图像与掩膜的交集窗口（以图像坐标系为准）
        # 计算重叠区域的边界（在掩膜投影坐标系下）
        img_bounds = img_src.bounds
        mask_bounds = mask_src.bounds
        overlap_left = max(img_bounds.left, mask_bounds.left)
        overlap_bottom = max(img_bounds.bottom, mask_bounds.bottom)
        overlap_right = min(img_bounds.right, mask_bounds.right)
        overlap_top = min(img_bounds.top, mask_bounds.top)
        if overlap_left >= overlap_right or overlap_bottom >= overlap_top:
            return []   # 无重叠

        # 将重叠边界转换到图像像素窗口
        img_window = from_bounds(overlap_left, overlap_bottom, overlap_right, overlap_top, img_src.transform)
        # 将重叠边界转换到掩膜像素窗口
        mask_window = from_bounds(overlap_left, overlap_bottom, overlap_right, overlap_top, mask_src.transform)

        # 读取图像的重叠区域波段数据
        img_data = img_src.read(band_index + 1, window=img_window)
        # 读取掩膜的重叠区域（单波段）
        mask_data = mask_src.read(1, window=mask_window)

        # 提取掩膜值为1的像素位置（在重叠窗口内）
        rows, cols = np.where(mask_data == 1)
        if len(rows) == 0:
            return []

        # 按行主序排序（先按行，再按列）
        idx = np.lexsort((cols, rows))
        sorted_rows = rows[idx]
        sorted_cols = cols[idx]

        # 从图像数据中获取对应像素值
        values = img_data[sorted_rows, sorted_cols].tolist()
        return values


def fit_segment(x, y, degree):
    """对一段数据执行多项式拟合，返回多项式函数和拟合点（用于绘图）"""
    if len(x) < degree + 1:
        # 有效点数不足，降低阶数
        degree = max(1, len(x) - 1)
        if degree < 1:
            return None, None, None
    coeffs = np.polyfit(x, y, degree)
    poly = np.poly1d(coeffs)
    x_fine = np.linspace(x[0], x[-1], 200)
    y_fine = poly(x_fine)
    return poly, x_fine, y_fine


def main():
    # ==================== 用户可修改的参数 ====================
    folder = r"D:\BaiduNetdiskDownload\20250902\_3_52_52\316\R\test"   # 存放BSQ图像的文件夹
    mask_path = r"E:\is2\yaopu\925\roi_mask.dat"                 # 单波段掩膜 DAT 文件（值为1的区域为ROI）
    band_idx = 10                                               # 要提取的波段索引（0‑based）
    poly_degree = 2                                               # 多项式拟合阶数（全局，各航带独立拟合）
    # ========================================================

    # 打开掩膜文件（保持打开状态以便重复读取）
    if not os.path.exists(mask_path):
        print(f"掩膜文件不存在: {mask_path}")
        return
    mask_src = rasterio.open(mask_path)

    # 获取文件夹内所有高光谱文件（支持 .dat, .img, .bsq）
    files = [os.path.join(folder, f) for f in os.listdir(folder)
             if f.lower().endswith(('.dat', '.img', '.bsq'))]
    files.sort()

    # 存储全局信息
    all_filtered_values = []      # 所有有效像素值（按顺序）
    segment_info = []             # 每个航带的 (全局起始索引, 全局结束索引, 有效像素值数组, 拟合结果)
    cum_len = 0                   # 累积的有效像素数（全局x坐标）

    for i, f in enumerate(files):
        print(f"Processing {f}...")
        values = extract_mask_values(f, mask_src, band_idx)
        if values:
            # 去除零值，保留顺序
            seg_vals = [v for v in values if v != 0]
            n_seg = len(seg_vals)
            if n_seg > 0:
                start_idx = cum_len + 1
                end_idx = cum_len + n_seg
                cum_len += n_seg
                all_filtered_values.extend(seg_vals)

                # 对该段进行拟合（局部坐标1..n_seg）
                x_local = np.arange(1, n_seg + 1)
                y_local = np.array(seg_vals)
                # 调用拟合函数
                poly, x_fine_local, y_fine_local = fit_segment(x_local, y_local, poly_degree)
                # 将局部拟合坐标转换为全局坐标
                if x_fine_local is not None:
                    x_fine_global = x_fine_local + start_idx - 1
                    segment_info.append({
                        'start_idx': start_idx,
                        'end_idx': end_idx,
                        'seg_vals': seg_vals,
                        'poly': poly,
                        'x_fine_global': x_fine_global,
                        'y_fine': y_fine_local,
                        'n_vals': n_seg
                    })
                else:
                    # 拟合失败（点数太少），仍记录段信息但无拟合曲线
                    segment_info.append({
                        'start_idx': start_idx,
                        'end_idx': end_idx,
                        'seg_vals': seg_vals,
                        'poly': None,
                        'x_fine_global': None,
                        'y_fine': None,
                        'n_vals': n_seg
                    })
            else:
                # 该航带无有效像素
                print(f"  -> 航带 {i+1} 无有效像素（所有值为0）")
                segment_info.append(None)   # 占位，便于后续处理（可选）
        else:
            print(f"  -> 无重叠或掩膜内无像素")
            segment_info.append(None)

    mask_src.close()

    if not all_filtered_values:
        print("没有有效的采样值。")
        return

    n_total = len(all_filtered_values)
    x_global = np.arange(1, n_total + 1)

    # 绘制图形
    plt.figure(figsize=(14, 6))

    # 绘制原始数据折线（所有有效像素）
    plt.plot(x_global, all_filtered_values, 'b-', linewidth=1, label='Measured values (zeros removed)')

    # 绘制各航带拟合曲线（不同颜色）
    colors = plt.cm.tab10(np.linspace(0, 1, len([seg for seg in segment_info if seg is not None])))
    color_idx = 0
    for seg in segment_info:
        if seg is None or seg['poly'] is None:
            continue
        # 使用不同颜色绘制拟合曲线
        plt.plot(seg['x_fine_global'], seg['y_fine'], '--', linewidth=2, color=colors[color_idx],
                 label=f'Segment {color_idx+1} fit (deg={poly_degree})')
        color_idx += 1

    # 添加图例（自动合并重复项，这里不合并）
    plt.legend()

    plt.xlabel("Pixel index after removing zeros (row-major order, cumulative across scenes)")
    plt.ylabel(f"Band {band_idx + 1} value")
    plt.title("Spectral variation across water pixels in ROI (per‑segment polynomial fit)")
    plt.grid(True, linestyle='--', alpha=0.5)

    # 标注每个航带的分界线（基于有效像素累积）
    ylim = plt.ylim()
    for seg in segment_info:
        if seg is None:
            continue
        end_idx = seg['end_idx']
        n_vals = seg['n_vals']
        if end_idx < n_total:   # 非最后一个段画竖线
            plt.axvline(x=end_idx, color='gray', linestyle='--', alpha=0.7)
            plt.text(end_idx, ylim[1] * 0.95, f"{n_vals}", ha='center', fontsize=8)
        else:
            # 最后一个段：在右边界标注数量
            plt.text(end_idx - 0.5, ylim[1] * 0.95, f"{n_vals}", ha='right', fontsize=8)

    plt.tight_layout()
    out_path = r"E:\is2\yaopu\ocbrdf\test\roi_variation_original_10band.png"
    plt.savefig(out_path)
    print(f"图形已保存至: {out_path}")

    # 可选：打印各段拟合系数
    print("\n各航带拟合多项式系数（从最高次到常数项）：")
    for i, seg in enumerate(segment_info):
        if seg is not None and seg['poly'] is not None:
            coeffs = seg['poly'].coeffs
            print(f"航带 {i+1}（有效像素数 {seg['n_vals']}）：")
            for j, c in enumerate(coeffs):
                print(f"  x^{len(coeffs)-1-j}: {c:.6f}")
        elif seg is not None:
            print(f"航带 {i+1}：无有效像素或拟合失败")

if __name__ == "__main__":
    main()