#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author ：hhx
@Date ：2022/5/13 19:13 
@Description ：
"""
import matplotlib.pyplot as plt
import numpy as np
import torch
import os
from tqdm import tqdm

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 P(outputs1, outputs2):
    """皮尔逊距离"""
    outputs1 = preprocess_div(outputs1)
    outputs2 = preprocess_div(outputs2)
    x = np.vstack((outputs1, outputs2))
    return np.corrcoef(x)[0][-1]


def C(outputs1, outputs2):
    """余弦夹角距离"""
    return outputs1 @ outputs2 / (
            np.sqrt(np.sum(np.power(outputs1, 2))) * np.sqrt(np.sum(np.power(outputs2, 2))))


if __name__ == '__main__':

    index = 3
    if index == 0:
        title = 'Veg'
    elif index == 1:
        title = 'IS'
    elif index == 2:
        title = 'WB'
    else:
        title = 'BLI'

    path = 'result/指数/洛阳'
    Result_P_Resnext = np.zeros([len(os.listdir(path)) - 6, len(os.listdir(path)) - 6])
    for idx, i in enumerate(tqdm(os.listdir(path)[6:])):
        for idj, j in enumerate(os.listdir(path)[6:]):
            d1 = np.load(os.path.join(path, i))
            d2 = np.load(os.path.join(path, j))

            P_value = P(d1[index], d2[index])
            Result_P_Resnext[idx, idj] = P_value
            Result_P_Resnext[idj, idx] = P_value

    fig = plt.figure(figsize=(6, 6))
    ax = plt.subplot(111)
    cax = ax.matshow(1 - Result_P_Resnext, cmap='coolwarm', vmin=0, vmax=1)
    # cax = ax.matshow(Result_P_Resnext, cmap='hot', vmin=0, vmax=1)
    ax.set_title('{}_{}'.format(path.split('/')[-1], title))
    fig.colorbar(cax)

    plt.tight_layout()
    plt.savefig('result/指数/结果/{}_{}.png'.format(path.split('/')[-1], title))
    plt.show()