import numpy as np

from util import *
from osgeo import gdal


def calculate_glint_spatial_distribution(imgpath, left_shoulder_wave, valley_wave, right_shoulder_wave, shoulder_window=1, valley_window=1):
    left_shoulder_band = average_bands(left_shoulder_wave - shoulder_window, left_shoulder_wave + shoulder_window, imgpath)
    valley_band = average_bands(valley_wave - valley_window, valley_wave + valley_window, imgpath)
    right_shoulder_band = average_bands(right_shoulder_wave - shoulder_window, right_shoulder_wave + shoulder_window, imgpath)

    glint_spatial_distribution = (left_shoulder_band + right_shoulder_band) / 2 - valley_band
    glint_spatial_distribution = glint_spatial_distribution / glint_spatial_distribution.max()

    glint_spatial_distribution[np.where(glint_spatial_distribution < 0)] = 0

    write_bands(imgpath, r"D:\PycharmProjects\sun_glint\test_data\glint_spatial_distribution", glint_spatial_distribution)

    return glint_spatial_distribution


@timeit
def deglint_oxygen_absorption_valley(imgpath, water_mask, out_imgpath, left_shoulder_wave, valley_wave, right_shoulder_wave, shoulder_window=1, valley_window=1):
    glint_spatial_distribution = calculate_glint_spatial_distribution(imgpath, left_shoulder_wave, valley_wave,
                                                                      right_shoulder_wave, shoulder_window,
                                                                      valley_window)

    dataset = gdal.Open(imgpath)
    im_width = dataset.RasterXSize
    im_height = dataset.RasterYSize
    num_bands = dataset.RasterCount
    geotransform = dataset.GetGeoTransform()
    im_proj = dataset.GetProjection()

    # 通过data ignore value获取影像有效区域

    format = "ENVI"
    driver = gdal.GetDriverByName(format)
    dst_ds = driver.Create(out_imgpath, im_width, im_height, num_bands, gdal.GDT_Int16,
                           options=["INTERLEAVE=BSQ"])
    dst_ds.SetGeoTransform(geotransform)
    dst_ds.SetProjection(im_proj)

    dataset_water_mask = gdal.Open(water_mask)
    data_water_mask = dataset_water_mask.GetRasterBand(1).ReadAsArray()
    del dataset_water_mask
    for i in range(num_bands):
        band = dataset.GetRasterBand(i + 1).ReadAsArray()

        # 找波段最大最下值时，需要掩膜掉没有影像的区域（一般为0值的区域）
        glint = glint_spatial_distribution * (band.max() - min(band[np.where(band > 0)])) * 1
        glint[np.where(data_water_mask == 0)] = 0
        band_deglint = band - glint

        # 确保值大于等于0
        band_deglint[np.where(band_deglint < 0)] = 0

        dst_ds.GetRasterBand(i + 1).WriteArray(band_deglint)

    del dataset, dst_ds

    write_fields_to_hdrfile(get_hdr_file_path(imgpath), get_hdr_file_path(out_imgpath))