from util import *
from osgeo import gdal
import argparse


@timeit
def otsu(img, max_value, data_water_mask, ignore_value=0, foreground=1, background=0):
    height = img.shape[0]
    width = img.shape[1]

    hist = np.zeros([max_value], np.float32)

    # 计算直方图
    invalid_counter = 0
    for i in range(height):
        for j in range(width):
            if img[i, j] == ignore_value or img[i, j] < 0 or data_water_mask[i, j] == 0:
                invalid_counter = invalid_counter + 1
                continue

            hist[img[i, j]] += 1

    hist /= (height * width - invalid_counter)

    threshold = 0
    deltaMax = 0
    # 遍历像素值，计算最大类间方差
    for i in range(max_value):
        wA = 0
        wB = 0
        uAtmp = 0
        uBtmp = 0
        uA = 0
        uB = 0
        u = 0
        for j in range(max_value):
            if j <= i:
                wA += hist[j]
                uAtmp += j * hist[j]
            else:
                wB += hist[j]
                uBtmp += j * hist[j]
        if wA == 0:
            wA = 1e-10
        if wB == 0:
            wB = 1e-10
        uA = uAtmp / wA
        uB = uBtmp / wB
        u = uAtmp + uBtmp

        # 计算类间方差
        deltaTmp = wA * ((uA - u)**2) + wB * ((uB - u)**2)
        # 找出最大类间方差以及阈值
        if deltaTmp > deltaMax:
            deltaMax = deltaTmp
            threshold = i

    # 二值化
    det_img = img.copy()
    det_img[img > threshold] = foreground
    det_img[img <= threshold] = background
    det_img[np.where(data_water_mask == 0)] = background
    return det_img


def find_overexposure_area(img_path, threhold=4095):
    # 第一步通过某个像素的光谱找到信号最强的波段

    # 根据上步所得的波段号检测过曝区域
    pass


@timeit
def find_severe_glint_area(img_path, water_mask, glint_wave=750, output_path=None):
    if output_path is None:
        output_path = append2filename(img_path, "_severe_glint_area")

    glint_band_number = find_band_number(glint_wave, img_path)

    dataset = gdal.Open(img_path)
    num_bands = dataset.RasterCount
    im_geotrans = dataset.GetGeoTransform()
    im_proj = dataset.GetProjection()

    tmp = dataset.GetRasterBand(glint_band_number + 1)  # 波段计数从1开始
    band_flare = tmp.ReadAsArray().astype(np.int16)
    band_flare_stretch = (band_flare / band_flare.max() * 255).astype(np.int32)

    dataset_water_mask = gdal.Open(water_mask)
    data_water_mask = dataset_water_mask.GetRasterBand(1).ReadAsArray()
    del dataset_water_mask

    # 不加1报错：IndexError: index 73 is out of bounds for axis 0 with size 73
    flare_binary = otsu(band_flare_stretch, band_flare_stretch.max() + 1, data_water_mask)

    write_bands(img_path, output_path, flare_binary)

    del dataset

    return output_path


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    img_path = r"D:\PycharmProjects\0water_rlx\test_data\ref_mosaic_1m_bsq"

    parser = argparse.ArgumentParser(description="此程序通过大律法分割图像，提取耀斑最严重的区域，输出的栅格和输入的影像具有相同的行列数。")

    parser.add_argument('-i1', '--input', type=str, required=True, help='输入影像文件的路径')
    parser.add_argument('-i2', '--input_water_mask', type=str, required=True, help='输入水域掩膜文件的路径')
    parser.add_argument('-gw', '--glint_wave', type=float, default=750.0, help='用于提取耀斑严重区域的波段.')
    parser.add_argument('-o', '--output', type=str, help='输出文件的路径')

    parser.add_argument('-v', '--verbose', action='store_true', help='启用详细模式')

    args = parser.parse_args()

    find_severe_glint_area(args.input, args.input_water_mask, args.glint_wave, args.output)