import warnings

import matplotlib.pyplot as plt

warnings.filterwarnings('ignore')
warnings.simplefilter('ignore')
import torch
from torch import nn
from models.model_RW import Res_Deeplab
from utils import *
from pylab import *

mpl.rcParams['font.sans-serif'] = ['SimSun']
mpl.rcParams['axes.unicode_minus'] = False
mpl.rcParams["font.size"] = 13
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

from GradCAM import GradCAM

# 1.加载模型
model = Res_Deeplab(True)
model = model.to(device=device)
model_state_dict = torch.load(
    r'model_data\model_0.116.pth',
    map_location='cpu')
model.load_state_dict(model_state_dict)
model.eval()


def normalizer(data):
    std = torch.std(data, dim=0)
    mean = torch.mean(data, dim=0)
    return (data - mean) / std


class SemanticSegmentationTarget:
    def __init__(self, category, mask):
        self.category = category
        self.mask = torch.from_numpy(mask)
        if torch.cuda.is_available():
            self.mask = self.mask.cuda()

    def __call__(self, model_output):
        return (model_output[self.category, :] * self.mask).sum()


if __name__ == '__main__':

    train_image = torch.tensor(np.load('dataset/data.npy')).to(device).float()
    train_fulllabels = torch.tensor(np.load('dataset/label.npy')).to(device).float()

    pred = model(train_image)
    normalized_masks = torch.nn.functional.softmax(pred, dim=1).cpu()

    car_category = 4   # 4代表不透水面
    car_mask = normalized_masks[0, :, :].argmax(axis=0).detach().cpu().numpy()
    car_mask_float = np.float32(car_mask == car_category)

    target_layers = [model.conv5a]
    targets = [SemanticSegmentationTarget(car_category, car_mask_float)]
    with GradCAM(model=model,
                 target_layers=target_layers) as cam:
        grayscale_cam = cam(input_tensor=train_image, targets=targets, size=200)[0]
    if np.mean(grayscale_cam) != 0:
        plt.figure(figsize=(9, 8))

        plt.subplot(511)
        plt.title('标签')
        pred1 = torch.argmax(input=pred, dim=1).cpu().numpy()
        plt.plot(pred1[0])
        xlabel = [0, 12, 24, 36, 47]  # 设置横轴的自定义点
        plt.xticks(xlabel, [DateList[i] for i in xlabel], rotation=0)


        plt.subplot(512)
        plt.title('Sentinel-2时间序列')
        plt.plot(train_image[0].cpu().T)

        plt.xticks(xlabel, [DateList[i] for i in xlabel], rotation=0)
        plt.grid(axis='y')
        target_layers = [model.conv6]
        with GradCAM(model=model, target_layers=target_layers) as cam:
            grayscale_cam1 = cam(input_tensor=train_image, targets=targets, size=200)[0]

        plt.subplot(513)

        m1 = np.tile(grayscale_cam[0], (5)).reshape(5, 200)
        plt.imshow(m1, cmap='rainbow')
        plt.title('随机游走前CAM')
        plt.xticks([])
        plt.yticks([])

        plt.subplot(514)
        m2 = np.tile(grayscale_cam1[0], (5)).reshape(5, 200)
        plt.imshow(m2, cmap='rainbow')
        plt.title('随机游走后CAM')
        plt.xticks([])
        plt.yticks([])
        plt.tight_layout()

        plt.subplot(515)
        m = nn.Softmax(dim=0)
        y = pred[0]
        y = normalizer(y)
        y = m(y)
        y = y.detach().cpu().numpy()
        colorlist = ['#00BFFF', '#006400', '#3CB371', '#DAA520', '#FF0000', '#FFDEAD']
        namelist = ['水体', '林地', '草地', '裸地', '不透水面', '耕地']
        for id, out in enumerate(y):
            plt.plot(out, color=colorlist[id])
            plt.fill_between(x=range(48), y1=0, y2=out, facecolor=colorlist[id], alpha=0.3, label=namelist[id])

        plt.legend(loc=2,
                   bbox_to_anchor=(1, 1),
                   borderaxespad=0.2,
                   ncol=1,
                   frameon=True, fontsize=13)
        xlabel = [0, 12, 24, 36, 47]  # 设置横轴的自定义点
        plt.xticks(xlabel, [DateList[i] for i in xlabel], rotation=0)
        plt.title('预测结果')
        plt.tight_layout()
        plt.savefig('img.png', dpi=200)
        plt.show()