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


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

    # 计算波长
    @classmethod
    def calculate_wavelength(cls, x):
        wavelength = x * 1.999564 - 279.893
        return wavelength

wavelength = np.empty(639 - 339)
for i in range(339, 639):
    wavelength[i - 339] = GRID.calculate_wavelength(i)


# 等效于ENVI拉伸：No stretch
def stretch(img, minimum=0, maximum=255):
    if len(img.shape) == 2:
        img_new = (img - minimum) / (maximum - minimum)
        img_new[img_new < 0] = 0
        img_new[img_new > 1] = 1
        return img_new
    else:
        img_new = np.empty(img.shape)
        for i in range(img.shape[2]):
            img_new[:, :, i] = (img[:, :, i] - minimum) / (maximum - minimum)
            img_new[:, :, i][img_new[:, :, i] < 0] = 0
            img_new[:, :, i][img_new[:, :, i] > 1] = 1
            return img_new


# 当lowPercentile=0, highPercentile=100时，等效于Min-Max Stretching
# lowPercentile=2, highPercentile=98时，等效于ENVI拉伸：Linear 2%
# https://blog.csdn.net/LEILEI18A/article/details/80180483
def percentile_stretching(img, lowPercentile=0, highPercentile=100, minout=0, maxout=255):
    if len(img.shape) == 2:
        low = np.percentile(img, lowPercentile)
        up = np.percentile(img, highPercentile)

        img_new = ((img - low) / (up - low)) * (maxout - minout) + minout
        img_new[img_new < minout] = minout
        img_new[img_new > maxout] = maxout
        img_out = np.uint8(img_new)
        return img_out
    else:  # 对于彩色照片，需要先单独对每个波段拉伸
        img_new = np.empty(img.shape)
        for i in range(img.shape[2]):
            low = np.percentile(img[:, :, i], lowPercentile)
            up = np.percentile(img[:, :, i], highPercentile)

            img_new[:, :, i] = minout + ((img[:, :, i] - low) / (up - low)) * (maxout - minout)
            img_new[:, :, i][img_new[:, :, i] < minout] = minout
            img_new[:, :, i][img_new[:, :, i] > maxout] = maxout
        img_out = np.uint8(img_new)
        return img_out


# 画出图像直方图
# https://blog.csdn.net/fly_wt/article/details/83904207
def image_hist(image):  # 画三通道图像的直方图
    color = ("blue", "green", "red")  # 画笔颜色的值可以为大写或小写或只写首字母或大小写混合
    for i, color in enumerate(color):
        hist = cv2.calcHist([image], [i], None, [256], [0, 256])
        plt.plot(hist, color=color)
        plt.xlim([0, 256])
    plt.show()


print("读取影像")
#image = r'D:\py_program\corning410\2%拉伸显示问题 + 漏帧\x270\dn值\corning410_test10'
image = r'D:\corning410_test10'
im_proj, im_geotrans, im_data = GRID.read_img(image)


print("挑取波段用于真彩色显示")
rgb_raw = np.dstack((im_data[121], im_data[76], im_data[36]))
rgb = rgb_raw.astype(np.uint8)


x1 = stretch(rgb_raw)  #等效于ENVI拉伸：No stretch
x2 = percentile_stretching(rgb_raw)  # Min-Max Stretching
x3 = percentile_stretching(rgb_raw, 2, 98)  # 2% Stretching

print("画出影像")
plt.imshow(x3)
plt.show()