import matplotlib.pyplot as plt
import torch
import numpy as np

def show_data(x,y):

    if x.is_cuda:
        x = x.cpu().detach()

    if y.is_cuda:
        y = y.cpu()
        
    plt.figure(figsize=(12,6))

    #Draw the content of sub-window 1
    plt.subplot(121)
    plt.imshow(x)
    ##Draw the content of sub-window 2
    plt.subplot(122)
    plt.imshow(y)
    plt.show()

def calculate_acc(pred,label):
    cor = 0.
    (row,col) = pred.shape
    for i in range(len(pred)):
        pred_row = pred[i]
        label_row = label[i]

        pred_max_index = torch.argmax(pred_row)
        label_max_index = torch.argmax(label_row)

        if pred_max_index == label_max_index:
            cor += 1

    acc = float(cor/row)
    return acc

def calculate_acc_Dict(pred,label,N_k_types):

    pred=np.array(pred)
    ConfusionMatrix = np.zeros((N_k_types, N_k_types))
    (row,col) =pred.shape

    for i in range(len(pred)):
        pred_row = pred[i]
        label_row = label[i]

        pred_max_index = np.argmax(pred_row)
        label_max_index = np.argmax(label_row)

        ConfusionMatrix[label_max_index][pred_max_index]+=1
    Sum_row_CfM = np.sum(ConfusionMatrix, axis=1)
    Sum_col_CfM = np.sum(ConfusionMatrix, axis=0)
    AAs = np.diagonal(ConfusionMatrix) / Sum_row_CfM
    AO = np.trace(ConfusionMatrix) / row
    Po = AO
    Pe = np.sum(Sum_row_CfM * Sum_col_CfM) / row / row
    AA = np.mean(AAs)
    kappa = (AO-Pe) / (1-Pe)

    acc_Dict={'AO':AO, 'AA':AA, 'AAs':AAs, 'Po':Po, 'Pe':Pe, 'kappa':kappa, 'ConfusionMatrix':ConfusionMatrix}

    return acc_Dict


                

