# -*- coding: utf-8 -*-

"""

@File    : rs_dataset.py

@Time    : 2019/6/22 10:57

@Author  : Parker

@Email   : now_cherish@163.com

@Software: PyCharm

@Des     : data set

"""

import csv

import torch

from torch.utils.data import Dataset

import torchvision.transforms as transforms

import pydicom

import os.path as osp

import os

from PIL import Image

import numpy as np

import random

import cv2

from tqdm import tqdm

import matplotlib.pyplot as plt

import time

from skimage.morphology import remove_small_holes, remove_small_objects

from skimage.measure import label, regionprops

from skimage.filters import threshold_otsu

def data_understanding():

    labels = prepare_label()

    s, ss = {}, {}

    for key, one in tqdm(zip(list(labels.keys()), list(labels.values()))):

        lb = int("".join(map(str, one)), 2)

        if lb not in s.keys():

            s[lb] = []

        s[lb].append(key)

    for one in labels.values():

        for idx, t in enumerate(one):

            if idx not in ss.keys():

                ss[idx] = 0

            if t == 1:

                ss[idx] += 1

    for one in s.keys():

        print(bin(one)[2:].zfill(6), len(s[one]))

def prepare_label():

    labels = ["epidural", "intraparenchymal", "intraventricular",

              "subarachnoid", "subdural", "any"]

    label_ranks = {}

    for i in range(len(labels)):

        label_ranks[labels[i]] = i

    all_true_labels = {}

    with open(osp.join('/media/tiger/zzr/rsna/stage_1_train.csv'), 'r') as fp:

        csv_reader = csv.reader(fp, delimiter=',')

        next(csv_reader, None)

        print('processing data ...')

        for row in tqdm(csv_reader):

            id = "_".join(row[0].split('_')[:2])

            label_id = label_ranks[row[0].split('_')[2]]

            if id not in all_true_labels:

                all_true_labels[id] = [0] * 6

            all_true_labels[id][label_id] = int(row[1])

    return all_true_labels

class RSDataset(Dataset):

    def __init__(self, rootpth='/media/tiger/zzr/rsna', des_size=(512, 512), mode='train'):

        """

        :param rootpth: 根目录

        :param re_size: 数据同一resize到这个尺寸再后处理

        :param crop_size: 剪切

        :param erase: 遮罩比例

        :param mode: train/val/test

        """

        self.root_path = rootpth

        self.des_size = des_size

        self.mode = mode

        self.name = None

        # 处理对应标签

        assert (mode == 'train' or mode == 'val' or mode == 'test')

        labels = ["epidural", "intraparenchymal", "intraventricular",

                  "subarachnoid", "subdural", "any"]

        self.label_ranks = {}

        for i in range(len(labels)):

            self.label_ranks[labels[i]] = i

        self.labels = self.prepare_label()

        # 读取文件名称

        self.file_names = []

        for root,dirs,names in os.walk(osp.join(rootpth, mode)):

            for name in names:

                if name == 'ID_6431af929.dcm':

                    continue

                self.file_names.append(osp.join(root,name))

        # 确定分隔符号

        self.split_char = '\\' if '\\' in self.file_names[0] else '/'

        # totensor 转换n

        self.to_tensor = transforms.Compose([ # 32.98408291578699 33.70147134726827

            transforms.ToTensor(),

            transforms.Normalize(32.98408291578699, 33.70147134726827)

        ])

    def data_loader(self, fname):

        """

        load data

        :param fname:

        :return:

        """

        ds = pydicom.dcmread(fname)

        try:

            windowCenter = int(ds.WindowCenter[0])

            windowWidth = int(ds.WindowWidth[0])

        except:

            windowCenter = int(ds.WindowCenter)

            windowWidth = int(ds.WindowWidth)

        intercept = ds.RescaleIntercept

        slope = ds.RescaleSlope

        data = ds.pixel_array

        data = np.clip(data * slope + intercept, windowCenter - windowWidth / 2, windowCenter + windowWidth / 2).astype(np.float32)

        data = self.preprocess(data)

        return data

    def preprocess(self, data):

        """

        otsu threshold

        :param data:

        :return:

        """

        try:

            thres = threshold_otsu(data)

        except:

            thres = np.min(data)

        data1 = data > thres

        data1 = remove_small_objects(data1)

        label_data = label(data1)

        props = regionprops(label_data)

        area = 0

        bbox = (0, 0, np.shape(data)[0], np.shape(data)[1])

        for idx, i in enumerate(props):

            if i.area > area:

                area = i.area

                bbox = i.bbox

        data1 = data[bbox[0]:bbox[2]+1, bbox[1]:bbox[-1]+1]

        return data1

    def prepare_label(self):

        all_true_labels = {}

        import csv

        with open(osp.join(self.root_path, 'stage_1_train.csv'), 'r') as fp:

            csv_reader = csv.reader(fp, delimiter=',')

            next(csv_reader, None)

            for row in tqdm(csv_reader):

                id = "_".join(row[0].split('_')[:2])

                label_id = self.label_ranks[row[0].split('_')[2]]

                if id not in all_true_labels:

                    all_true_labels[id] = [0] * 6

                all_true_labels[id][label_id] = float(row[1])

        return all_true_labels

    def __getitem__(self, idx):

        self.name = self.file_names[idx]

        category = self.labels[self.name.split(self.split_char)[-1].split('.')[0]]

        img = cv2.resize(self.data_loader(self.name), dsize=self.des_size, interpolation=cv2.INTER_LINEAR)

        # plt.imshow(img)

        # plt.show()

        return self.to_tensor(img), torch.tensor(category)

    def __len__(self):

        return len(self.file_names)

    def calculateMeanStd(self, idx):

        """

        :param idx:

        :return:

        """

        self.name = self.file_names[idx]

        img = self.data_loader(self.name)

        return np.mean(img), np.std(img)

class RSDataset_test(RSDataset):

    def __init__(self, rootpth='/media/tiger/zzr/rsna', des_size=(512, 512), mode='test'):

        super().__init__()

        # 读取文件名称

        self.file_names = []

        for root, dirs, names in os.walk(osp.join(rootpth, mode)):

            for name in names:

                self.file_names.append(osp.join(root, name))

    def __getitem__(self, idx):

        self.name = self.file_names[idx]

        img = cv2.resize(self.data_loader(self.name), dsize=self.des_size, interpolation=cv2.INTER_LINEAR)

        return self.to_tensor(img), self.name.split(self.split_char)[-1].split('.')[0]

    def __len__(self):

        return len(self.file_names)

if __name__ == '__main__':

    data = RSDataset_test()

    for i in tqdm(range(len(data))):

        a, b = data.__getitem__(i)

        print(data.name)

        print(b)

    # mean, std = 0, 0

    # for i in tqdm(range(len(data))):

    #     u, d = data.calculateMeanStd(i)

    #     u /= len(data)

    #     d /= len(data)

    #     mean += u

    #     std += d

    #

    # print(mean, std)