import numpy as np
from util import *
from osgeo import gdal


@timeit
def deglint_subtract_nir(imgpath, water_mask, out_imgpath, start_wave, end_wave):
    glint = average_bands_in_mask(start_wave, end_wave, imgpath, water_mask)

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

    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)

    for i in range(num_bands):
        band = dataset.GetRasterBand(i + 1).ReadAsArray()
        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))