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_multi_band(image_path, mask_src, band_indices):
    """
    从图像中提取与掩膜重叠区域内掩膜值为1的像素值（多个波段），按行主序返回。
    返回值包含零值（保留原始顺序）。
    """
    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_indices, window=img_window)
        mask_data = mask_src.read(1, window=mask_window)

        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]

        result = {}
        for i, band in enumerate(band_indices):
            band_data = img_data[i]
            values = band_data[sorted_rows, sorted_cols].tolist()
            result[band] = values
        return result


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 plot_single_band(band, x_global, y_original, seg_info_list, segment_sizes_nonzero,
                     out_dir, poly_degree):
    """绘制单个波段的原始数据和分段拟合曲线，并保存图片。"""
    plt.figure(figsize=(14, 6))

    # 原始数据
    plt.plot(x_global, y_original, 'b-', linewidth=1.5, alpha=0.7, label='Original data')

    # 分段拟合曲线，只给第一个有效段添加图例
    labeled = False
    for seg in seg_info_list:
        if seg is not None:
            label = 'Fitted curve' if not labeled else ""
            plt.plot(seg['x_fine_global'], seg['y_fine'], 'r--', linewidth=2, label=label)
            labeled = True

    plt.xlabel("Pixel index after removing zeros (row-major order, cumulative across scenes)")
    plt.ylabel(f"Band {band+1} value")
    plt.title(f"Band {band+1}: Original vs. Fitted (per‑segment polynomial, deg={poly_degree})")
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.legend()

    # 标注航带分界线
    ylim = plt.ylim()
    cum_nz = 0
    total_nz = len(y_original)
    for i, nz_cnt in enumerate(segment_sizes_nonzero):
        if nz_cnt == 0:
            continue
        cum_nz += nz_cnt
        if cum_nz < total_nz:
            plt.axvline(x=cum_nz, color='gray', linestyle='--', alpha=0.7)
            plt.text(cum_nz, ylim[1] * 0.95, f"{nz_cnt}", ha='center', fontsize=8)
        else:
            plt.text(cum_nz - 0.5, ylim[1] * 0.95, f"{nz_cnt}", ha='right', fontsize=8)

    plt.tight_layout()
    out_path = os.path.join(out_dir, f"band_{band+1}_original_and_fitted.png")
    plt.savefig(out_path)
    plt.close()
    print(f"  -> 波段 {band+1} 图已保存至: {out_path}")


def main():
    # ==================== 用户可修改的参数 ====================
    folder = r"D:\BaiduNetdiskDownload\20250902\_3_52_52\316\R\test"
    mask_path = r"E:\is2\yaopu\925\roi_mask.dat"
    band_indices = [35,36,37]          # 波段 9,10,11 (0‑based)
    poly_degree = 2
    out_dir = r"E:\is2\yaopu\ocbrdf\test"
    # ========================================================

    os.makedirs(out_dir, exist_ok=True)

    if not os.path.exists(mask_path):
        print(f"掩膜文件不存在: {mask_path}")
        return
    mask_src = rasterio.open(mask_path)

    files = [os.path.join(folder, f) for f in os.listdir(folder)
             if f.lower().endswith(('.dat', '.img', '.bsq'))]
    files.sort()

    band_values_with_zeros = {band: [] for band in band_indices}
    segment_sizes_total = []

    for f in files:
        print(f"Processing {f}...")
        extracted = extract_mask_values_multi_band(f, mask_src, band_indices)
        if extracted:
            seg_len = len(extracted[band_indices[0]])
            segment_sizes_total.append(seg_len)
            for band in band_indices:
                band_values_with_zeros[band].extend(extracted[band])
        else:
            print(f"  -> 无重叠或掩膜内无像素")
            segment_sizes_total.append(0)

    mask_src.close()

    if not band_values_with_zeros[band_indices[0]]:
        print("没有有效的采样值。")
        return

    # 去除零值
    band_values_filtered = {}
    nonzero_masks = {}
    for band in band_indices:
        vals = np.array(band_values_with_zeros[band])
        mask_nonzero = (vals != 0)
        band_values_filtered[band] = vals[mask_nonzero].tolist()
        nonzero_masks[band] = mask_nonzero

    # 计算每个航带有效像素数
    segment_sizes_nonzero = []
    cum = 0
    for sz in segment_sizes_total:
        seg_mask = nonzero_masks[band_indices[0]][cum:cum+sz]
        nz_count = np.sum(seg_mask)
        segment_sizes_nonzero.append(nz_count)
        cum += sz

    # 构建每个波段的分段拟合信息
    segment_info_per_band = {}
    for band in band_indices:
        seg_info = []
        cum_nz = 0
        for nz_cnt in segment_sizes_nonzero:
            if nz_cnt == 0:
                seg_info.append(None)
                continue
            start_nz = cum_nz
            end_nz = cum_nz + nz_cnt
            seg_vals = band_values_filtered[band][start_nz:end_nz]
            x_local = np.arange(1, nz_cnt + 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 + cum_nz
                seg_info.append({
                    'x_fine_global': x_fine_global,
                    'y_fine': y_fine_local,
                    'nz_count': nz_cnt,
                    'start_idx': cum_nz + 1,
                    'end_idx': cum_nz + nz_cnt
                })
            else:
                seg_info.append(None)
            cum_nz += nz_cnt
        segment_info_per_band[band] = seg_info

    # 1. 绘制三个波段拟合曲线对比图
    plt.figure(figsize=(14, 6))
    colors = ['blue', 'green', 'red']
    labels = ['Band 9', 'Band 10', 'Band 11']
    for idx, band in enumerate(band_indices):
        seg_info = segment_info_per_band[band]
        labeled = False
        for seg in seg_info:
            if seg is not None:
                label = labels[idx] if not labeled else ""
                plt.plot(seg['x_fine_global'], seg['y_fine'],
                         color=colors[idx], linestyle='--', linewidth=2, label=label)
                labeled = True
    plt.xlabel("Pixel index after removing zeros (row-major order, cumulative across scenes)")
    plt.ylabel("Pixel value")
    plt.title(f"Fitted variation for bands 9, 10, 11 (per-segment polynomial fit, deg={poly_degree})")
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.legend()
    ylim = plt.ylim()
    cum_nz = 0
    total_nz = len(band_values_filtered[band_indices[0]])
    for i, nz_cnt in enumerate(segment_sizes_nonzero):
        if nz_cnt == 0:
            continue
        cum_nz += nz_cnt
        if cum_nz < total_nz:
            plt.axvline(x=cum_nz, color='gray', linestyle='--', alpha=0.7)
            plt.text(cum_nz, ylim[1] * 0.95, f"{nz_cnt}", ha='center', fontsize=8)
        else:
            plt.text(cum_nz - 0.5, ylim[1] * 0.95, f"{nz_cnt}", ha='right', fontsize=8)
    plt.tight_layout()
    combined_path = os.path.join(out_dir, "roi_fitted_bands_9_10_11_zeros_removed.png")
    plt.savefig(combined_path)
    plt.close()
    print(f"组合对比图已保存至: {combined_path}")

    # 2. 为每个波段单独绘制原始数据+拟合曲线图
    for band in band_indices:
        x_global = np.arange(1, len(band_values_filtered[band]) + 1)
        y_original = np.array(band_values_filtered[band])
        seg_info = segment_info_per_band[band]
        plot_single_band(band, x_global, y_original, seg_info,
                         segment_sizes_nonzero, out_dir, poly_degree)

    # 打印拟合系数摘要
    print("\n各航带拟合多项式系数（仅显示每个波段第一个有拟合的航带示例）")
    for band in band_indices:
        print(f"\n波段 {band+1}:")
        for i, seg in enumerate(segment_info_per_band[band]):
            if seg is not None:
                nz_cnt = seg['nz_count']
                seg_vals = band_values_filtered[band][seg['start_idx']-1:seg['end_idx']]
                x_local = np.arange(1, nz_cnt+1)
                coeffs = np.polyfit(x_local, seg_vals, poly_degree)
                print(f"  航带 {i+1}: {coeffs}")
                break


if __name__ == "__main__":
    main()