'''
1、读取影像
2、bin
3、去除暗电流 + 转反射率
4、保存光谱
'''
import numpy as np
import matplotlib.pyplot as plt
import sys
from osgeo import gdal   #读写影像数据

class GRID:

    #读图像文件
    @classmethod
    def read_img(cls, filename):
        try:
            dataset = gdal.Open(filename)  # 打开文件
            im_width = dataset.RasterXSize  # 栅格矩阵的列数
            im_height = dataset.RasterYSize  # 栅格矩阵的行数
            num_bands = dataset.RasterCount  # 栅格矩阵的波段数
            im_geotrans = dataset.GetGeoTransform()  # 仿射矩阵
            im_proj = dataset.GetProjection()  # 地图投影信息
            im_data = dataset.ReadAsArray(0, 0, im_width, im_height)  # 将数据写成数组，对应栅格矩阵
            del dataset
            return im_proj, im_geotrans, im_data
        except:
            sys.exit()

    #写文件，以写成tif为例
    @classmethod
    def write_img(cls, dst_filename, data):
        format = "ENVI"
        driver = gdal.GetDriverByName(format)
        RasterXSize = data.shape[2]  # 遥感影像的sample（列数）
        RasterYSize = data.shape[1]  # 遥感影像的line（行数）
        band = data.shape[0]
        dst_ds = driver.Create(dst_filename, RasterXSize, RasterYSize,
                               band,
                               gdal.GDT_Float32)  # driver.Create()函数中RasterXSize代表影像的sample（列数），RasterYSize代表影像的line（行数）
        for i in range(band):
            dst_ds.GetRasterBand(i + 1).WriteArray(data[i, :, :])  # gdal的band从1开始，所以dst_ds.GetRasterBand(i+1)
        dst_ds = None

    # bin
    @classmethod
    def bin(cls, img, nBin):
        if nBin == 1:
            return img
        image_bin = np.empty((int(img.shape[0] / nBin), img.shape[1], img.shape[2]))
        k = np.arange(img.shape[0])[0::nBin]
        for i in range(image_bin.shape[0]):
            for j in range(nBin):
                image_bin[i] += img[k[i] + j]
        return image_bin

# 文件名
image = r'D:\py_program\corning410\test_spectral\lib_spectral2_20帧\delete\leaf'
dark_image = r'D:\py_program\corning410\test_spectral\lib_spectral2_20帧\delete\dark'
baiban_image = r'D:\py_program\corning410\test_spectral\lib_spectral2_20帧\delete\baiban'

# 读取影像
print('读取影像')
im_proj, im_geotrans, im_data = GRID.read_img(image)
d1, d2, dark_noise = GRID.read_img(dark_image)
b1, b2, baiban = GRID.read_img(baiban_image)



n_bin = 1
im_data = GRID.bin(im_data, n_bin)
dark_noise = GRID.bin(dark_noise, n_bin)
baiban = GRID.bin(baiban, n_bin)


# （1）去除暗电流；（2）转反射率
dark_noise_mean = dark_noise.mean(axis=1)
baiban_mean = baiban.mean(axis=1)
im_data = im_data.astype(np.float)
for i in range(im_data.shape[1]):
    im_data[:, i, :] = (im_data[:, i, :] - dark_noise_mean).astype(np.float) / baiban_mean.astype(np.float)


print('将影像写入到文件')
GRID.write_img(image + '_reflectivity', im_data)


# 计算波长
def calculate_wavelength(x):
    wavelength = x * 1.999564 - 279.893
    return wavelength
wavelength_tmp = np.empty(300)
for i in range(340, 640):
    wavelength_tmp[i - 340] = calculate_wavelength(i)
wavelength = np.empty(im_data.shape[0])
k = np.arange(300)[0::n_bin]
for i in range(wavelength.shape[0]):
    tmp = 0
    for j in range(n_bin):
        tmp += wavelength_tmp[k[i] + j]
    wavelength[i] = tmp / n_bin




# 保存光谱为txt文件
spectralNumber = 1
spectral_container = np.empty((im_data.shape[0], spectralNumber)).astype(np.float)
spectral_container = np.insert(spectral_container, 0, wavelength, axis=1)


spectral = im_data.mean(1).mean(1)
spectral_container = np.insert(spectral_container, 1, spectral, axis=1)
np.savetxt(r'D:\py_program\corning410\test_spectral\lib_spectral2_20帧\delete\spectral.txt', spectral_container[:, [0, 1]], fmt='%f')

plt.plot(wavelength, spectral)
plt.show()