import cv2
import datetime
import gc
import glob
import imageio
import matplotlib
import matplotlib.pyplot as plt
import multiprocessing.pool
import numpy as np
import openslide
import os
import os.path
import pandas as pd
import random
import seaborn as sns
import shutil
import sys
import timeit
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import warnings
import psutil

from PIL import Image
from albumentations import *
from albumentations.torch import ToTensor
from efficientnet_pytorch import EfficientNet
from functools import partial
from itertools import cycle
from multiprocessing import Pool
from optparse import OptionParser
from scipy import interp
from sklearn import metrics
from sklearn.metrics import classification_report,accuracy_score,confusion_matrix,auc,roc_curve,roc_auc_score
from sklearn.svm import SVC
from torch.nn import functional as F
from torch.utils.data import DataLoader, Dataset, sampler


#### get_opt ####
parser = OptionParser()
parser.add_option("--slidePath", default=None, dest="slidePath", help="input slide path", metavar="SLIDE_PATH")
parser.add_option("--maskPath", default=None, dest="maskPath", help="output mask path", metavar="MASK_PATH")
parser.add_option("--patchPath", default=None, dest="patchPath", help="output patch path", metavar="PATCH_PATH")
parser.add_option("--nonZeroProp", default=0.2, type="float", dest="nonZeroProp", help="non zero proption of patch", metavar="NON_ZERO_PROP")
parser.add_option("--judgePatch", default='yes', type="string", dest="judgePatch", help="judge patch's nonZeroProp or not", metavar="JUDGE_OR_NOT")
parser.add_option("--patchAbsPathLabelCSV", default=None, dest="patchAbsPathLabelCSV", help="csv path of patch abspath label", metavar="PATCH_ABSPATH_LABEL_CSV_PATH")
parser.add_option("--predictTmpResultCSV", default=None, dest="predictTmpResultCSV", help="csv path of png predict label", metavar="PNG_PREDICT_LABEL_CSV_PATH")
parser.add_option("--wholeSlideInferDir", default=None, dest="wholeSlideInferDir", help="path of whole slide infer", metavar="WHOLE_SLIDE_INFER_PATH")
parser.add_option("--tifPredClassCSV" ,default=None, dest="tifPredClassCSV", help="csv path of tif pred", metavar="WHOLE_SLIDE_PREDICT_CSV_PATH")
parser.add_option("--heatmapDir", default=None, dest="heatmapDir", help="path of heatmap", metavar="HEATMAP_PATH")
parser.add_option("--leftDis", default=0.82, type="float", dest="leftDis", help="left distance", metavar="LEFT_DISTANCE")

(options, args) = parser.parse_args()

slidePath = options.slidePath
maskPath = options.maskPath
patchPath = options.patchPath
nonZeroProp = options.nonZeroProp
judgePatch = options.judgePatch
patchAbsPathLabelCSV = options.patchAbsPathLabelCSV
predictTmpResultCSV = options.predictTmpResultCSV
wholeSlideInferDir = options.wholeSlideInferDir
tifPredClassCSV = options.tifPredClassCSV
heatmapDir = options.heatmapDir
leftDis = options.leftDis


#### util_hist.py ####
cores = 16  # hp.CORES


#### ops_hist.py ####
def get_wsi_path(path):
    '''
    get sorted normal wsi paths for get mask
    :param path: normal wsi path
    :return: sort wsi paths
    '''
    wsi_paths = glob.glob(os.path.join(path, "*.tif"))
    wsi_paths.sort()
    return wsi_paths

def name_and_exist(path, maskPath, last_filename):
    normal_map_name = os.path.split(path)[-1].split('.')[0] + last_filename
    normal_map_exist = chk_file(maskPath, normal_map_name)
    return normal_map_name, normal_map_exist

def chk_file(dir_path, filename):  # 检查mask是否已经生成过
    exsit = False
    for file_name in os.listdir(dir_path):
        if file_name == filename:
            exsit = True
    return exsit

def saved_tissue_mask(slide, mask_dir, tissue_mask_name, level):
    normal_slide_lv = slide.read_region((0, 0), level, slide.level_dimensions[level])
    R, G, B, _ = normal_slide_lv.split()
    normal_slide_lv_rgb = np.array(Image.merge('RGB', [R, G, B]))
    normal_slide_lv_rgb = cv2.cvtColor(normal_slide_lv_rgb, cv2.COLOR_BGR2HSV)
    normal_slide_lv_rgb = normal_slide_lv_rgb[:, :, 1]
    _, normal_tissue_mask = cv2.threshold(normal_slide_lv_rgb, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    cv2.imwrite(os.path.join(mask_dir, tissue_mask_name), np.array(normal_tissue_mask))  # 保存mask到……/mask/SLIDENAME_tissue_mask.png


#### mask.py ####
def get_mask(wsi_path, level, mask_dir):
    '''
    :param path: normal slide paths
    :param level: level
    :return: get tissue mask and slide map
    '''
    tissue_mask_name, tissue_mask_exist = name_and_exist(wsi_path, mask_dir, "_tissue_mask.png")  # 返回如('MG1802789-2020-04-26_23_56_30_tissue_mask.png', False)
        
    if tissue_mask_exist:
        return
    else:
        #print('Get mask from: %s' % wsi_path)
        try:
            slide = openslide.OpenSlide(wsi_path)  #get wsi picture
            #normal_slide_map = np.array(slide.get_thumbnail(slide.level_dimensions[level]))
        except:
            #print('Warning: fail to  get mask from {0}'.format(wsi_path))
            #fail_mask.write(wsi_path+'\n')
            return
        saved_tissue_mask(slide, mask_dir, tissue_mask_name, level)

#slidePath = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/tif'
#maskPath = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/mask'
#fail_mask = open(r'/home/jwang/preciMedWorkFlow/deepIHC/fail_to_get_mask.txt','a')
level = 1
#print('Getting masks now...')
p = Pool(cores)
slideAbsPathLst = get_wsi_path(slidePath)  # 以'.tif'结尾的slide的绝对路径列表
if not os.path.exists(maskPath):  # 创建mask/目录
    os.makedirs(maskPath)
partial_get_mask = partial(get_mask, level=level, mask_dir=maskPath)
p.map(partial_get_mask, slideAbsPathLst)
p.close()
p.join()
#fail_mask.close()


#### util_new.py ####
level = 1
PATCH_SIZE = 256


#### ops_wsi_new.py ####
def get_wsi_path(path):
    '''
    get sorted normal wsi paths for get mask
    :param path: normal wsi path
    :return: sort wsi paths
    '''
    normal_wsi_paths = glob.glob(os.path.join(path, "*.tif"))
    # normal_wsi_paths.sort()
    return normal_wsi_paths

def name_and_exist(path, dir, last_filename):
    normal_map_name = os.path.split(path)[-1].split('.')[0] + last_filename  # 'MG1802789-2020-04-26_23_56_30_tissue_mask.png'
    normal_map_exist = chk_file(dir, normal_map_name)  # mask是否已存在
    return normal_map_name, normal_map_exist
    
def chk_file(dir_path, filename):
    exsit = False
    for file_name in os.listdir(dir_path):  # 遍历/……/mask里的所有文件
        if file_name == filename:
            exsit = True
    return exsit

def name_and_exist_fuzzy(path, dir, last_filename):
    normal_map_name = os.path.split(path)[-1].split('.')[0] + last_filename  # 'MG1802789-2020-04-26_23_56_30_*_*.png'
    normal_map_exist = chk_file_fuzzy(dir, normal_map_name)
    return normal_map_name, normal_map_exist
    
def chk_file_fuzzy(dir_path,file_name):
    exist = False
    results = glob.glob(os.path.join(dir_path,file_name))  # 匹配 /……/patch/SLIDENAME/MG1802789-2020-04-26_23_56_30_*_*.png 的列表
    if results:
        exist = True
    return exist
    
def get_bbox_from_mask_image(mask_image):
    contours, hierarchy = cv2.findContours(mask_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    bounding_boxes = [cv2.boundingRect(c) for c in contours]
    return bounding_boxes


#### my_patch_extractor_pool.py ####
#def get_normal_patches_from_normal_slide(slide_path, maskPath, level, patchPath, nonZeroProp, judgePatch):
#    extract_patches_from_slide(slide_path, maskPath, level, patchPath, nonZeroProp, judgePatch)
    
def get_samples_of_patch_starting_points_with_stride(mask_image):
    bounding_boxes = get_bbox_from_mask_image(mask_image)
    #print('bounding_boxes:', bounding_boxes)
    list_starting_points = []
    for x, y, w, h in bounding_boxes:  # 如：x=251,y=3042,w=2,h=3
        stride_limit = min(w, h)  # stride_limit=2
        #print(stride_limit)
        if PATCH_SIZE > stride_limit:
            continue
        else:
            devider = 1
            stride = int(PATCH_SIZE / devider)
            #print(stride)
        X = range(x, x + w - PATCH_SIZE, stride)  # np.arange(x, x+w)
        Y = range(y, y + h - PATCH_SIZE, stride)  # np.arange(y, y+h)
        for row_starting_point in X:
            for col_starting_point in Y:
                list_starting_points.append((row_starting_point, col_starting_point))
    return list_starting_points
    
def extract_patches_from_slide(slide_path_, maskPath, level, patchPath, nonZeroProp=0.2, judgePatch=True):
    for i in range(len(slide_path_)):
        slide_path = slide_path_[i]  # 列表的每个tif文件的绝对路径
        #print('slide_path:', slide_path)
        patchSlideDir = os.path.join(patchPath, os.path.split(slide_path)[-1].split('.')[0])  # /……/patch/SLIDENAME
        #print('patchSlideDir:', patchSlideDir)
        if not os.path.exists(patchSlideDir):  # 以SLIDENAME为目录名创建一个目录
            os.makedirs(patchSlideDir)
        tissue_mask_name, tissue_mask_exist = name_and_exist(slide_path, maskPath, '_tissue_mask.png')  # 返回如('MG1802789-2020-04-26_23_56_30_tissue_mask.png', False)，可修改丢弃tissue_mask_exist
        #print('tissue_mask_name:', tissue_mask_name)
        #print('tissue_mask_exist:', tissue_mask_exist)
        #patch_name, patch_exist = name_and_exist_fuzzy(slide_path, patchSlideDir, '_*_*.png')  # 返回如('MG1802789-2020-04-26_23_56_30_*_*.png', False)，可修改丢弃patch_exist
        #print('patch_name:', patch_name)
        #print('patch_exist:', patch_exist)
        slide_name = os.path.split(slide_path)[-1]  # 返回‘MG1802789-2020-04-26_23_56_30.tif’
        #print('slide_name:', slide_name)
        if slide_name:
            #if not patch_exist:
            slide = openslide.OpenSlide(slide_path)
            tissue_mask = cv2.imread(os.path.join(maskPath, tissue_mask_name), 0)
            tissue_patches_start_points = get_samples_of_patch_starting_points_with_stride(tissue_mask)
            down_samples = round(slide.level_downsamples[level])
            sample_accepted = 0
            sample_rejected = 0
            #nonZeroProp = 0.2
            #print('tissue_patches_start_points:', tissue_patches_start_points)
            for x, y in tissue_patches_start_points:
                if judgePatch:
                    if np.sum(tissue_mask[y:y+PATCH_SIZE,x:x+PATCH_SIZE])/255.0 >= PATCH_SIZE*PATCH_SIZE*nonZeroProp:
                        patch_read_from_wsi_at_level = slide.read_region((x * down_samples, y * down_samples),level, (PATCH_SIZE, PATCH_SIZE))
                        r, g, b, _ = patch_read_from_wsi_at_level.split()
                        patch_rgb = Image.merge("RGB", [r, g, b])
                        patch_name = str(x) + '_' + str(y) + '.png'
                        patch_fullname = os.path.join(patchSlideDir, patch_name)
                        patch_rgb.save(patch_fullname)
                        sample_accepted += 1
                    else:
                        sample_rejected += 1
                elif not judgePatch:
                    patch_read_from_wsi_at_level = slide.read_region((x * down_samples, y * down_samples),level, (PATCH_SIZE, PATCH_SIZE))
                    r, g, b, _ = patch_read_from_wsi_at_level.split()
                    patch_rgb = Image.merge("RGB", [r, g, b])
                    patch_name = str(x) + '_' + str(y) + '.png'
                    patch_fullname = os.path.join(patchSlideDir, patch_name)
                    patch_rgb.save(patch_fullname)
                    sample_accepted += 1
            slide.close()
            #print('File: %s; Accept: %d; Reject: %d' % (slide_name, sample_accepted, sample_rejected))
        else:
            pass
            #print('{} existed'.format(slide_name))

#start_date = datetime.datetime.now()
#start_time = timeit.default_timer()
#print('test class: >>>>')
slideAbsPathLst = get_wsi_path(slidePath)  # slidePath下的所有tif文件的列表
#print('slideAbsPathLst:', slideAbsPathLst)
#patchPath = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/patch'
if judgePatch == 'yes':
    judgePatch = True
else:
    judgePatch = False
extract_patches_from_slide(slideAbsPathLst, maskPath, level, patchPath, nonZeroProp, judgePatch)
#end_date = datetime.datetime.now()
#end_time = timeit.default_timer()
#hour = (end_time - start_time) // 3600
#minute = (end_time - start_time) % 3600 // 60
#second = (end_time - start_time) % 3600 % 60
#start_date = start_date.strftime('%Y-%m-%d %H:%M:%S')
#end_date = end_date.strftime('%Y-%m-%d %H:%M:%S')
#print("\nStart date: {0}".format(start_date))
#print("End date: {0}".format(end_date))
#print('CPU time is {0}h {1}m {2}s'.format(hour, minute, second))


#### get_fname.py ####
def _iter_valide_files(directory, white_list_formats, follow_links):
    def _recursive_list(subpath):
        return sorted(os.walk(subpath, followlinks=follow_links),
                      key=lambda x: x[0])  # 返回os.walk输出的元组根据第一个元素排序后的列表
    for root, _, files in _recursive_list(directory):
        for fname in sorted(files):  # 每个png排序后
            for extension in white_list_formats:
                if fname.lower().endswith('.tiff'):
                    warnings.warn('Using \'.tiff\' files with multiple bands '
                                  'will cause distortion. '
                                  'Please verify your output.')
                if fname.lower().endswith('.' + extension):
                    yield root, fname  # ('/home/jwang/preciMedWorkFlow/deepIHC/testTiff/patch/P28642-9-2020-05-27_11_56_47', 'P28642-9-2020-05-27_0_1_0.png')。每次为不同png

def _list_valid_filenames_in_directory(directory, white_list_formats,
                                       class_indices,
                                       follow_links):
    dirname = os.path.basename(directory)  # PATCH_DIR1
    valid_files = _iter_valide_files(directory,  # '/home/jwang/preciMedWorkFlow/deepIHC/testTiff/patch/PATCH_DIR1'等所有文件夹
                                     white_list_formats,  # {'png'}
                                     follow_links)  # False
    classes = []
    filenames = []
    for root, fname in valid_files:
        classes.append(class_indices[dirname])
        absolute_path = os.path.join(root, fname)
        # relative_path = os.path.join(
            # dirname, os.path.relpath(absolute_path, directory))
        filenames.append(absolute_path)
    return classes, filenames

def get_class_fnames(directory, classes=None,
                     white_list_formats = {'png', 'jpg', 'jpeg'},
                     follow_links=False):
    '''
    get the one by one order, for its correspond. such as predict1 -- file1
    :param directory:the base directory such as tumor0, and tumor0 contains tumor, normal subdirectory
    :param classes:if None, it will be [subdir1, subdir0, subdir2.....], else it will be classes=[tumor,normal]
    :param white_list_formats:the file format need feed to the iter
    :param follow_links:always False.
    :return:classes, filenames  ======  [0,1,1,0,1,0,0,1], [normal.png, tumor.png, tumor.png, ..]
    '''
    if not classes:
        classes = []
        for subdir in sorted(os.listdir(directory)):  # 每个存有各自patch的文件夹 SLIDENAME1、SLIDENAME2
            if os.path.isdir(os.path.join(directory, subdir)):
                classes.append(subdir)  # [SLIDENAME1, SLIDENAME2]即[PATCH_DIR1, PATCH_DIR2]
    #num_classes = len(classes)
    class_indices = dict(zip(classes, range(len(classes))))  # {PATCH_DIR1: 0, PATCH_DIR1: 1, …}
    pool = multiprocessing.pool.ThreadPool()
    results = []
    filenames = []
    i = 0
    for dirpath in (os.path.join(directory, subdir) for subdir in classes):  # dirpath为PATCH_DIR所有文件夹的绝对路径
        results.append(pool.apply_async(_list_valid_filenames_in_directory,
                                        (dirpath,  # dirpath如'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/patch/PATCH_DIR1'等所有文件夹
                                        white_list_formats,  # {'png'}
                                         class_indices,  # {PATCH_DIR1: 0, PATCH_DIR1: 1, …}
                                         follow_links)))  # False
    classes_list = []
    filenames_list = []
    for res in results:
        classes, filenames = res.get()
        classes_list += classes
        filenames_list += filenames
    return classes_list, filenames_list

#patchAbsPathLabelCSV = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/tmpResult/patch2tif.csv'
if not os.path.exists(os.path.dirname(patchAbsPathLabelCSV)):  # 创建tmpResult目录
    os.makedirs(os.path.dirname(patchAbsPathLabelCSV))
classes,fnames = get_class_fnames(patchPath, classes = None, white_list_formats = {'png'}, follow_links=False)
print('fnames:', fnames)
with open(patchAbsPathLabelCSV,'w') as f:  # patchPath下的每个文件夹内的png文件的绝对路径，写入同一列
    f.write('TestPatchName\n')
    lines = '\n'.join(fnames)
    f.writelines(lines)


#### yanghuan_pred_eff_softmax_task.py ####
gc.collect()
#start_time = timeit.default_timer()
device = torch.device("cpu")
#warnings.filterwarnings("ignore")
seed = 69
random.seed(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
np.random.seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True
model = EfficientNet.from_name('efficientnet-b5')
model.load_state_dict(torch.load(r'/home/jwang/preciMedWorkFlow/deepIHC/efficientnet-b5-b6417697.pth'))
for p in  model.parameters():
    p.requires_grad = False
model._fc = nn.Sequential(nn.Linear(in_features=2048, out_features=128, bias=True),
             nn.Linear(in_features=128, out_features=6, bias=True),nn.Softmax(1))
ckpt_path = r'/home/jwang/preciMedWorkFlow/deepIHC/sysu1_20X_0.5_0overlap_eff_1e-5.pth'  # 模型文件
model.to(device)
model.eval()
state = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
model.load_state_dict(state["state_dict"])
class TestDataset(Dataset):
    '''Dataset for test prediction'''
    def __init__(self, df, mean, std):
        self.fnames = df.index.tolist()
        self.num_samples = len(self.fnames)
        self.transform = Compose([Normalize(mean=mean, std=std, p=1), ToTensor(),])
    def __getitem__(self, idx):
        fname = self.fnames[idx]
        path = fname
        image = cv2.imread(path)
        images = self.transform(image=image)["image"]
        return fname, images
    def __len__(self):
        return self.num_samples
#patchAbsPathLabelCSV = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/tmpResult/patch2tif.csv'
df_tmp = pd.read_csv(patchAbsPathLabelCSV, encoding='utf_8_sig')
print('df_tmp:', df_tmp)
df_tmp = df_tmp.set_index('TestPatchName')
print('df_tmp:', df_tmp)
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)
print('psutil.virtual_memory().available:', psutil.virtual_memory().available/1024/1024)
testset = DataLoader(
    TestDataset(df_tmp, mean, std),
    batch_size=100,
    shuffle=False,
    num_workers=0,
    pin_memory=True)
print('testset:', testset)
print('psutil.virtual_memory().available:', psutil.virtual_memory().available/1024/1024)
preds = []
fnames_list = []
#flag = 0
print('enumerate(testset):', enumerate(testset))
for i, batch in enumerate(testset):
    #print('i:',i)
    #print('batch:', batch)
    #if i//100 > flag:
        #print(i)
        #flag += 1
    fnames, images = batch
    batch_preds = model(images.to(device))
    batch_preds = batch_preds.detach().cpu().numpy()
    preds.append(batch_preds)
    fnames_list.extend(fnames)
preds_cat = np.concatenate(preds,axis=0)
fnames_cat = np.array(fnames_list)
result = pd.DataFrame(preds_cat)
result.index = fnames_cat
#predictTmpResultCSV = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/tmpResult/predictTmpResult.csv'
result.to_csv(predictTmpResultCSV,encoding='utf_8_sig')  # 每个slide的png的6个分类的预测结果写入同一csv，各列为png绝对路径、第0类概率、第1类概率、……、第5类概率。
#end_time = timeit.default_timer()
#hour = (end_time - start_time) // 3600
#minute = (end_time - start_time) % 3600 // 60
#second = (end_time - start_time) % 3600 % 60
#print('over. CPU time is {0}h {1}m {2}s'.format(hour, minute, second))


#### wsi_level_inference.py ####
def reorganize_patch_prob(predictTmpResultCSV,wholeSlideInferDir):
    if not os.path.exists(wholeSlideInferDir):  # 创建wsiInferResult
        os.makedirs(wholeSlideInferDir)
    with open(predictTmpResultCSV) as f:
        SNAME = ''
        head = f.readline()
        cols = head.strip().split(',')
        cols = ['Probability_' + class_index[c] for c in cols[1:]] + ['PatchName']
        new_head = ','.join(cols) + '\n'    
        wsi_patch_dic = {}
        for line in f:
            items = line.strip().split(',')
            patch_full_name = items[0]
            slide_name,patch_name = patch_full_name.split('/')[-2:]
            patch_full_name = '/'.join(patch_full_name.split('/')[-2:])
            res_line = ','.join(items[1:] + [patch_full_name]) + '\n'
            if slide_name in wsi_patch_dic:
                wsi_patch_dic[slide_name].append(res_line)
            else:
                wsi_patch_dic[slide_name] = [res_line]
        for wsi in wsi_patch_dic:
            slide_result = os.path.join(wholeSlideInferDir,wsi + '_result.csv')
            with open(slide_result,'w') as resf:
                resf.write(new_head)
                resf.writelines(wsi_patch_dic[wsi])

class_index = {'0':'ac','1':'g','2':'ng','3':'normal','4':'scc','5':'smcc'}
#predictTmpResultCSV = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/tmpResult/preditTmpResult.csv'
#wholeSlideInferDir = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/tmpResult/wsiInferResult'
reorganize_patch_prob(predictTmpResultCSV, wholeSlideInferDir)


#### get_roc_acc_sysu_update.py ####
def infer_line_class(prob_list):
    tumor_cutoff = 0.1
    g_cutoff = 0.3
    ng_cutoff = 0.4
    normal_cutoff = 0.9
    ac,g,ng,normal,scc,smcc = prob_list  # 右边如[0.048, 0.095, 0.143, 0.19, 0.238, 0.286]
    if normal >= normal_cutoff:
        line_class = 'normal'
    tumor_group = [ac,scc,smcc]
    tumor_index = ['ac','scc','smcc']
    tumor_max = max(tumor_group)  # 0.286
    sup = normal-0.2  # -0.01
    if tumor_max >= tumor_cutoff:
        line_class = tumor_index[tumor_group.index(tumor_max)]  # smcc
    elif normal >= normal_cutoff or (g<0.1 and ng<0.2) :
        line_class = 'normal'
    elif g >= g_cutoff:
        line_class = 'g'
    elif ng>ng_cutoff:
        line_class = 'ng'
    elif ng >=0.2 and ng>=sup and  g<0.2:
        line_class = 'ng'
    elif g >= 0.2:
        line_class = 'g'
    else:
        line_class = 'normal'
    return line_class

def infer_line_class_old(prob_list):
    tumor_cutoff = 0.1
    g_cutoff = 0.3
    ng_cutoff = 0.4
    normal_cutoff = 0.9
    ac,g,ng,normal,scc,smcc = prob_list  # 右边[0.0024823523, 9.748552e-05, 0.0016754136, 0.9953035, 0.00025055165, 0.00019066907]
    if normal >= normal_cutoff:
        line_class = 'normal'
    else:
        tumor_group = [ac,scc,smcc]
        tumor_index = ['ac','scc','smcc']
        tumor_max = max(tumor_group)
        if tumor_max >= tumor_cutoff:
            line_class = tumor_index[tumor_group.index(tumor_max)]        
        else:
            if g >= g_cutoff:
                line_class = 'g'
            else:
                sup = normal-0.2
                if ng >= ng_cutoff or ng >= sup:
                    line_class = 'ng'
                else:
                    line_class_prob = max(prob_list)
                    class_index = ['ac','g','ng','normal','scc','smcc']
                    line_class =  class_index[prob_list.index(line_class_prob)]
    return line_class  # 类别字符串，如'normal'

def get_prediction_balance_tumor_nontumor_sensitivity(patches_pred_file, conf_thresholds,normal_percent_thred=0.85):
    classes = list(conf_thresholds.keys())  # ['ac', 'g', 'ng', 'normal', 'scc', 'smcc']
    class_to_count = {_class:0 for _class in classes} #各类别计数，{'ac': 0, 'g': 0, 'ng': 0, 'normal': 0, 'scc': 0, 'smcc': 0}
    patches_pred_lines = open(patches_pred_file, 'r').readlines()[1:]  # ['0.0024823523, 9.748552e-05, 0.0016754136, 0.9953035, 0.00025055165, 0.00019066907, P28642-9-2020-05-27_11_56_47/P28642-9-2020-05-27_0_1_0.png\n', ……]
    for line in patches_pred_lines:
        line_items = line[:-1].split(',')[:-1]  # 每个png的，如['0.0024823523', ' 9.748552e-05', ' 0.0016754136', ' 0.9953035', ' 0.00025055165', ' 0.00019066907']
        confs = [float(x) for x in line_items]  # 转成浮点数，如[0.0024823523, 9.748552e-05, 0.0016754136, 0.9953035, 0.00025055165, 0.00019066907]
        line_class = infer_line_class_old(confs)  # patch水平类别判断，如'normal'
        class_to_count[line_class] += 1  # 类别计数+1
    if sum(class_to_count.values()) > 0:
        class_to_percent = {_class:class_to_count[_class]/sum(class_to_count.values()) for _class in class_to_count}  # 计算各类个数占所有的比重，如{'ac': 0.047619047619047616, 'g': 0.09523809523809523, 'ng': 0.14285714285714285, 'normal': 0.19047619047619047, 'scc': 0.23809523809523808, 'smcc': 0.2857142857142857}
    else:
        class_to_percent = {_class:0.0 for _class in class_to_count}
    wsi_prob_list = list(class_to_percent.values())  # 如[0.047619047619047616, 0.09523809523809523, 0.14285714285714285, 0.19047619047619047, 0.23809523809523808, 0.2857142857142857]
    wsi_prob_list = [round(x,3) for x in wsi_prob_list]  # 保留3位小数，如[0.048, 0.095, 0.143, 0.19, 0.238, 0.286]
    wsi_class = infer_line_class(wsi_prob_list)  # slide水平类别判断，如'normal'
    return wsi_prob_list, wsi_class  # 如([0.048, 0.095, 0.143, 0.19, 0.238, 0.286], 'smcc')
    
def write_wsi_agg_pred(wsi_label_dic,input_dir,tifPredClassCSV,conf_thresholds,normal_thred):
    cols = ['WSI_Name', 'Pred_Label'] + ['Probability_' + x for x in {'0':'ac','1':'g','2':'ng','3':'normal','4':'scc','5':'smcc'}.values()]
    new_head = ','.join(cols) + '\n'
    f = open(tifPredClassCSV,'w')
    f.write(new_head)
    for res in os.listdir(input_dir):  # 每个如'P28642-9-2020-05-27_11_56_47_result.csv'的文件
        if res.endswith('_result.csv'):
            wsi_name = res.split('/')[-1].split('_')[0] + '.tif'  # 如'P28642-9-2020-05-27.tif'
            in_file = os.path.join(input_dir,res)  # 如'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/tmpResult/wsiInferResult/P28642-9-2020-05-27_11_56_47_result.csv'
            sequenced_patch_probs,str_wsi_label = get_prediction_balance_tumor_nontumor_sensitivity(in_file, conf_thresholds,normal_thred)  # 如([0.048, 0.095, 0.143, 0.19, 0.238, 0.286], 'smcc')
            sequenced_patch_probs_str = [str(x) for x in sequenced_patch_probs]  # 如['0.048', '0.095', '0.143', '0.19', '0.238', '0.286']
            items = [wsi_name,str_wsi_label] + sequenced_patch_probs_str 
            outline = ','.join(items) + '\n'  # 'P28642-9-2020-05-27.tif,smcc,0.048,0.095,0.143,0.19,0.238,0.286\n'
            f.write(outline)
    f.close()

if not os.path.exists(os.path.dirname(tifPredClassCSV)):  # 创建getTifPredClass目录
    os.makedirs(os.path.dirname(tifPredClassCSV))
normal_thred = 0.9
#my_target_names = ['LUAD','PTB','OP','NL','LUSC','SCLC']
conf_thresholds = {'ac':0.1, 'g': 0.3,'ng':0.4,'normal':0.9,'scc':0.1,'smcc':0.1}
#class_index = {'0':'ac','1':'g','2':'ng','3':'normal','4':'scc','5':'smcc'}
#wholeSlideInferDir = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/tmpResult/wsiInferResult'
#tifPredClassCSV = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/tmpResult/getTifPredClass/tifPredClass.csv'
write_wsi_agg_pred({}, wholeSlideInferDir, tifPredClassCSV, conf_thresholds, normal_thred)


#### heatmap.py ####
def dict_tiles_stats(pred_df_path):
    pred_df = pd.read_csv(pred_df_path)
    #print(pred_df)
    pred_df = pred_df[['PatchName','Probability_normal','Probability_ac','Probability_g','Probability_ng','Probability_scc',
                       'Probability_smcc']]
    #print(pred_df['PatchName'])
    pred_df['PatchName'] = pred_df['PatchName'].map(lambda x: os.path.basename(x).split('.')[0] )
    #print(pred_df)
    stats_dict = pred_df.set_index('PatchName').T.to_dict('list')
    #print(stats_dict)
    return stats_dict

def get_heatmap(stats_dict,patch_size,resample_factor,base_path,divide=True):
    stats_dict = stats_dict
    flag = False
    for tile in sorted(stats_dict.keys()):
        # extract coordinates of the tile
        im2 = imageio.imread(os.path.join(base_path,tile+'.png'))
        ixTile = int(tile.split('_')[-2])
        iyTile = int(tile.split('_')[-1])
        # check how bif the "re-combined" slide should be (invert col/row because of the swapaxes required)
        rTile = patch_size
        cTile = patch_size
        xTile =  ixTile
        yTile =  iyTile
        req_xLength = xTile + rTile
        req_yLength = yTile + cTile
        if resample_factor > 0:
            rTile = int(rTile / resample_factor)
            cTile = int(cTile / resample_factor)
            im2s = np.array(Image.fromarray(im2).resize((cTile, rTile)))
            xTile = int(xTile / resample_factor)
            yTile = int(yTile / resample_factor)
            req_xLength = xTile + rTile
            req_yLength = yTile + cTile
        else:
            im2s = im2
        if flag:
            NewSlide = False
        else:
            # Moved to a new slide
            flag = True
            NewSlide = True
            # create a new re-combined slide    
            WholeSlide_0 = np.zeros([req_xLength, req_yLength, 3])
            HeatMap_0 = np.zeros([req_xLength, req_yLength, 3])
            HeatMap_divider = np.zeros([req_xLength, req_yLength, 3])        
        oClass, cmap, current_score = get_inference_from_file(tile, stats_dict)
        # prepare heatmap
#         print("current score")
#         print(current_score)
        if current_score < 0:
            print("No probability found")
        else:
            if NewSlide == False:
                # append to the new re-combined slide        
                WholeSlide_temp = np.zeros([max(WholeSlide_0.shape[0], req_xLength), max(WholeSlide_0.shape[1],req_yLength), 3])
                WholeSlide_temp[0:WholeSlide_0.shape[0],0:WholeSlide_0.shape[1],:] = WholeSlide_0
                WholeSlide_temp[xTile:req_xLength, yTile:req_yLength,:] = np.swapaxes(im2s,0,1)
                WholeSlide_0 = WholeSlide_temp
                del WholeSlide_temp
                HeatMap_temp = WholeSlide_0 * 0
                HeatMap_temp[0:HeatMap_0.shape[0], 0:HeatMap_0.shape[1],:] = HeatMap_0
                HeatMap_0 = HeatMap_temp
                del HeatMap_temp
                HeatMap_divider_tmp = WholeSlide_0 * 0
                HeatMap_divider_tmp[0:HeatMap_divider.shape[0], 0:HeatMap_divider.shape[1],:] = HeatMap_divider
                HeatMap_divider = HeatMap_divider_tmp
                del HeatMap_divider_tmp
            else:
                WholeSlide_0[xTile:req_xLength, yTile:req_yLength,:] = np.swapaxes(im2s,0,1)
            #heattile = np.ones([512,512]) * current_score
            heattile = np.ones([req_xLength-xTile,req_yLength-yTile]) * current_score
            heattile = cmap(heattile)
            heattile = heattile[:,:,0:3]
            '''divider: get mean of overlapped patches'''
            if divide:
                HeatMap_0[xTile:req_xLength, yTile:req_yLength,:] = HeatMap_0[xTile:req_xLength, yTile:req_yLength,:] + heattile
            else:
                HeatMap_0[xTile:req_xLength, yTile:req_yLength,:] = heattile#HeatMap_0[xTile:req_xLength, yTile:req_yLength,:] + heattile
            HeatMap_divider[xTile:req_xLength, yTile:req_yLength,:] = HeatMap_divider[xTile:req_xLength, yTile:req_yLength,:] + 1
    return HeatMap_divider,HeatMap_0,WholeSlide_0

def get_inference_from_file(test_filename, stats_dict):
    basename = test_filename
    # remove extension to basename:
    current_score = -1
    score_correction = 0
    oClass = -1
    cmap = plt.get_cmap('binary')
    line = stats_dict[basename]
    NumberOfClasses = len(line)
    class_all = []
    sum_class = 0
    for nC in range(NumberOfClasses):
        class_all.append(float(line[nC]))
        sum_class = sum_class + float(line[nC])
    for nC in range(NumberOfClasses):
        class_all[nC] = class_all[nC] / sum_class
    current_score = max(class_all)
    oClass = class_all.index(max(class_all)) + 1
    score_correction = 1.0 / len(class_all)    
#     '''normal'''
    if oClass == 1:        
#         cmap = sns.light_palette("snow", as_cmap=True)
        cmap = sns.light_palette("snow", as_cmap=True)
#     '''ac'''
    elif oClass == 2:
        cmap = sns.light_palette("deeppink", as_cmap=True)
#     '''g'''
    elif oClass == 3:
#         cmap = sns.light_palette("y", as_cmap=True)
        cmap = sns.light_palette("orange", as_cmap=True)
#     '''ng'''
    elif oClass == 4:
        cmap = sns.light_palette("palegreen", as_cmap=True)
#     '''scc'''
    elif oClass == 5:
        cmap = sns.light_palette("royalblue", as_cmap=True)
#     '''smcc'''
    else:
#         cmap = sns.light_palette("chocolate", as_cmap=True)
        cmap = sns.light_palette("darkviolet", as_cmap=True)
#     print(oClass, current_score, (current_score-score_correction)/(1.0-score_correction))
    return oClass, cmap, (current_score-score_correction)/(1.0-score_correction)

def get_inference_from_file_matplotlib_color(test_filename, stats_dict):
    basename = test_filename
    # remove extension to basename:
    current_score = -1
    score_correction = 0
    oClass = -1
    cmap = plt.get_cmap('binary')
    line = stats_dict[basename]
    NumberOfClasses = len(line)
    class_all = []
    sum_class = 0
    for nC in range(NumberOfClasses):
        class_all.append(float(line[nC]))
        sum_class = sum_class + float(line[nC])
    for nC in range(NumberOfClasses):
        class_all[nC] = class_all[nC] / sum_class
    current_score = max(class_all)
    oClass = class_all.index(max(class_all)) + 1
    score_correction = 1.0 / len(class_all)
    if oClass == 1:
        cmap = plt.get_cmap('binary')
    elif oClass == 2:
        cmap = plt.get_cmap('Reds')
    elif oClass == 3:
        cmap = plt.get_cmap('Wistia')
    elif oClass == 4:
        cmap = plt.get_cmap('Greens')
    elif oClass == 5:
        cmap = plt.get_cmap('Blues')
#         cmap = plt.get_cmap(plt.get_cmap("PiYG", 1))
    else:
        cmap = plt.get_cmap('Purples')
#     print(oClass, current_score, (current_score-score_correction)/(1.0-score_correction))
    return oClass, cmap, (current_score-score_correction)/(1.0-score_correction)

def get_plot(alpha=0.33,divide=True):
    result = get_heatmap(stats_dict=stats_dict,patch_size=256,resample_factor=12,base_path=base_path,divide=divide)
    HeatMap_divider = result[0] * 1.0 + 0.0
    HeatMap_0 = result[1]
    HeatMap_divider[HeatMap_divider == 0] = 1.0
    if divide:
        HeatMap_0 = np.divide(HeatMap_0, HeatMap_divider)
    WholeSlide_0 = result[2]
#     alpha = 0.33
#     alpha = 0.99
    out = HeatMap_0 * 255 * (1.0 - alpha) + WholeSlide_0 * alpha
    out = out.transpose((1, 0, 2))
    out[out==0]=255
#     scipy.misc.imsave(outfile,out)
    return(Image.fromarray(out.astype('uint8')))

matplotlib.use('Agg')
#wholeSlideInferDir = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/tmpResult/wsiInferResult'
#patchPath = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/patch'
#heatmapDir = r'/home/jwang/preciMedWorkFlow/deepIHC/testTiff/heatMap'
#slide_name = 'P24221-7E-2020-05-27_13_12_53'
if not os.path.exists(heatmapDir):  # 创建heatmap目录
    os.makedirs(heatmapDir)
for slide_name in os.listdir(patchPath):
    pred_df_path = os.path.join(wholeSlideInferDir,slide_name + '_result.csv')  # '/home/jwang/preciMedWorkFlow/deepIHC/testTiff/tmpResult/wsiInferResult/P28642-9-2020-05-27_11_56_47_result.csv'
    base_path = os.path.join(patchPath,slide_name)  # '/home/jwang/preciMedWorkFlow/deepIHC/testTiff/patch/P28642-9-2020-05-27_11_56_47'
    stats_dict = dict_tiles_stats(pred_df_path)  # 转为字典，如{'P28642-9-2020-05-27_0_1_0': [0.9953035, 0.0024823523000000003, 9.748552e-05, 0.0016754136, 0.00025055165, 0.00019066907], ……}
    #out = get_plot(alpha=1,divide=True)
    #heatmap_array = np.array(out)
    #heatmap_name = os.path.join(heatmapDir,slide_name+'_tissue.png')
    #imageio.imwrite(heatmap_name,heatmap_array)
    out = get_plot(alpha=0.4,divide=True)
    #heatmap_array = np.array(out)
    #heatmap_name = os.path.join(heatmapDir,slide_name+'_pure_heatmap.png')
    #imageio.imwrite(heatmap_name,heatmap_array)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    frame = plt.gca()
    frame.axes.get_yaxis().set_visible(True)
    frame.axes.get_xaxis().set_visible(True)
    plt.grid(True)
    ax.xaxis.grid(True, which='major') #x坐标轴的网格使用主刻度 
    ax.yaxis.grid(True, which='minor') #x坐标轴的网格使用主刻度
    ax.imshow(out)
    plt.xlabel('Width-x',fontsize=12)
    plt.ylabel('Height-y',fontsize=12)
    plt.title('Heatmap',fontsize=14)
    norm = matplotlib.colors.Normalize(vmin=-1, vmax=1)
    #leftDis = 0.82
    cbar_ax_ac = fig.add_axes([leftDis, 0.8, 0.16, 0.035])
    cbar_ax_scc = fig.add_axes([leftDis, 0.7, 0.16, 0.035])
    cbar_ax_smcc = fig.add_axes([leftDis, 0.6, 0.16, 0.035])
    cbar_ax_g = fig.add_axes([leftDis, 0.5, 0.16, 0.035])
    cbar_ax_ng = fig.add_axes([leftDis, 0.4, 0.16, 0.035])
    cbar_ax_normal = fig.add_axes([leftDis, 0.3, 0.16, 0.035])
    im_ac = ax.imshow(out,cmap = sns.light_palette("deeppink", as_cmap=True))
    im_scc = ax.imshow(out,cmap = sns.light_palette("royalblue", as_cmap=True))
    im_smcc = ax.imshow(out,cmap = sns.light_palette("darkviolet", as_cmap=True))
    im_g = ax.imshow(out,cmap = sns.light_palette("orange", as_cmap=True))
    im_ng = ax.imshow(out,cmap = sns.light_palette("palegreen", as_cmap=True))
    im_normal = ax.imshow(out, cmap = sns.light_palette("snow", as_cmap=True))
    plt.colorbar(im_ac, norm=norm,orientation='horizontal',ticks=[],label='LUAD',ticklocation='top',cax=cbar_ax_ac)
    plt.colorbar(im_scc, norm=norm,orientation='horizontal',ticks=[],label='LUSC',ticklocation='top',cax=cbar_ax_scc)
    plt.colorbar(im_smcc, norm=norm,orientation='horizontal',ticks=[],label='SCLC',ticklocation='top',cax=cbar_ax_smcc)
    plt.colorbar(im_g, norm=norm,orientation='horizontal',ticks=[],label='PTB',ticklocation='top',cax=cbar_ax_g)
    plt.colorbar(im_ng, norm=norm,orientation='horizontal',ticks=[],label='OP',ticklocation='top',cax=cbar_ax_ng)
    plt.colorbar(im_normal,norm=norm,orientation='horizontal',ticks=[],label='NL',ticklocation='top',cax=cbar_ax_normal)
    heatmap_name = os.path.join(heatmapDir,slide_name+'_heatmap_with_legend.png')
    plt.savefig(heatmap_name, bbox_inches = 'tight', format='png', dpi=1200)