#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author ：hhx
@Date ：2022/5/13 20:34 
@Description ：根据不同波段指数生成特征图
"""
from utils import *
import numpy as np
import matplotlib.pyplot as plt
import os
from PIL import Image, ImageEnhance

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False


def preprocess_div(data1):
    data = data1.ravel()
    id = np.where((data.astype('str') == 'nan'))
    data = np.delete(data, id)
    return data


def Cal(B1, B2):
    return (B1 - B2) / (B1 + B2)


def Cal_SI(B1, B2, B3, B4):
    return ((B1 + B2) - (B3 + B4)) / ((B1 + B2) + (B3 + B4))


def plot_(i, data, id, title):
    temp = np.zeros_like(data)
    plt.subplot(220 + i)
    temp[id] = 255
    img1 = Image.fromarray(np.uint8(temp))
    plt.title(title)
    plt.imshow(img1)


def plot_index(i, id, title, fig):
    ax = plt.subplot(220 + i)
    cax = ax.matshow(id, cmap='coolwarm')
    ax.set_title(title)
    fig.colorbar(cax)

if __name__ == '__main__':
    Path = r'G:\哨兵2号数据\长江存储'
    c = 0
    Veg = []
    IS = []
    BL = []
    WB = []
    for i in os.listdir(Path):
        if i.endswith('tif'):

            data = readGeoTIFF(os.path.join(Path, i))
            temp = np.zeros([4, data.shape[1],  data.shape[2]])
            Blue, Green, Red, NIR, SWIR = data[1], data[2], data[3], data[7], data[11]

            IBI = Cal((2 * SWIR) / (SWIR + NIR), (NIR / (NIR + Red)) + (Green / (Green + SWIR)))   # 建筑指数

            NDBI = Cal(SWIR, NIR)
            EVI = (Red - NIR) * 2.5 / (NIR + 6 * Red - 7.5 * Blue + 1)  # >0.6
            MSAVI = 0.5 * (2 * NIR + 1 - np.sqrt((2 * NIR + 1) ** 2 - 8 * (NIR - Red)))  # 是个废物
            ARVI = (NIR - (2 * Red) + Blue) / (NIR + (2 * Red) + Blue)

            NDVI = Cal(NIR, Red)
            BUAI = NDBI - NDVI  # 通过这个提取植被
            MNDWI = Cal(data[2], data[11])  # 水体

            BLI = Cal(data[3], data[2])  # 裸土

            index1 = Cal(data[7], data[3])  # NDVI 归一化植被指数
            index2 = Cal(data[11], data[7])  # NDBI 归一化建筑指数
            index3 = Cal(data[2], data[7])  # NDWI 归一化水指数
            index4 = Cal_SI(data[11], data[3], data[7], data[1])  # SI 裸土指数

            temp[0] = ARVI
            temp[1] = index2
            temp[2] = MNDWI
            temp[3] = BLI
            np.save('result/指数/长江存储/{}.npy'.format(i.split('.')[0]), temp)
            c += 1