import os
import glob
import torch
import random
import numpy as np
from tqdm import tqdm
import torch.nn as nn
from PIL import Image
from SODData import SODData
import torch.optim as optim
from skimage import morphology
import multiprocessing as multi_p
import pydensecrf.densecrf as dcrf
from torchvision import transforms
from alisuretool.Tools import Tools
import torch.backends.cudnn as cudnn
from skimage.io import imread, imsave
from torch.utils.data import DataLoader, Dataset
from torchvision.models import resnet
from torchvision.models._utils import IntermediateLayerGetter
from pydensecrf.utils import unary_from_softmax, unary_from_labels


#######################################################################################################################
# 0 CRF
class CRFTool(object):

    @staticmethod
    def get_ratio(img, ratio_th=0.5):
        black = np.count_nonzero(img < ratio_th) + 1
        white = np.count_nonzero(img > ratio_th) + 1
        ratio = white / (white + black)
        return ratio

    @staticmethod
    def _get_k1_k2(img, black_th=0.25, white_th=0.75, ratio_th=16):
        black = np.count_nonzero(img < black_th) + 1
        white = np.count_nonzero(img > white_th) + 1
        ratio = black / white
        if ratio > 1:  # 物体小
            ratio = int(ratio)
            ratio = ratio_th if ratio > ratio_th else ratio
            k1 = (ratio_th + ratio) // 2  # 多膨胀
            k2 = (ratio_th - ratio) // 2  # 少腐蚀
        else:  # 物体大
            ratio = int(1 / ratio)
            ratio = ratio_th if ratio > ratio_th else ratio
            k1 = (ratio_th - ratio) // 2  # 少膨胀
            k2 = (ratio_th + ratio) // 2  # 多腐蚀
            pass
        return k1, k2

    @classmethod
    def get_uncertain_area(cls, annotation, black_th=0.3, white_th=0.5, ratio_th=16):
        annotation = np.copy(annotation)
        k1, k2 = cls._get_k1_k2(annotation, black_th=black_th, white_th=white_th, ratio_th=ratio_th)
        result1 = morphology.closing(annotation, morphology.disk(3))  # 闭运算：消除噪声
        result2 = morphology.dilation(result1, morphology.disk(k1))  # 膨胀：增加不确定边界
        result3 = morphology.erosion(result1, morphology.disk(k2))  # 腐蚀：增加不确定边界

        other = (black_th + white_th) / 2
        result_annotation = np.zeros_like(annotation) + other
        result_annotation[result2 < black_th] = 0  # black
        result_annotation[result3 > white_th] = 1  # white
        change = result_annotation == other
        return result_annotation, change

    @staticmethod
    def crf(image, annotation, t=5, n_label=2, a=0.1, b=0.9):  # [3, w, h], [1, w, h]
        image = np.ascontiguousarray(image)
        annotation = np.concatenate([annotation, 1 - annotation], axis=0)
        h, w = image.shape[:2]

        d = dcrf.DenseCRF2D(w, h, 2)
        unary = unary_from_softmax(annotation)
        unary = np.ascontiguousarray(unary)
        d.setUnaryEnergy(unary)
        d.addPairwiseGaussian(sxy=3, compat=3)
        d.addPairwiseBilateral(sxy=80, srgb=13, rgbim=np.copy(image), compat=10)
        q = d.inference(t)

        result = np.array(q).reshape((2, h, w))
        return result[0]

    @staticmethod
    def crf_label(image, annotation, t=5, n_label=2, a=0.3, b=0.5):
        image = np.ascontiguousarray(image)
        h, w = image.shape[:2]
        annotation = np.squeeze(np.array(annotation))

        a, b = (a * 255, b * 255) if np.max(annotation) > 1 else (a, b)
        label_extend = np.zeros_like(annotation, dtype=np.int)
        label_extend[annotation >= b] = 2
        label_extend[annotation <= a] = 1
        _, label = np.unique(label_extend, return_inverse=True)

        d = dcrf.DenseCRF2D(w, h, n_label)
        u = unary_from_labels(label, n_label, gt_prob=0.7, zero_unsure=True)
        u = np.ascontiguousarray(u)
        d.setUnaryEnergy(u)
        d.addPairwiseGaussian(sxy=(3, 3), compat=3)
        d.addPairwiseBilateral(sxy=(80, 80), srgb=(13, 13, 13), rgbim=np.copy(image), compat=10)
        q = d.inference(t)
        map_result = np.argmax(q, axis=0)
        result = map_result.reshape((h, w))
        return result

    pass


#######################################################################################################################
# 1 Data


class FixedResized(object):

    def __init__(self, img_w=300, img_h=300):
        self.img_w, self.img_h = img_w, img_h
        pass

    def __call__(self, img, label, image_crf=None, param=None):
        img = img.resize((self.img_w, self.img_h))
        label = label.resize((self.img_w, self.img_h))
        if image_crf is not None:
            image_crf = image_crf.resize((self.img_w, self.img_h))
        return img, label, image_crf, param

    pass


class RandomHorizontalFlip(transforms.RandomHorizontalFlip):
    def __call__(self, img, label, image_crf=None, param=None):
        if random.random() < self.p:
            img = transforms.functional.hflip(img)
            label = transforms.functional.hflip(label)
            if image_crf is not None:
                image_crf = transforms.functional.hflip(image_crf)
            if param is not None:
                param.append(1)
            pass
        else:
            if param is not None:
                param.append(0)
        return img, label, image_crf, param

    pass


class ToTensor(transforms.ToTensor):
    def __call__(self, img, label, image_crf=None, param=None):
        img = super().__call__(img)
        label = super().__call__(label)
        if image_crf is not None:
            image_crf = super().__call__(image_crf)
        return img, label, image_crf, param

    pass


class Normalize(transforms.Normalize):
    def __call__(self, img, label, image_crf=None, param=None):
        img = super().__call__(img)
        return img, label, image_crf, param

    pass


class Compose(transforms.Compose):
    def __call__(self, img, label, image_crf=None, param=None):
        for t in self.transforms:
            img, label, image_crf, param = t(img, label, image_crf, param)
        return img, label, image_crf, param

    pass


class DatasetUSOD(Dataset):

    def __init__(self, img_name_list, lab_name_list, cam_lbl_name_list, his_train_lbl_name_list,
                 his_save_lbl_name_list, size_train=224, label_a=0.2, label_b=0.5, has_crf=False):
        self.img_name_list = img_name_list
        self.tra_lab_name_list = lab_name_list
        self.cam_lbl_name_list = cam_lbl_name_list
        self.his_train_lbl_name_list = his_train_lbl_name_list
        self.his_save_lbl_name_list = his_save_lbl_name_list
        self.transform = Compose([FixedResized(size_train, size_train), RandomHorizontalFlip(), ToTensor(),
                                  Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])

        self.label_a = label_a
        self.label_b = label_b
        self.has_crf = has_crf

        self.lbl_name_list_for_train = None
        self.lbl_name_list_for_save = self.his_save_lbl_name_list

        self.image_size_list = [Image.open(image_name).size for image_name in self.img_name_list]
        Tools.print("DatasetUSOD: size_train={}".format(size_train))
        pass

    def set_label(self, is_supervised, cam_for_train=True):
        Tools.print("DatasetUSOD change label: is_supervised={} cam_for_train={}".format(is_supervised, cam_for_train))
        if is_supervised:
            self.lbl_name_list_for_train = self.tra_lab_name_list
        else:
            self.lbl_name_list_for_train = self.cam_lbl_name_list if cam_for_train else self.his_train_lbl_name_list
            pass
        pass

    def __len__(self):
        return len(self.img_name_list)

    def __getitem__(self, idx):
        assert self.lbl_name_list_for_train is not None

        image = Image.open(self.img_name_list[idx]).convert("RGB")
        if os.path.exists(self.lbl_name_list_for_train[idx]):
            label = Image.open(self.lbl_name_list_for_train[idx]).convert("L")
        else:
            label = Image.fromarray(np.zeros_like(np.asarray(image), dtype=np.uint8)).convert("L")
            pass

        image, label, image_for_crf, param = self.transform(image, label, image, [])

        # 重要参数1: 处理 label
        label_final = torch.zeros_like(label)
        label_final[label >= self.label_b] = 1
        label_final[label <= self.label_a] = 0
        label_final[(self.label_a < label) & (label < self.label_b)] = 255

        return image, label_final, image_for_crf, idx, param

    def save_history(self, history, idx):
        h_path = self.lbl_name_list_for_save[idx]
        h_path = Tools.new_dir(h_path)

        history = np.asarray(np.squeeze(history) * 255, dtype=np.uint8)
        im = Image.fromarray(history).resize(self.image_size_list[idx])
        im.save(h_path)
        pass

    def _crf_one_pool(self, pool_id, epoch, img_name_list, his_save_lbl_name_list, his_train_lbl_name_list):
        for i, (img_name, save_lbl_name, train_lbl_name) in enumerate(zip(
                img_name_list, his_save_lbl_name_list, his_train_lbl_name_list)):
            try:
                img = np.asarray(Image.open(img_name).convert("RGB"))  # 图像
                ann = np.asarray(Image.open(save_lbl_name).convert("L")) / 255  # 训练的输出
                if self.has_crf:
                    # 0.0001
                    # 2_CAM_123_224_256_A5_SFalse_DFalse_224_256_cam_up_norm_C23_crf_History_DieDai_CRF_0.3_0.5_211
                    # 2020-07-31 01:21:26 Test 29 avg mae=0.10352051467412994 score=0.6654844658526717
                    # 2020-07-31 01:24:02 Train 29 avg mae=0.07711194122040814 score=0.8706896792050232
                    # ann_label = CRFTool.crf(img, np.expand_dims(ann, axis=0))
                    # ann = (0.75 * ann + 0.25 * ann_label)

                    # 不确定区域; 对不确定区域进行CRF; 修改不确定区域
                    # 1_Morphology_Train_CAM_123_224_256_A5_SFalse_DFalse_224_256_cam_up_norm_C23_crf_History_DieDai_CRF_0.3_0.5_211
                    # 2020-08-01 13:09:27 Test 27 avg mae=0.11535412584032331 score=0.648697955898359
                    # 2020-08-01 13:11:38 Train 27 avg mae=0.07713918824764815 score=0.8691519831500819
                    # if epoch <= 10:
                    #     ann_label = CRFTool.crf(img, np.expand_dims(ann, axis=0))
                    #     ann = (0.75 * ann + 0.25 * ann_label)
                    # else:
                    #     ann, change = CRFTool.get_uncertain_area(ann, black_th=self.label_a,
                    #                                              white_th=self.label_b, ratio_th=16)
                    #     ann2 = CRFTool.crf_label(img, np.expand_dims(ann, axis=0), a=self.label_a, b=self.label_b)
                    #     ann[change] = ann2[change]
                    #     pass

                    # 2_Morphology_Train_CAM_123_224_256_A5_SFalse_DFalse_224_256_cam_up_norm_C23_crf_History_DieDai_CRF_0.3_0.5_211
                    # 2020-08-01 22:47:48 Test 25 avg mae=0.12813315100613096 score=0.6673295398451972
                    # 2020-08-01 22:49:49 Train 25 avg mae=0.08134926897896962 score=0.872017064360255
                    # if epoch <= 10:
                    #     ann_label = CRFTool.crf(img, np.expand_dims(ann, axis=0))
                    #     ann = (0.75 * ann + 0.25 * ann_label)
                    # else:
                    #     ann, change = CRFTool.get_uncertain_area(ann, black_th=self.label_a,
                    #                                              white_th=self.label_b, ratio_th=10)
                    #     ann2 = CRFTool.crf_label(img, np.expand_dims(ann, axis=0), a=self.label_a, b=self.label_b)
                    #     ann[change] = ann2[change]

                    # 3_Morphology_Train_CAM_123_224_256_A5_SFalse_DFalse_224_256_cam_up_norm_C23_crf_History_DieDai_CRF_0.3_0.5_211
                    # 2020-08-02 01:38:45 Test 15 avg mae=0.1323865410827455 score=0.6618741786757457
                    # 2020-08-02 01:40:51 Train 15 avg mae=0.08220166247338057 score=0.8720659252577703
                    # if epoch <= 2:
                    #     ann_label = CRFTool.crf(img, np.expand_dims(ann, axis=0))
                    #     ann = (0.75 * ann + 0.25 * ann_label)
                    # else:
                    #     ann, change = CRFTool.get_uncertain_area(ann, black_th=self.label_a,
                    #                                              white_th=self.label_b, ratio_th=10)
                    #     ann2 = CRFTool.crf_label(img, np.expand_dims(ann, axis=0), a=self.label_a, b=self.label_b)
                    #     ann[change] = ann2[change]

                    # 4
                    # ann_label = CRFTool.crf(img, np.expand_dims(ann, axis=0))
                    # ratio3, lbl = 0, None
                    # if os.path.exists(train_lbl_name):
                    #     lbl = np.asarray(Image.open(train_lbl_name).convert("L")) / 255  # 训练的标签
                    #     ratio3 = CRFTool.get_ratio(lbl)
                    # ratio1 = CRFTool.get_ratio(ann)
                    # ratio2 = CRFTool.get_ratio(ann_label)
                    # if ratio1 > 0.1 and ratio2 > 0.1 and ratio3 > 0.1 and np.abs(ratio1 - ratio3) < 0.05 and \
                    #         np.abs(ratio1 - ratio2) < 0.05 and np.abs(ratio2 - ratio3) < 0.05:
                    #     ann = 0.5 * lbl + 0.5 * (0.5 * ann + 0.5 * ann_label)
                    # else:
                    #     if np.abs(ratio1 - ratio2) < 0.05:
                    #         ann = 0.5 * ann + 0.5 * ann_label
                    #     else:
                    #         ann = 0.75 * ann + 0.25 * ann_label
                    #     pass

                    # 1_FLoss_Morphology_Train_CAM_123_224_256_A5_SFalse_DFalse_224_256_cam_up_norm_C23_crf_History_DieDai_CRF_0.3_0.5_211
                    # loss=bce+0.1*f 31_train_0.058.pth
                    # 2020-08-04 03:06:01 Test 31 avg mae=0.11399387341170084 score=0.6844463873396547
                    # 2020-08-04 03:08:01 Test 31 avg mae=0.07444360389966856 score=0.8787626041157726

                    # 1_FLoss_Morphology_Train_CAM_123_224_256_A5_SFalse_DFalse_320_320_cam_up_norm_C23_crf_History_DieDai_CRF_0.3_0.5_211
                    # 23_train_0.065.pth
                    # 2020-08-05 01:36:52 Test 23 avg mae=0.10584519406080625 score=0.6931652778949701
                    # 2020-08-05 01:40:45 Test 23 avg mae=0.07252354135906154 score=0.8806913579074058
                    # 3_FLoss_Morphology_Train_CAM_123_224_256_A5_SFalse_DFalse_320_320_cam_up_norm_C23_crf_History_DieDai_CRF_0.3_0.5_211
                    # 29_train_0.062.pth
                    # 2020-08-05 05:43:12 Test 29 avg mae=0.10267482305265915 score=0.7039479408676812
                    # 2020-08-05 05:47:09 Test 29 avg mae=0.07275413212112405 score=0.882560829883874
                    if epoch <= 10:
                        ann_label = CRFTool.crf(img, np.expand_dims(ann, axis=0))
                        ann = (0.75 * ann + 0.25 * ann_label)
                    else:
                        ann, change = CRFTool.get_uncertain_area(ann, black_th=self.label_a,
                                                                 white_th=self.label_b, ratio_th=10)
                        ann2 = CRFTool.crf_label(img, np.expand_dims(ann, axis=0), a=self.label_a, b=self.label_b)
                        ann[change] = ann2[change]
                        pass

                    # 2_FLoss_Morphology_Train_CAM_123_224_256_A5_SFalse_DFalse_224_256_cam_up_norm_C23_crf_History_DieDai_CRF_0.3_0.5_211
                    # 45_train_0.021.pth
                    # 2020-08-04 17:00:15 Test 45 avg mae=0.10881691630042735 score=0.6741292126678432
                    # 2020-08-04 17:03:43 Test 45 avg mae=0.07142853449860757 score=0.8794887318974295
                    # 2_FLoss_Morphology_Train_CAM_123_224_256_A5_SFalse_DFalse_320_320_cam_up_norm_C23_crf_History_DieDai_CRF_0.3_0.5_211
                    # 45_train_0.022.pth
                    # 2020-08-04 20:45:27 Test 45 avg mae=0.0999465329797974 score=0.6952895536245347
                    # 2020-08-04 20:49:32 Test 45 avg mae=0.0697628950389723 score=0.8812714221550307
                    # ann_label = CRFTool.crf(img, np.expand_dims(ann, axis=0))
                    # ann = (0.75 * ann + 0.25 * ann_label)
                    pass

                imsave(Tools.new_dir(train_lbl_name), np.asarray(ann * 255, dtype=np.uint8), check_contrast=False)
            except Exception:
                Tools.print("{} {} {} {}".format(pool_id, epoch, img_name, save_lbl_name))
            pass
        pass

    def crf_dir(self, epoch=0):
        Tools.print("DatasetUSOD crf_dir form {}".format(self.his_save_lbl_name_list[0]))
        Tools.print("DatasetUSOD crf_dir   to {}".format(self.his_train_lbl_name_list[0]))

        pool_num = multi_p.cpu_count()
        pool = multi_p.Pool(processes=pool_num)
        one_num = len(self.img_name_list) // pool_num + 1
        for i in range(pool_num):
            img_name_list = self.img_name_list[one_num*i: one_num*(i+1)]
            his_save_lbl_name_list = self.his_save_lbl_name_list[one_num*i: one_num*(i+1)]
            his_train_lbl_name_list = self.his_train_lbl_name_list[one_num*i: one_num*(i+1)]
            pool.apply_async(self._crf_one_pool, args=(i, epoch, img_name_list,
                                                       his_save_lbl_name_list, his_train_lbl_name_list))
            pass
        pool.close()
        pool.join()

        Tools.print("DatasetUSOD crf_dir OVER")
        pass

    pass


class DatasetEvalUSOD(Dataset):

    def __init__(self, img_name_list, lab_name_list, save_lbl_name_list, size_test=256):
        self.image_name_list = img_name_list
        self.label_name_list = lab_name_list
        self.transform = Compose([FixedResized(size_test, size_test), ToTensor(),
                                  Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])

        self.save_lbl_name_list = save_lbl_name_list
        self.image_size_list = [Image.open(image_name).size for image_name in self.image_name_list]
        pass

    def __len__(self):
        return len(self.image_name_list)

    def __getitem__(self, idx):
        image = Image.open(self.image_name_list[idx]).convert("RGB")
        label = Image.open(self.label_name_list[idx]).convert("L")
        image, label, image_for_crf, _ = self.transform(image, label, image, None)
        return image, label, image_for_crf, idx

    def save_history(self, history, idx):
        h_path = self.save_lbl_name_list[idx]
        h_path = Tools.new_dir(h_path)

        history = np.asarray(np.squeeze(history) * 255, dtype=np.uint8)
        im = Image.fromarray(history).resize(self.image_size_list[idx])
        im.save(h_path)
        pass

    @staticmethod
    def eval_mae(y_pred, y):
        return np.abs(y_pred - y).mean()

    @staticmethod
    def eval_pr(y_pred, y, th_num):
        prec, recall = np.zeros(shape=(th_num,)), np.zeros(shape=(th_num,))
        th_list = np.linspace(0, 1 - 1e-10, th_num)
        for i in range(th_num):
            y_temp = y_pred >= th_list[i]
            tp = (y_temp * y).sum()
            prec[i], recall[i] = tp / (y_temp.sum() + 1e-20), tp / y.sum()
            pass
        return prec, recall

    pass


#######################################################################################################################
# 2 Model


"""
is_supervised_pre_train
2020-08-05 17:29:57 Test 33 avg mae=0.04582420507388141 score=0.8697536064569166
2020-08-05 17:32:57 Test 33 avg mae=0.010893938894679938 score=0.9843851488063832

is_unsupervised_pre_train
2020-08-05 18:28:53 Test 39 avg mae=0.0489034632394237 score=0.8598464423336043
2020-08-05 18:31:45 Test 39 avg mae=0.013386835949495435 score=0.9809714979566246

R50_1_Morphology_Train_CAM_123_224_256_A5_SFalse_DFalse_320_320_cam_up_norm_C23_crf_History_DieDai_CRF_0.3_0.5_211
19_train_0.063.pth
2020-08-06 00:15:24 Test 19 avg mae=0.10735431700269192 score=0.6911668524814638
2020-08-06 00:20:32 Test 19 avg mae=0.0695407475260171 score=0.8907156497510346
"""


class ConvBlock(nn.Module):

    def __init__(self, cin, cout, stride=1, has_relu=True):
        super(ConvBlock, self).__init__()
        self.has_relu = has_relu

        self.conv = nn.Conv2d(cin, cout, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn = nn.BatchNorm2d(cout)
        self.relu = nn.ReLU(inplace=True)
        pass

    def forward(self, x):
        out = self.conv(x)
        out = self.bn(out)
        if self.has_relu:
            out = self.relu(out)
        return out

    pass


class BASNet(nn.Module):

    def __init__(self, is_supervised_pre_train=False,
                 is_unsupervised_pre_train=True, unsupervised_pre_train_path="./pre_model/MoCov2.pth"):
        super(BASNet, self).__init__()

        # -------------Encoder--------------
        backbone = resnet.__dict__["resnet50"](pretrained=is_supervised_pre_train,
                                               replace_stride_with_dilation=[False, True, True])
        if is_unsupervised_pre_train:
            backbone = self.load_unsupervised_pre_train(backbone, unsupervised_pre_train_path)
        return_layers = {'relu': 'e0', 'layer1': 'e1', 'layer2': 'e2', 'layer3': 'e3', 'layer4': 'e4'}
        self.backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)

        # Convert
        self.convert_5 = ConvBlock(2048, 512)
        self.convert_4 = ConvBlock(1024, 512)
        self.convert_3 = ConvBlock(512, 256)
        self.convert_2 = ConvBlock(256, 256)
        self.convert_1 = ConvBlock(64, 128)

        # -------------Decoder-------------
        self.decoder_1_b1 = resnet.BasicBlock(512, 512)  # 40
        self.decoder_1_b2 = resnet.BasicBlock(512, 512)  # 40
        self.decoder_1_c = ConvBlock(512, 512, has_relu=True)  # 40

        self.decoder_2_b1 = resnet.BasicBlock(512, 512)  # 40
        self.decoder_2_b2 = resnet.BasicBlock(512, 512)  # 40
        self.decoder_2_c = ConvBlock(512, 256, has_relu=True)  # 40

        self.decoder_3_b1 = resnet.BasicBlock(256, 256)  # 80
        self.decoder_3_b2 = resnet.BasicBlock(256, 256)  # 80
        self.decoder_3_c = ConvBlock(256, 256, has_relu=True)  # 80

        self.decoder_4_b1 = resnet.BasicBlock(256, 256)  # 160
        self.decoder_4_b2 = resnet.BasicBlock(256, 256)  # 160
        self.decoder_4_c = ConvBlock(256, 128, has_relu=True)  # 160

        self.decoder_5_b1 = resnet.BasicBlock(128, 128)  # 160
        self.decoder_5_b2 = resnet.BasicBlock(128, 128)  # 160
        self.decoder_5_out = nn.Conv2d(128, 1, 3, padding=1, bias=False)  # 160
        pass

    def forward(self, x):
        # -------------Encoder-------------
        feature = self.backbone(x)  # (64, 160), (256, 80), (512, 40), (1024, 40), (2048, 40)
        e0 = self.convert_1(feature["e0"])  # 128
        e1 = self.convert_2(feature["e1"])  # 256
        e2 = self.convert_3(feature["e2"])  # 256
        e3 = self.convert_4(feature["e3"])  # 512
        e4 = self.convert_5(feature["e4"])  # 512

        # -------------Decoder-------------
        d1 = self.decoder_1_b2(self.decoder_1_b1(e4))  # 512 * 40 * 40
        d1_d2 = self._up_to_target(self.decoder_1_c(d1), e3) + e3  # 512 * 40 * 40

        d2 = self.decoder_2_b2(self.decoder_2_b1(d1_d2))  # 512 * 21 * 21
        d2_d3 = self._up_to_target(self.decoder_2_c(d2), e2) + e2  # 512 * 40 * 40

        d3 = self.decoder_3_b2(self.decoder_3_b1(d2_d3))  # 256 * 40 * 40
        d3_d4 = self._up_to_target(self.decoder_3_c(d3), e1) + e1  # 256 * 80 * 80

        d4 = self.decoder_4_b2(self.decoder_4_b1(d3_d4))  # 256 * 80 * 80
        d4_d5 = self._up_to_target(self.decoder_4_c(d4), e0) + e0  # 128 * 160 * 160

        d5 = self.decoder_5_b2(self.decoder_5_b1(d4_d5))  # 128 * 160 * 160
        d5_out = self.decoder_5_out(d5)  # 1 * 160 * 160
        d5_out_sigmoid = torch.sigmoid(d5_out)  # 1 * 160 * 160  # 小输出
        d5_out_up = self._up_to_target(d5_out, x)  # 1 * 320 * 320
        d5_out_up_sigmoid = torch.sigmoid(d5_out_up)  # 1 * 320 * 320  # 大输出

        return_result = {"out": d5_out, "out_sigmoid": d5_out_sigmoid,
                         "out_up": d5_out_up, "out_up_sigmoid": d5_out_up_sigmoid}
        return return_result

    @staticmethod
    def _up_to_target(source, target):
        if source.size()[2] != target.size()[2] or source.size()[3] != target.size()[3]:
            source = torch.nn.functional.interpolate(
                source, size=[target.size()[2], target.size()[3]], mode='bilinear', align_corners=False)
            pass
        return source

    @staticmethod
    def load_unsupervised_pre_train(model, pre_train_path, change_key="module.encoder."):
        pre_train = torch.load(pre_train_path)["model"]
        checkpoint = {key.replace(change_key, ""): pre_train[key] for key in pre_train.keys() if change_key in key}
        result = model.load_state_dict(checkpoint, strict=False)
        if len(result.unexpected_keys) == 0:
            Tools.print("Success Load Unsupervised pre train from {}".format(pre_train_path))
        else:
            Tools.print("Error Load Unsupervised pre train from {}".format(pre_train_path))
            pass
        return model

    pass


#######################################################################################################################
# 3 Runner


class BASRunner(object):

    def __init__(self, batch_size=8, size_train=224, size_test=256, is_f_loss=False,
                 label_a=0.2, label_b=0.5, has_crf=True,
                 tra_img_name_list=None, tra_lbl_name_list=None, tra_data_name_list=None, learning_rate=None,
                 cam_label_dir="../BASNetTemp/cam/CAM_123_224_256", cam_label_name='cam_up_norm_C123',
                 his_label_dir="../BASNetTemp/his/CAM_123_224_256", model_dir="./saved_models/model"):
        self.batch_size = batch_size
        self.size_train = size_train
        self.size_test = size_test
        self.is_f_loss = is_f_loss

        self.label_a = label_a
        self.label_b = label_b
        self.has_crf = has_crf

        # Dataset
        self.model_dir = model_dir
        self.img_name_list = tra_img_name_list
        self.lbl_name_list = tra_lbl_name_list
        self.dataset_name_list = tra_data_name_list
        (self.cam_lbl_name_list, self.his_train_lbl_name_list,
         self.his_save_lbl_name_list) = self.get_tra_img_label_name(self.img_name_list, self.dataset_name_list,
                                                                    cam_label_dir, cam_label_name, his_label_dir)
        self.dataset_sod = DatasetUSOD(
            img_name_list=self.img_name_list, lab_name_list=self.lbl_name_list,
            cam_lbl_name_list=self.cam_lbl_name_list, his_train_lbl_name_list=self.his_train_lbl_name_list,
            his_save_lbl_name_list=self.his_save_lbl_name_list, size_train=self.size_train,
            label_a=self.label_a, label_b=self.label_b, has_crf=self.has_crf)
        self.data_loader_sod = DataLoader(self.dataset_sod, self.batch_size, shuffle=True, num_workers=8)
        self.data_batch_num = len(self.data_loader_sod)

        # Model
        self.net = BASNet()
        self.net = nn.DataParallel(self.net).cuda()
        cudnn.benchmark = True

        # Loss and optimizer
        self.bce_loss = nn.BCELoss().cuda()
        self.learning_rate = [[0, 0.0001], [20, 0.00001]] if learning_rate is None else learning_rate
        self.optimizer = optim.Adam(self.net.parameters(), lr=self.learning_rate[0][1])
        pass

    def _adjust_learning_rate(self, epoch):
        for param_group in self.optimizer.param_groups:
            for lr in self.learning_rate:
                if epoch == lr[0]:
                    learning_rate = lr[1]
                    param_group['lr'] = learning_rate
            pass
        pass

    def load_model(self, model_file_name):
        checkpoint = torch.load(model_file_name)
        # checkpoint = {key: checkpoint[key] for key in checkpoint.keys() if "_c1." not in key}
        self.net.load_state_dict(checkpoint, strict=False)
        Tools.print("restore from {}".format(model_file_name))
        pass

    @staticmethod
    def get_tra_img_label_name( img_name_list, dataset_name_list, cam_label_dir, cam_label_name, his_label_dir):
        cam_lbl_name_list = [os.path.join(
            cam_label_dir, '{}_{}_{}.bmp'.format(dataset_name, os.path.splitext(os.path.basename(
                img_path))[0], cam_label_name)) for img_path, dataset_name in zip(img_name_list, dataset_name_list)]

        his_train_lbl_name_list = [os.path.join(
            his_label_dir, "train", '{}_{}_{}.bmp'.format(dataset_name, os.path.splitext(os.path.basename(
                img_path))[0], cam_label_name)) for img_path, dataset_name in zip(img_name_list, dataset_name_list)]
        his_save_lbl_name_list = [os.path.join(
            his_label_dir, "save", '{}_{}_{}.bmp'.format(dataset_name, os.path.splitext(os.path.basename(
                img_path))[0], cam_label_name)) for img_path, dataset_name in zip(img_name_list, dataset_name_list)]

        Tools.print("train images: {}".format(len(img_name_list)))
        return cam_lbl_name_list, his_train_lbl_name_list, his_save_lbl_name_list

    def all_loss_fusion(self, sod_output, sod_label, ignore_label=255.0):
        positions = sod_label.view(-1, 1) != ignore_label
        loss = self.bce_loss(sod_output.view(-1, 1)[positions], sod_label.view(-1, 1)[positions])
        if self.is_f_loss:
            loss += 0.1 * self.f_loss(sod_output.view(-1, 1)[positions], sod_label.view(-1, 1)[positions])
        return loss

    def all_loss_fusion2(self, sod_output, sod_label, ignore_label=255.0):
        positions = sod_label.view(-1, 1) != ignore_label
        if self.is_f_loss:
            loss_bce = self.f_loss(sod_output.view(-1, 1)[positions], sod_label.view(-1, 1)[positions])
        else:
            loss_bce = self.bce_loss(sod_output.view(-1, 1)[positions], sod_label.view(-1, 1)[positions])
        return loss_bce

    @staticmethod
    def f_loss(sod_output, sod_label):
        tp = torch.sum(sod_output * sod_label)
        fp = torch.sum(sod_output * (1 - sod_label))
        tn = torch.sum((1 - sod_output) * sod_label)
        precision = tp / (tp + fp)
        recall = tp / (tp + tn)
        loss = 1 - (1 + 0.3) * (precision * recall) / (0.3 * precision + recall)
        return loss

    def save_histories(self, histories, indexes):
        for history, index in zip(histories, indexes):
            self.dataset_sod.save_history(idx=int(index), history=np.asarray(history.squeeze()))
        pass

    def train(self, epoch_num=200, start_epoch=0, save_epoch_freq=2,
              is_supervised=False, has_history=False, history_epoch_start=10, history_epoch_freq=10):
        all_loss = 0
        for epoch in range(start_epoch, epoch_num+1):
            Tools.print()
            self._adjust_learning_rate(epoch)
            Tools.print('Epoch:{:03d}, lr={:.5f}'.format(epoch, self.optimizer.param_groups[0]['lr']))

            ###########################################################################
            # 0 准备
            if epoch == 0:  # 0
                self.dataset_sod.set_label(is_supervised=is_supervised, cam_for_train=True)
            elif epoch == history_epoch_start:  # 5
                self.dataset_sod.set_label(is_supervised=is_supervised, cam_for_train=False)

            if not is_supervised and epoch >= history_epoch_start and \
                    (epoch - history_epoch_start) % history_epoch_freq == 0:  # 5, 10, 15
                self.dataset_sod.crf_dir(epoch)
            ###########################################################################

            ###########################################################################
            # 1 训练模型
            all_loss = 0.0
            self.net.train()
            for i, (inputs, targets, image_for_crf, indexes, params) in tqdm(enumerate(self.data_loader_sod),
                                                                             total=self.data_batch_num):
                inputs = inputs.type(torch.FloatTensor).cuda()
                targets = targets.type(torch.FloatTensor).cuda()

                self.optimizer.zero_grad()
                return_m = self.net(inputs)
                sod_output = return_m["out_up_sigmoid"]

                loss = self.all_loss_fusion(sod_output, targets)
                loss.backward()
                self.optimizer.step()
                all_loss += loss.item()

                ##############################################
                if has_history:
                    histories = np.asarray(targets.detach().cpu())
                    sod_output = np.asarray(sod_output.detach().cpu())

                    # 处理翻转
                    history_list, sod_output_list = [], []
                    for index, (history, sod_one) in enumerate(zip(histories, sod_output)):
                        if params[0][index] == 1:
                            history = np.expand_dims(np.fliplr(history[0]), 0)
                            sod_one = np.expand_dims(np.fliplr(sod_one[0]), 0)
                        history_list.append(history)
                        sod_output_list.append(sod_one)
                        pass
                    histories, sod_output = np.asarray(history_list), np.asarray(sod_output_list)

                    # 正式开始
                    self.save_histories(indexes=indexes, histories=sod_output)
                    pass
                ##############################################
                pass
            Tools.print("[E:{:3d}/{:3d}] loss:{:.3f}".format(epoch, epoch_num, all_loss/self.data_batch_num))
            ###########################################################################

            ###########################################################################
            # 2 保存模型
            if (epoch + 1) % save_epoch_freq == 0:
                save_file_name = Tools.new_dir(os.path.join(
                    self.model_dir, "{}_train_{:.3f}.pth".format(epoch, all_loss/self.data_batch_num)))
                torch.save(self.net.state_dict(), save_file_name)

                Tools.print()
                Tools.print("Save Model to {}".format(save_file_name))
                Tools.print()

                ###########################################################################
                # 3 评估模型
                self.eval(self.net, epoch=epoch, is_test=True, batch_size=self.batch_size, size_test=self.size_test)
                self.eval(self.net, epoch=epoch, is_test=False, batch_size=self.batch_size, size_test=self.size_test)
                ###########################################################################
                pass
            ###########################################################################

            ###########################################################################
            # 3 评估模型
            # self.eval(self.net, epoch=epoch, is_test=True, batch_size=self.batch_size, size_test=self.size_test)
            ###########################################################################
            pass

        # Final Save
        save_file_name = Tools.new_dir(os.path.join(
            self.model_dir, "{}_train_{:.3f}.pth".format(epoch_num, all_loss/self.data_batch_num)))
        torch.save(self.net.state_dict(), save_file_name)

        Tools.print()
        Tools.print("Save Model to {}".format(save_file_name))
        Tools.print()
        pass

    @classmethod
    def eval(cls, net, epoch=0, is_test=True, size_test=256, batch_size=16, th_num=25, beta_2=0.3, save_path=None):
        which = "TE" if is_test else "TR"
        data_dir = '/media/ubuntu/4T/ALISURE/Data/DUTS/DUTS-{}'.format(which)
        image_dir, label_dir = 'DUTS-{}-Image'.format(which), 'DUTS-{}-Mask'.format(which)

        # 数据
        img_name_list = glob.glob(os.path.join(data_dir, image_dir, '*.jpg'))
        lbl_name_list = [os.path.join(data_dir, label_dir, '{}.png'.format(
            os.path.splitext(os.path.basename(img_path))[0])) for img_path in img_name_list]

        # EVAL
        cls.eval_by_image_label(net, img_name_list, lbl_name_list, epoch=epoch, size_test=size_test,
                                batch_size=batch_size, th_num=th_num, beta_2=beta_2, save_path=save_path)
        pass

    @staticmethod
    def eval_by_image_label(net, img_name_list, lbl_name_list,
                            epoch=0, size_test=256, batch_size=16, th_num=25, beta_2=0.3, save_path=None):
        # 数据
        save_lbl_name_list = [os.path.join(save_path, '{}.bmp'.format(
            os.path.splitext(os.path.basename(img_path))[0])) for img_path in img_name_list] if save_path else None
        dataset_eval_sod = DatasetEvalUSOD(img_name_list=img_name_list, save_lbl_name_list=save_lbl_name_list,
                                           lab_name_list=lbl_name_list, size_test=size_test)
        data_loader_eval_sod = DataLoader(dataset_eval_sod, batch_size, shuffle=False, num_workers=4)

        # 执行
        avg_mae = 0.0
        avg_prec = np.zeros(shape=(th_num,)) + 1e-6
        avg_recall = np.zeros(shape=(th_num,)) + 1e-6
        net.eval()
        with torch.no_grad():
            for i, (inputs, labels, _, indexes) in tqdm(enumerate(data_loader_eval_sod),
                                                        total=len(data_loader_eval_sod)):
                inputs = inputs.type(torch.FloatTensor)
                inputs = inputs.cuda() if torch.cuda.is_available() else inputs

                now_label = labels.squeeze().data.numpy()
                return_m = net(inputs)

                now_pred = return_m["out_up_sigmoid"].squeeze(dim=1).cpu().data.numpy()

                if save_path:
                    for history, index in zip(now_pred, indexes):
                        dataset_eval_sod.save_history(idx=int(index), history=np.asarray(history.squeeze()))
                        pass
                    pass

                mae = dataset_eval_sod.eval_mae(now_pred, now_label)
                prec, recall = dataset_eval_sod.eval_pr(now_pred, now_label, th_num)

                avg_mae += mae
                avg_prec += prec
                avg_recall += recall
                pass
            pass

        # 结果
        avg_mae = avg_mae / len(data_loader_eval_sod)
        avg_prec = avg_prec / len(data_loader_eval_sod)
        avg_recall = avg_recall / len(data_loader_eval_sod)
        score = (1 + beta_2) * avg_prec * avg_recall / (beta_2 * avg_prec + avg_recall)
        score[score != score] = 0
        Tools.print("Test {} avg mae={} score={}".format(epoch, avg_mae, score.max()))
        pass

    pass


#######################################################################################################################
# 4 Main


"""
dut_dmron_5166 320 2020-08-06 11:34:45 Test 0 avg mae=0.08417008735204892 score=0.7514609826751497
dut_dmron_5166 256 2020-08-06 11:35:37 Test 0 avg mae=0.09943417199907888 score=0.7357637453850571
dut_dmron_5166 240 2020-08-06 11:38:19 Test 0 avg mae=0.08378190836251544 score=0.7527351531616690

msra_b 320 2020-08-06 12:13:52 Test 0 avg mae=0.05882423051367168 score=0.8876016764520471
msra_b 240 2020-08-06 12:12:24 Test 0 avg mae=0.06310220548047793 score=0.8792714795165546

ecssd 400 2020-08-06 12:17:47 Test 0 avg mae=0.09875074808602918 score=0.8380909748089568
ecssd 320 2020-08-06 12:15:25 Test 0 avg mae=0.09106177961855323 score=0.8471667136215534
ecssd 240 2020-08-06 12:15:46 Test 0 avg mae=0.11315359232088405 score=0.8279910151378724
"""


if __name__ == '__main__':
    # os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2, 3"
    # os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1"
    os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"
    # os.environ["CUDA_VISIBLE_DEVICES"] = "2"

    # _size_train, _size_test = 224, 256
    # _batch_size = 6 * len(os.environ["CUDA_VISIBLE_DEVICES"].split(","))

    _size_train, _size_test = 320, 400
    _batch_size = 6 * len(os.environ["CUDA_VISIBLE_DEVICES"].split(","))

    _is_all_data = False
    _is_supervised = False
    _has_history = True
    _is_f_loss = False

    ####################################################################################################
    _history_epoch_start, _history_epoch_freq, _save_epoch_freq = 2, 1, 1
    _label_a, _label_b, _has_crf = 0.3, 0.5, True

    _learning_rate = [[0, 0.0001], [30, 0.00001]]
    _cam_label_dir = "../BASNetTemp/cam/CAM_123_224_256_A5_SFalse_DFalse"
    _cam_label_name = 'cam_up_norm_C23_crf'
    ####################################################################################################

    _name_model = "R50_1{}_Morphology_Train_{}_{}_{}{}{}{}_DieDai{}_{}_{}".format(
        "_FLoss" if _is_f_loss else "", os.path.basename(_cam_label_dir), _size_train, _size_test,
        "_{}".format(_cam_label_name), "_Supervised" if _is_supervised else "", "_History" if _has_history else "",
        "_CRF" if _has_crf else "",  "{}_{}".format(_label_a, _label_b),
        "{}{}{}".format(_history_epoch_start, _history_epoch_freq, _save_epoch_freq))
    _his_label_dir = "../BASNetTemp/R50_his/{}".format(_name_model)

    Tools.print()
    Tools.print(_name_model)
    Tools.print(_cam_label_name)
    Tools.print(_cam_label_dir)
    Tools.print(_his_label_dir)
    Tools.print()

    sod_data = SODData(data_root_path="/media/ubuntu/4T/ALISURE/Data/SOD")
    all_image, all_mask, all_dataset_name = sod_data.get_all_train_and_mask() if _is_all_data else sod_data.duts_tr()

    bas_runner = BASRunner(batch_size=_batch_size, size_train=_size_train, size_test=_size_test, is_f_loss=_is_f_loss,
                           cam_label_dir=_cam_label_dir, cam_label_name=_cam_label_name, his_label_dir=_his_label_dir,
                           label_a=_label_a, label_b=_label_b, has_crf=_has_crf,
                           tra_img_name_list=all_image, tra_lbl_name_list=all_mask,
                           tra_data_name_list=all_dataset_name, learning_rate=_learning_rate,
                           model_dir="../BASNetTemp/saved_models/{}".format(_name_model))

    # TRAIN
    # bas_runner.load_model(model_file_name="../BASNetTemp/saved_models/CAM_123_224_256/930_train_1.172.pth")
    # bas_runner.load_model(model_file_name="../BASNetTemp/saved_models/CAM_123_224_256_DTrue/1000_train_1.072.pth")
    # bas_runner.load_model(model_file_name="../BASNetTemp/saved_models/CAM_123_224_256_DFalse/1000_train_1.154.pth")
    bas_runner.train(epoch_num=40, start_epoch=0, history_epoch_start=_history_epoch_start,
                     history_epoch_freq=_history_epoch_freq, is_supervised=_is_supervised, has_history=_has_history)

    # EVAL
    # _model_name = "CAM_123_SOD_224_256_cam_up_norm_C123_crf_Filter_History_DieDai_CRF"
    # bas_runner.load_model(model_file_name="../BASNetTemp/saved_models/{}/30_train_0.011.pth".format(_model_name))
    # bas_runner.eval(bas_runner.net, epoch=0, is_test=True, batch_size=_batch_size, size_test=_size_test,
    #                 save_path="/media/ubuntu/4T/ALISURE/USOD/BASNetTemp/his/{}/test".format(_model_name))

    # EVAL
    # _model_name = "R50_1_Morphology_Train_CAM_123_224_256_A5_SFalse_DFalse_320_320_cam_up_norm_C23_crf_History_DieDai_CRF_0.3_0.5_211"
    # _model_file = "19_train_0.063.pth"
    # bas_runner.load_model(model_file_name="../BASNetTemp/saved_models/{}/{}".format(_model_name, _model_file))
    # sod_data = SODData(data_root_path="/media/ubuntu/4T/ALISURE/Data/SOD")
    # img_name_list, lbl_name_list, dataset_name_list = sod_data.ecssd()
    # bas_runner.eval_by_image_label(
    #     bas_runner.net, img_name_list, lbl_name_list, epoch=0, batch_size=_batch_size,
    #     size_test=_size_test, save_path="../BASNetTemp/eval/{}/{}/test".format(dataset_name_list[0], _model_name))
    pass
