import torch
import torch.nn as nn
import numpy as np
from model.rpn.bbox import bbox_transform_batch, bbox_overlaps_batch
from logger import get_logger

logger = get_logger()

class _ProposalTargetLayer(nn.Module):
    """
    对应GT目标的检测proposals 分类label和bbox回归目标
    选择128个ROIs用以训练
    """

    def __init__(self, num_classes):
        super(_ProposalTargetLayer, self).__init__()
        self._num_classes = num_classes
        # Normalize the targets using "precomputed" (or made up) means and stdevs
        self.BBOX_NORMALIZE_MEANS = torch.FloatTensor((0.0, 0.0, 0.0, 0.0))
        self.BBOX_NORMALIZE_STDS = torch.FloatTensor((0.1, 0.1, 0.2, 0.2))
        # Deprecated (inside weights)
        self.BBOX_INSIDE_WEIGHTS = torch.FloatTensor((1.0, 1.0, 1.0, 1.0))

    def forward(self, all_rois, gt_boxes, num_boxes, batch_size):

        self.batch_size = batch_size
        self.BBOX_NORMALIZE_MEANS = self.BBOX_NORMALIZE_MEANS.type_as(gt_boxes)
        self.BBOX_NORMALIZE_STDS = self.BBOX_NORMALIZE_STDS.type_as(gt_boxes)
        self.BBOX_INSIDE_WEIGHTS = self.BBOX_INSIDE_WEIGHTS.type_as(gt_boxes)

        # 创建一个和gt_boxes形状一致的全零矩阵 [1, 20, 5]
        gt_boxes_append = gt_boxes.new(gt_boxes.size()).zero_()
        gt_boxes_append[:, :, 1:5] = gt_boxes[:, :, :4]  # 将gt_boxes第三维前三列赋值给gt_boxes_append第三维1~4列

        # Include ground-truth boxes in the set of candidate rois
        all_rois = torch.cat([all_rois, gt_boxes_append], 1)  # 如果输入的all_rois=[1,2000,5],输出[1,2020,5]

        num_images = 1
        rois_per_image = int(self.batch_size / num_images)
        # Fraction of minibatch that is labeled foreground
        fg_rois_per_image = int(np.round(0.25 * rois_per_image))  # if batch_size=1, fg_rois_per_image=0
        fg_rois_per_image = 1 if fg_rois_per_image == 0 else fg_rois_per_image
        labels, rois, bbox_targets, bbox_inside_weights = self._sample_rois_pytorch(
            all_rois, gt_boxes, fg_rois_per_image,
            rois_per_image, self._num_classes)

        bbox_outside_weights = (bbox_inside_weights > 0).float()

        return rois, labels, bbox_targets, bbox_inside_weights, bbox_outside_weights

    def backward(self, top, propagate_down, bottom):
        """This layer does not propagate gradients."""
        pass

    def reshape(self, bottom, top):
        """Reshaping happens during the call to forward."""
        pass

    def _get_bbox_regression_labels_pytorch(self, bbox_target_data, labels_batch, num_classes):
        """Bounding box 回归目标(bbox_target_data) 以紧凑的形式 b*N*(类别, tx, ty, tw, th)
        该函数将目标扩展为使用4-of-4*k表示(即：只有一个类具有非零目标)
        Returns:
            bbox_target (ndarray): b x N x 4K blob of regression targets
            bbox_inside_weights (ndarray): b x N x 4K blob of loss weights
        """
        batch_size = labels_batch.size(0)
        rois_per_image = labels_batch.size(1)
        clss = labels_batch
        bbox_targets = bbox_target_data.new(batch_size, rois_per_image, 4).zero_()
        bbox_inside_weights = bbox_target_data.new(bbox_targets.size()).zero_()

        for b in range(batch_size):
            # assert clss[b].sum() > 0
            if clss[b].sum() == 0:
                continue
            inds = torch.nonzero(clss[b] > 0).view(-1)
            for i in range(inds.numel()):
                ind = inds[i]
                bbox_targets[b, ind, :] = bbox_target_data[b, ind, :]
                bbox_inside_weights[b, ind, :] = self.BBOX_INSIDE_WEIGHTS

        return bbox_targets, bbox_inside_weights

    def _compute_targets_pytorch(self, ex_rois, gt_rois):
        """计算单张图片的bounding box 回归目标"""

        assert ex_rois.size(1) == gt_rois.size(1)
        assert ex_rois.size(2) == 4
        assert gt_rois.size(2) == 4

        batch_size = ex_rois.size(0)
        rois_per_image = ex_rois.size(1)

        targets = bbox_transform_batch(ex_rois, gt_rois)

        # Optionally normalize targets by a precomputed mean and stdev
        targets = ((targets - self.BBOX_NORMALIZE_MEANS.expand_as(targets))
                   / self.BBOX_NORMALIZE_STDS.expand_as(targets))

        return targets

    def _sample_rois_pytorch(self, all_rois, gt_boxes, fg_rois_per_image, rois_per_image, num_classes):
        """Generate a random sample of RoIs comprising foreground and background
        examples.
        """
        # all_rois=[1,2000,5], gt_boxes=[1,20,5], fg_rois_per_image=1,
        # rois_per_image=1, self._num_classes=21

        # overlaps: (rois x gt_boxes)
        overlaps = bbox_overlaps_batch(all_rois, gt_boxes)
        logger.debug("proposal_target_layer->_sample_rois_pytorch: overlaps = {}".format(overlaps.size()))

        max_overlaps, gt_assignment = torch.max(overlaps, 2)
        logger.debug("proposal_target_layer->_sample_rois_pytorch: max_overlaps = {}, "
              "gt_assignment = {}".format(overlaps.size(), gt_assignment))

        batch_size = overlaps.size(0)
        num_proposal = overlaps.size(1)
        num_boxes_per_img = overlaps.size(2)

        offset = torch.arange(0, batch_size) * gt_boxes.size(1)  # offset=tensor([0])
        offset = offset.view(-1, 1).type_as(gt_assignment) + gt_assignment
        logger.debug("proposal_target_layer->_sample_rois_pytorch: offset = {}, "
              "offset.view(-1)={}".format(offset, offset.view(-1)))

        # labels = gt_boxes[:, :, 4].contiguous().view(-1).index(offset.view(-1)) \
        #     .view(batch_size, -1)
        labels = gt_boxes[:, :, 4].contiguous().view(-1)[offset.view(-1)] \
            .view(batch_size, -1)
        logger.debug("labels = {}".format(labels))

        labels_batch = labels.new(batch_size, rois_per_image).zero_()
        rois_batch = all_rois.new(batch_size, rois_per_image, 5).zero_()
        gt_rois_batch = all_rois.new(batch_size, rois_per_image, 5).zero_()
        # Guard against the case when an image has fewer than max_fg_rois_per_image
        # foreground RoIs
        for i in range(batch_size):
            # ROI判定为foreground的阈值
            fg_inds = torch.nonzero(max_overlaps[i] >= 0.5).view(-1)
            fg_num_rois = fg_inds.numel()

            # Select background RoIs as those within [BG_THRESH_LO, BG_THRESH_HI)
            bg_inds = torch.nonzero((max_overlaps[i] < 0.5) &
                                    (max_overlaps[i] >= 0.1)).view(-1)
            bg_num_rois = bg_inds.numel()

            if fg_num_rois > 0 and bg_num_rois > 0:
                # sampling fg
                fg_rois_per_this_image = min(fg_rois_per_image, fg_num_rois)

                # torch.randperm seems has a bug on multi-gpu setting that cause the segfault.
                # See https://github.com/pytorch/pytorch/issues/1868 for more details.
                # use numpy instead.
                # rand_num = torch.randperm(fg_num_rois).long().cuda()
                rand_num = torch.from_numpy(np.random.permutation(fg_num_rois)).type_as(gt_boxes).long()
                fg_inds = fg_inds[rand_num[:fg_rois_per_this_image]]

                # sampling bg
                bg_rois_per_this_image = rois_per_image - fg_rois_per_this_image

                # Seems torch.rand has a bug, it will generate very large number and make an error.
                # We use numpy rand instead.
                # rand_num = (torch.rand(bg_rois_per_this_image) * bg_num_rois).long().cuda()
                rand_num = np.floor(np.random.rand(bg_rois_per_this_image) * bg_num_rois)
                rand_num = torch.from_numpy(rand_num).type_as(gt_boxes).long()
                bg_inds = bg_inds[rand_num]

            elif fg_num_rois > 0 and bg_num_rois == 0:
                # sampling fg
                # rand_num = torch.floor(torch.rand(rois_per_image) * fg_num_rois).long().cuda()
                rand_num = np.floor(np.random.rand(rois_per_image) * fg_num_rois)
                rand_num = torch.from_numpy(rand_num).type_as(gt_boxes).long()
                fg_inds = fg_inds[rand_num]
                fg_rois_per_this_image = rois_per_image
                bg_rois_per_this_image = 0
            elif bg_num_rois > 0 and fg_num_rois == 0:
                # sampling bg
                # rand_num = torch.floor(torch.rand(rois_per_image) * bg_num_rois).long().cuda()
                rand_num = np.floor(np.random.rand(rois_per_image) * bg_num_rois)
                rand_num = torch.from_numpy(rand_num).type_as(gt_boxes).long()

                bg_inds = bg_inds[rand_num]
                bg_rois_per_this_image = rois_per_image
                fg_rois_per_this_image = 0
            else:
                raise ValueError("bg_num_rois = 0 and fg_num_rois = 0, this should not happen!")

            # The indices that we're selecting (both fg and bg)
            keep_inds = torch.cat([fg_inds, bg_inds], 0)

            # Select sampled values from various arrays:
            labels_batch[i].copy_(labels[i][keep_inds])

            # Clamp labels for the background RoIs to 0
            if fg_rois_per_this_image < rois_per_image:
                labels_batch[i][fg_rois_per_this_image:] = 0

            rois_batch[i] = all_rois[i][keep_inds]
            rois_batch[i, :, 0] = i

            gt_rois_batch[i] = gt_boxes[i][gt_assignment[i][keep_inds]]

        bbox_target_data = self._compute_targets_pytorch(
            rois_batch[:, :, 1:5], gt_rois_batch[:, :, :4])

        bbox_targets, bbox_inside_weights = \
            self._get_bbox_regression_labels_pytorch(bbox_target_data, labels_batch, num_classes)

        return labels_batch, rois_batch, bbox_targets, bbox_inside_weights


#
# @lc app=leetcode.cn id=77 lang=python3
#
# [77] 组合
#感觉应该用递归的方法

# @lc code=start
from typing import List

class Solution:
    def combine(self, n: int, k: int) -> List[List[int]]:
        res = []

        def dfs(num:int,target,res_list):

            if target == 0:
                res.append(res_list)
            for i in range(num,n+1):

                dfs(i+1,target-1,res_list+[i])

        dfs(1,k,[])
        return res

        # @lc code=end


from knn.chart import Chart
from collections import Counter
from knn.knn_core import KnnCore
from knn.data_preparation import DataPreparation

if __name__ == '__main__':

    # ************************ DATA INITIALIZATION ************************ #

    data_preparation = DataPreparation(50, None)
    learnset_data, learnset_labels = data_preparation.create_learn_data()
    testset_data, testset_labels = data_preparation.create_test_data()
    print(len(learnset_data))
    print(len(testset_data))

    # ******************************* TESTS ******************************* #

    compatibility = Counter()
    print('Probki danych testowych na podstawie danych treningowych z glosowaniem: \n')
    for i in range(data_preparation.n_training_samples):
        neighbours = KnnCore.get_neighbours(learnset_data, learnset_labels, testset_data[i], 8, KnnCore.distance)
        print('{: >1} {: >3} {: >1} {: >50} {: >1} {: >1} {: >1} {: >1} {: >1} {: >1}'.format(
            'index: ', i, ', result of vote: ', str(KnnCore.vote_harmonic_weights(neighbours)),
            ', test label: ', str(testset_labels[i]), ', test data: ', str(testset_data[i]),
            ', compatibility: ', (KnnCore.vote(neighbours)[0] == str(testset_labels[i]))
        ))
        compatibility[KnnCore.vote(neighbours)[0] == testset_labels[i]] += 1
    compatibility_percentage = compatibility[True] / sum(compatibility.values())
    print('Compatibility percentage is: ', compatibility_percentage * 100, '%')
    chart = Chart()
    chart.gather_data(testset_data, testset_labels)
    chart.create_plot().show()


"""
This plots an f1-score tree for the DNN model of the project discovery/machine learning paper.
The code depends on networkx as well as matplotlib and numpy.

To run, make sure that the files
tree_level0-{1,2,3,4}, tree_level1-{1,2,3,4}, tree_level2-{1,2,3,4} and hpa_gold_standard_tree.csv is in the same folder
as this program and then run
python3 tree_plot.py

The program will output a pdf called tree_output5.pdf that will contain the result tree.
"""
import textwrap
import logging
import tree_utils
import matplotlib.colors
import csv
import networkx as nx
import matplotlib.pyplot as plt
import pickle
import collections
import numpy as np
import copy
from collections import OrderedDict

CUTOFF_VALUES = collections.defaultdict(lambda: 0.4, {
    # H0
    'Aggresome': 0.060,
    'Cell Junctions': 0.098,
    'Centrosome': 0.094,
    'Cytoplasm': 0.437,
    'Cytoskeleton (Actin filaments)': 0.134,
    'Cytoskeleton (Cytokinetic bridge)': 0.004,
    'Cytoskeleton (Intermediate filaments)': 0.085,
    'Cytoskeleton (Microtubule end)': 0.070,
    'Cytoskeleton (Microtubules)': 0.165,
    'Endoplasmic reticulum': 0.185,
    'Focal Adhesions': 0.126,
    'Golgi apparatus': 0.181,
    'Microtubule organizing center': 0.048,
    'Mitochondria': 0.264,
    'Negative': 0.306,
    'Nuclear bodies': 0.141,
    'Nuclear membrane': 0.125,
    'Nuclear speckles': 0.161,
    'Nucleoli (Fibrillar center)': 0.168,
    'Nucleoli': 0.183,
    'Nucleoplasm': 0.392,
    'Nucleus': 0.238,
    'Plasma membrane': 0.278,
    'Unspecific': 0.443,
    'Vesicles': 0.271,
    'NULL': 1.5,

    # H1
    'Cytoplasmic structures': 0.427,
    'Cytoskeleton_1': 0.213,
    'MTOC': 0.093,
    'Neg/Unspec': 0.307,
    'Nuclear bodies/speckles': 0.172,
    'Nuclear membrane_1': 0.129,
    'Nucleoli_1': 0.201,
    'Nucleus_1': 0.500,
    'Periphery': 0.282,
    'Secretory': 0.386,

    # H2
    'Cytoplasmic': 0.516,
    'Nuclear': 0.478,
    'Cytoskeleton': 0.216,
    'Unspecific/Negative': 0.296,
    })

# Hierarchy is defined as a string containing numbers
# Highest number indicates type, Nucleus:0, Cytoplasm:1, or Secretory:2
# Second highest indicates category within type
# Third indicates class within category
DEFAULT_HIERARCHY = {
        'ROOT': '',

        'Nuclear': '1',

        'Nucleus_1': '10',
        'Nucleus': '101',
        'Nucleoplasm': '102',

        'Nuclear bodies/speckles': '11',
        'Nuclear bodies': '111',
        'Nuclear bodies (many)': '112',
        'Nuclear speckles': '113',

        'Nucleoli_1': '12',
        'Nucleoli': '121',
        'Nucleoli (Fibrillar center)': '123',
        'Nucleoli (fib center)': '123',
        'Nucleoli (rim)': '124',

        'Nuclear membrane_1': '13',
        'Nuclear membrane': '131',

        'Cytoplasmic': '2',

        'Cytoplasmic structures': '20',
        'Cytoplasm': '201',
        'Aggresome': '202',
        'Mitochondria': '203',

        'Cytoskeleton_1': '21',
        'Intermediate filaments': '211',
        'Cytoskeleton (Intermediate filaments)': '211',
        'Microtubule end': '212',
        'Cytoskeleton (Microtubule end)': '212',
        'Microtubules': '213',
        'Cytoskeleton (Microtubules)': '213',
        'Actin filaments': '214',
        'Cytoskeleton (Actin filaments)': '214',
        'Cytokinetic bridge': '215',
        'Cytoskeleton (Cytokinetic bridge)': '215',
        'Cytoskeleton': '',

        'MTOC': '22',
        'Microtubule organizing center': '221',
        'Centrosome': '222',

        'Secretory': '23',
        'Endoplasmic reticulum': '231',
        'Golgi apparatus': '232',
        'Vesicles': '233',
        'Lipid droplets': '234',

        'Periphery': '24',
        'Cell Junctions': '240',
        'Focal Adhesions': '241',
        'Plasma membrane': '242',

        'Unspecific/Negative': '9',
        'Neg/Unspec': '90',
        'Negative': '901',
        'Unspecific': '902',
        'NULL': '903',

        'R&R': '',
        'CCD': '',
        'Periphery_1': '',

        'CCV': '',
        'CCV_2': '',
        ' ': '',
        }
DISP_HIERARCHY = {
        '': 'ROOT',

        '1': 'Nuclear',

        '10': 'Nucleus',
        '101': 'Nucleus',
        '102': 'Nucleoplasm',

        '11': 'Subnuclear structures',
        '111': 'Nuclear bodies',
        '112': 'Nuclear bodies (many)',
        '113': 'Nuclear speckles',

        '12': 'Nucleoli',
        '121': 'Nucleoli',
        '123': 'Nucleoli (Fibrillar center)',
        '124': 'Nucleoli (rim)',

        '13': 'Nuclear membrane',
        '131': 'Nuclear membrane',

        '2': 'Cytoplasmic',

        '20':  'Cytoplasm',
        '201': 'Cytosol',
        '202': 'Aggresome',
        '203': 'Mitochondria',

        '21': 'Cytoskeleton',
        '211': 'Intermediate filaments',
        '212': 'Microtubule end',
        '213': 'Microtubules',
        '214': 'Actin filaments',
        '215': 'Cytokinetic bridge',
        '': 'Cytoskeleton',

        '22': 'MTOC',
        '221': 'Microtubule organizing center',
        '222': 'Centrosome',

        '23': 'Secretory',
        '231': 'Endoplasmic reticulum',
        '232': 'Golgi apparatus',
        '233': 'Vesicles',
        '234': 'Lipid droplets',

        '24': 'Cell Periphery',
        '240': 'Cell Junctions',
        '241': 'Focal Adhesions',
        '242': 'Plasma membrane',

        '': 'Unspecific/Negative',
        '': 'Neg/Unspec',
        '': 'Negative',
        '': 'Unspecific',
        '': 'NULL',

        '': 'R&R',
        '': 'CCD',
        '': 'Periphery_1',

        '': 'CCV',
        '': 'CCV_2',

        }

skip_classes_in_output_l0 = {4, 11, 14}

def read_single_fov_counts_layer(preds, plates, classes, skip_classes, location_string):
    counts = collections.defaultdict(lambda: collections.defaultdict(int))
    for fov, fov_dict in preds.items():
        fov = '_'.join(fov.split('_')[:-1])
        fov_prediction = fov_dict['prediction']
        fov_prediction = [fov_prediction[oc] for oc in range(len(fov_prediction)) if oc not in skip_classes]
        fov_prediction = np.asarray(fov_prediction)

        fov_locations = set(plates[fov][location_string].split(','))
        fov_classes = set()
        for i in range(int(len(classes)/2)):
            c = classes[i]
            if c not in CUTOFF_VALUES:
                print(c, 'not in cutoff values. Using 0.4')
            if fov_prediction[i] >= CUTOFF_VALUES[c]:
                fov_classes.add(c)

        for c in fov_classes:
            code = DEFAULT_HIERARCHY[c]
            if c in fov_locations:
                counts[code]['tp'] += 1
            else:
                counts[code]['fp'] += 1

        for c in fov_locations:
            code = DEFAULT_HIERARCHY[c]
            if c not in fov_classes:
                counts[code]['fn'] += 1
    return counts

def get_if_info_classes(if_info, locations_header):
    classes = set()
    for p in if_info:
        unsplit = if_info[p][locations_header]
        split = unsplit.split(',')
        classes.update(split)
    if locations_header == 'locations_h1':
        classes.add('NULL')
    if locations_header == 'locations_h2':
        classes.add('Neg/Unspec')
    classes = sorted(classes)
    ordered_classes = collections.OrderedDict()
    for i, c in enumerate(classes):
        if isinstance(c, int):
            logging.critical('Class CANNOT be an int')
            raise ValueError('{} is int when it should be str'.format(c))

        ordered_classes[c] = i
        ordered_classes[i] = c
    return ordered_classes

def read_if_file(if_file, include={}, non_include={}, identifiers=['if_plate_id', 'position', 'sample'],
                 split_to_well=True):
    dictreader = csv.DictReader(open(if_file))
    if_info = dict()
    for line in dictreader:
        skipping = False

        for constraint in include:
            actual_line = line[constraint]
            actual_line = actual_line.split(',')

            any_in = [actual in include[constraint] for actual in actual_line]
            if not any(any_in):
                skipping = True
                break

        if skipping:
            continue

        for constraint in non_include:
            actual_line = line[constraint]
            value = actual_line.split(',')

            any_in = [actual in non_include[constraint] for actual in value]
            if any(any_in):
                skipping = True
                break

        if skipping:
            continue

        id_ = ''
        for i in identifiers:
            id_ += line[i] + '_'

        if split_to_well:
            id_ = '_'.join(id_.split('_')[:-2])
        else:
            id_ = '_'.join(id_.split('_')[:-1])
        if_info[id_] = line
    return if_info

if __name__ == '__main__':
    counts = collections.defaultdict(lambda: collections.defaultdict(int))

    print('H0')
    reader = read_if_file('hpa_gold_standard_tree.csv')
    classes = get_if_info_classes(reader, 'locations')
    for f in ['tree_level0-1', 'tree_level0-2', 'tree_level0-3', 'tree_level0-4']:
        h0_preds = pickle.load(open(f, 'rb'))
        h0_counts = read_single_fov_counts_layer(h0_preds, reader, classes, skip_classes_in_output_l0, 'locations')
        for c in h0_counts:
            counts[c]['tp'] += h0_counts[c]['tp']
            counts[c]['fp'] += h0_counts[c]['fp']
            counts[c]['fn'] += h0_counts[c]['fn']

    print('H1')
    reader = read_if_file('hpa_gold_standard_tree.csv')
    classes = get_if_info_classes(reader, 'locations_h1')
    for f in ['tree_level1-1', 'tree_level1-2', 'tree_level1-3', 'tree_level1-4']:
        h1_preds = pickle.load(open(f, 'rb'))
        h1_counts = read_single_fov_counts_layer(h1_preds, reader, classes, [], 'locations_h1')
        for c in h1_counts:
            if c == DEFAULT_HIERARCHY['Nucleus_1']:
                print(c, counts[c]['tp'])
            counts[c]['tp'] += h1_counts[c]['tp']
            counts[c]['fp'] += h1_counts[c]['fp']
            counts[c]['fn'] += h1_counts[c]['fn']

    print('H2')
    reader = read_if_file('hpa_gold_standard_tree.csv')
    classes = get_if_info_classes(reader, 'locations_h2')
    for f in ['tree_level2-1', 'tree_level2-2', 'tree_level2-3', 'tree_level2-4']:
        h2_preds = pickle.load(open(f, 'rb'))
        h2_counts = read_single_fov_counts_layer(h2_preds, reader, classes, [], 'locations_h2')
        for c in h2_counts:
            counts[c]['tp'] += h2_counts[c]['tp']
            counts[c]['fp'] += h2_counts[c]['fp']
            counts[c]['fn'] += h2_counts[c]['fn']

    precisions = {}
    recalls = {}
    lengths = {}
    f1s = {}
    print('making precision_recalls')
    for c in counts:
        tp = counts[c]['tp']
        fp = counts[c]['fp']
        fn = counts[c]['fn']
        precisions[c] = tp/((tp + fp) or 1)
        recalls[c] = tp/((tp + fn) or 1)
        lengths[c] = tp + fn
        f1s[c] = 2 * (precisions[c] * recalls[c])/(
                (precisions[c] + recalls[c]) or 1)

    G = nx.DiGraph()
    nodes = []
    node_sizes = []
    edges = []
    edge_colors = []
    node_colors = []

    labels = {k: v for (k, v) in DISP_HIERARCHY.items()}

    norm = matplotlib.colors.Normalize(0.3, 1.0)
    cmap = plt.cm.ScalarMappable(norm=norm, cmap=plt.cm.plasma)
    print('sorting')
    sort_f1s = OrderedDict(sorted(f1s.items()))
    sort_recalls = OrderedDict(sorted(recalls.items()))
    sort_precisions = OrderedDict(sorted(precisions.items()))
    sort_lengths = OrderedDict(sorted(lengths.items()))
    label_sizes = {}
    for code in sort_precisions:
        if code == '' or code == '9' or code == '90' or code == '901' or code == '902' or code == '903':
            continue
        G.add_edge(code[0:-1], code)
        edges.append((code[0:-1], code))
        edge_colors.append(plt.cm.viridis(sort_recalls[code]))

        node_sizes.append(max(300, np.sqrt(sort_lengths[code]) * 12.7))
        label_sizes[code] = node_sizes[-1]
        print(code)
        nodes.append(code)
        node_colors.append(cmap.to_rgba(sort_f1s[code]))

    pos = tree_utils.hierarchical_tree_pos(G, '')
    print(pos)
    sort_edges = sorted(edges)
    fig, ax = tree_utils.draw_hierarchical_tree(
            G, pos, nodes=nodes, node_size=node_sizes, node_color=node_colors,
            edges=sort_edges, edge_color='k', sm=cmap)
    print(node_sizes)

    labelpos = copy.deepcopy(pos)
    f_tex = "{:.2f}"
    l1 = len('(fibrillar center)')
    l2 = len('bodies/speckles')
    for c in labelpos:
        if c == '':
            continue
        originalpos = labelpos[c]
        if c in f1s:
            if f1s[c] <= 0.45:
                color = 'white'
            else:
                color = 'black'
            ax.text(originalpos[0]-0.010, originalpos[1]-0.0025,
                    f_tex.format(f1s[c]), fontsize=8, color=color)
            if len(c) > 2:
                ax.text(originalpos[0], originalpos[1]-0.032,
                        '\n'.join(textwrap.wrap(labels[c], l1)), rotation=-40, fontsize=8,
                        fontdict={'family': 'Sans-serif'})
            else:
                ax.text(originalpos[0]+0.0250, originalpos[1]-0.005,
                        '\n'.join(textwrap.wrap(labels[c], l2)), rotation=0, fontsize=8,
                        fontdict={'family': 'Sans-serif'})

    fig.savefig('tree_output5.pdf', dpi=100, transparent=True)


####################################################################
# The reconstruction of 2D uPTI                                    #
# This reconstruction is based on the uPTI paper                   #
# (https://www.biorxiv.org/content/10.1101/2020.12.15.422951v1)    #
#  ```L.-H. Yeh, I. E. Ivanov, B. B. Chhun, S.-M. Guo, E. Hashemi, #
#  J. R. Byrum, J. A. Pérez-Bermejo, H. Wang, Y. Yu,               #
#  P. G. Kazansky, B. R. Conklin, M. H. Han, and S. B. Mehta,      #
#  "uPTI: uniaxial permittivity tensor imaging of intrinsic        #
#  density and anisotropy," bioRxiv 2020.12.15.422951 (2020).```   #
####################################################################
import numpy as np
import matplotlib.pyplot as plt
from numpy.fft import fftshift
from waveorder import (
    optics,
    waveorder_reconstructor,
    visual,
)

## Initialization
## Load simulated images and parameters

file_name = "./PTI_simulation_data_NA_det_147_NA_illu_140_2D_spoke_discrete_no_1528_ne_1553_no_noise_Born.npz"

array_loaded = np.load(file_name)
list_of_array_names = sorted(array_loaded)

for array_name in list_of_array_names:
    globals()[array_name] = array_loaded[array_name]

print(list_of_array_names)
I_meas = I_meas[1:]
print(I_meas.shape)

### Setup reconstructor class
_, _, N, M, L = I_meas.shape
cali = False
bg_option = "global"
use_gpu = False
gpu_id = 0

A_matrix = 0.5 * np.array([[1, 1, 0], [1, 0, 1], [1, -1, 0], [1, 0, -1]])

setup = waveorder_reconstructor.waveorder_microscopy(
    (N, M),
    lambda_illu,
    ps,
    NA_obj,
    NA_illu,
    z_defocus,
    chi,
    n_media=n_media,
    cali=cali,
    A_matrix=A_matrix,
    inc_recon="3D",
    illu_mode="Arbitrary",
    Source=Source_cont,
    Source_PolState=Source_PolState,
    pad_z=5,
    use_gpu=use_gpu,
    gpu_id=gpu_id,
)

### Illumination patterns used
visual.plot_multicolumn(
    fftshift(Source_cont, axes=(1, 2)), origin="lower", num_col=5, size=5
)
plt.show()

## Reconstruct Stokes images and visualize them as a function of illumination pattern
S_image_recon = setup.Stokes_recon(I_meas)

S_bg_mean_0 = np.mean(S_image_recon[0, :, :, :, :], axis=(1, 2, 3))[
    :, np.newaxis, np.newaxis, np.newaxis
]
S_bg_mean_1 = np.mean(S_image_recon[1, :, :, :, :], axis=(1, 2, 3))[
    :, np.newaxis, np.newaxis, np.newaxis
]
S_bg_mean_2 = np.mean(S_image_recon[2, :, :, :, :], axis=(1, 2, 3))[
    :, np.newaxis, np.newaxis, np.newaxis
]

S_image_tm = np.zeros_like(S_image_recon)
S_image_tm[0] = S_image_recon[0] / S_bg_mean_0 - 1
S_image_tm[1] = (
    S_image_recon[1] / S_bg_mean_0
    - S_bg_mean_1 * S_image_recon[0] / S_bg_mean_0**2
)
S_image_tm[2] = (
    S_image_recon[2] / S_bg_mean_0
    - S_bg_mean_2 * S_image_recon[0] / S_bg_mean_0**2
)

## 3D uPTI reconstruction
### 3D volumes of the components of scattering potential tensor

# Regularization parameters.

# reg_inc for 3D data
# reg_inc = np.array([1, 1, 5e1, 5e1, 2.5e1, 2.5e1, 5e1])*1
#####################

# reg_inc for 2D data
reg_inc = np.array([1, 1, 5e1, 5e1, 5e1, 5e1, 5e1]) * 1
#####################

reg_ret_pr = 1e-2
f_tensor = setup.scattering_potential_tensor_recon_3D_vec(
    S_image_tm, reg_inc=reg_inc, cupy_det=True
)

visual.plot_multicolumn(
    f_tensor[..., L // 2],
    num_col=4,
    origin="lower",
    size=5,
    set_title=True,
    titles=[
        r"$f_{0r}$",
        r"$f_{0i}$",
        r"$f_{1c}$",
        r"$f_{1s}$",
        r"$f_{2c}$",
        r"$f_{2s}$",
        r"$f_{3}$",
    ],
)
plt.show()

### 3D volumes of principal retardance, 3D orientation, and optic sign

# reconstruct 3D anisotropy (principal retardance, 3D orientation, optic sign probability)
# material type:
# "positive" -> only solution of positively uniaxial material
# "negative" -> only solution of negatively uniaxial meterial
# "unknown" -> both solutions of positively and negatively uniaxial material + optic sign estimation

(
    retardance_pr,
    azimuth,
    theta,
    mat_map,
) = setup.scattering_potential_tensor_to_3D_orientation(
    f_tensor,
    S_image_tm,
    material_type="unknown",
    reg_ret_pr=reg_ret_pr,
    itr=20,
    step_size=0.1,
    fast_gpu_mode=True,
)
plt.show()

# scaling to the physical properties of the material

# optic sign probability
p_mat_map = optics.optic_sign_probability(mat_map, mat_map_thres=0.2)

# absorption and phase
phase = optics.phase_inc_correction(f_tensor[0], retardance_pr[0], theta[0])
absorption = f_tensor[1].copy()
phase_PT, absorption_PT, retardance_pr_PT = [
    optics.unit_conversion_from_scattering_potential_to_permittivity(
        SP_array, lambda_illu, n_media=n_media, imaging_mode="3D"
    )
    for SP_array in [phase, absorption, retardance_pr]
]

# clean up GPU memory leftorver

# import gc
# import cupy as cp

# gc.collect()
# cp.get_default_memory_pool().free_all_blocks()

## Visualize reconstructed physical properties of simulated sample
### Reconstructed phase, absorption, principal retardance, azimuth, and inclination assuming (+) and (-) optic sign

# browse the reconstructed physical properties
visual.plot_multicolumn(
    np.stack(
        [
            phase_PT[..., L // 2],
            retardance_pr_PT[0, :, :, L // 2],
            azimuth[0, :, :, L // 2],
            theta[0, :, :, L // 2],
            absorption_PT[..., L // 2],
            retardance_pr_PT[1, :, :, L // 2],
            azimuth[1, :, :, L // 2],
            theta[1, :, :, L // 2],
        ]
    ),
    num_col=4,
    origin="lower",
    set_title=True,
    size=5,
    titles=[
        r"phase",
        r"principal retardance (+)",
        r"$\omega$ (+)",
        r"$\theta$ (+)",
        r"absorption",
        r"principal retardance (-)",
        r"$\omega$ (-)",
        r"$\theta$ (-)",
    ],
)
plt.show()

### Phase, principal retardance, color-coded 3D orientation (azimuth and inclination), and optic signtruction
## display parameters for 3D dataset ##
# z_layer = 25
# y_layer = 100
# x_layer = 100
# phase_min = -0.02
# phase_max = 0.02
# ret_min = 0
# ret_max = 0.007
# p_min = 0.4
# p_max = 0.6
# abs_min = -0.02
# abs_max = 0.02
######################################

## display parameters for 2D dataset ##
z_layer = L // 2
y_layer = M // 2
x_layer = N // 2
phase_min = -0.012
phase_max = 0.012
ret_min = 0
ret_max = 0.002
p_min = 0.4
p_max = 0.6
abs_min = -0.01
abs_max = 0.01
######################################

fig, ax = plt.subplots(6, 2, figsize=(5, 15))
sub_ax = ax[0, 0].imshow(
    absorption_PT[:, :, z_layer],
    cmap="gray",
    origin="lower",
    vmin=abs_min,
    vmax=abs_max,
)
ax[0, 0].set_title("absorption (xy)")
plt.colorbar(sub_ax, ax=ax[0, 0])

sub_ax = ax[0, 1].imshow(
    np.transpose(absorption_PT[y_layer, :, :]),
    cmap="gray",
    origin="lower",
    vmin=abs_min,
    vmax=abs_max,
    aspect=psz / ps,
)
ax[0, 1].set_title("absorption (xz)")
plt.colorbar(sub_ax, ax=ax[0, 1])

sub_ax = ax[1, 0].imshow(
    phase_PT[:, :, z_layer],
    cmap="gray",
    origin="lower",
    vmin=phase_min,
    vmax=phase_max,
)
ax[1, 0].set_title("phase (xy)")
plt.colorbar(sub_ax, ax=ax[1, 0])

sub_ax = ax[1, 1].imshow(
    np.transpose(phase_PT[y_layer, :, :]),
    cmap="gray",
    origin="lower",
    vmin=phase_min,
    vmax=phase_max,
    aspect=psz / ps,
)
ax[1, 1].set_title("absorption (xz)")
plt.colorbar(sub_ax, ax=ax[1, 1])

sub_ax = ax[2, 0].imshow(
    np.abs(retardance_pr_PT[0, :, :, z_layer]),
    cmap="gray",
    origin="lower",
    vmin=ret_min,
    vmax=ret_max,
)
ax[2, 0].set_title("principal retardance (+) (xy)")
plt.colorbar(sub_ax, ax=ax[2, 0])

sub_ax = ax[2, 1].imshow(
    np.transpose(np.abs(retardance_pr_PT[0, y_layer, :, :])),
    cmap="gray",
    origin="lower",
    vmin=ret_min,
    vmax=ret_max,
    aspect=psz / ps,
)
ax[2, 1].set_title("principal retardance (+) (xz)")
plt.colorbar(sub_ax, ax=ax[2, 1])

sub_ax = ax[3, 0].imshow(
    np.abs(p_mat_map[:, :, z_layer]),
    cmap="gray",
    origin="lower",
    vmin=p_min,
    vmax=p_max,
)
ax[3, 0].set_title("optic sign probability (xy)")
plt.colorbar(sub_ax, ax=ax[3, 0])

sub_ax = ax[3, 1].imshow(
    np.transpose(np.abs(p_mat_map[y_layer, :, :])),
    cmap="gray",
    origin="lower",
    vmin=p_min,
    vmax=p_max,
    aspect=psz / ps,
)
ax[3, 1].set_title("optic sign probability (xz)")
plt.colorbar(sub_ax, ax=ax[3, 1])

sub_ax = ax[4, 0].imshow(
    azimuth[0, :, :, z_layer], cmap="gray", origin="lower", vmin=0, vmax=np.pi
)
ax[4, 0].set_title("in-plane orientation (+) (xy)")

sub_ax = ax[4, 1].imshow(
    np.transpose(azimuth[0, y_layer, :, :]),
    cmap="gray",
    origin="lower",
    vmin=0,
    vmax=np.pi,
    aspect=psz / ps,
)
ax[4, 1].set_title("in-plane orientation (+) (xz)")

sub_ax = ax[5, 0].imshow(
    theta[0, :, :, z_layer], cmap="gray", origin="lower", vmin=0, vmax=np.pi
)
ax[5, 0].set_title("inclination (+) (xy)")

sub_ax = ax[5, 1].imshow(
    np.transpose(theta[0, y_layer, :, :]),
    cmap="gray",
    origin="lower",
    vmin=0,
    vmax=np.pi,
    aspect=psz / ps,
)
ax[5, 1].set_title("inclination (+) (xz)")
plt.show()

### Render 3D orientation with 3D colorsphere (azimuth and inclination)
# create color-coded orientation images

## display parameters for 3D dataset ##
# ret_min_color = 0
# ret_max_color = 0.007
######################################

## display parameters for 2D dataset ##
ret_min_color = 0
ret_max_color = 0.002
######################################

orientation_3D_image = np.transpose(
    np.array(
        [
            azimuth[0] / 2 / np.pi,
            theta[0],
            (
                np.clip(
                    np.abs(retardance_pr_PT[0]), ret_min_color, ret_max_color
                )
                - ret_min_color
            )
            / (ret_max_color - ret_min_color),
        ]
    ),
    (3, 1, 2, 0),
)
orientation_3D_image_RGB = visual.orientation_3D_to_rgb(
    orientation_3D_image, interp_belt=20 / 180 * np.pi, sat_factor=1
)

plt.figure(figsize=(10, 10))
plt.imshow(orientation_3D_image_RGB[z_layer], origin="lower")
plt.figure(figsize=(10, 10))
plt.imshow(
    orientation_3D_image_RGB[:, y_layer], origin="lower", aspect=psz / ps
)
# plot the top view of 3D orientation colorsphere
plt.figure(figsize=(3, 3))
visual.orientation_3D_colorwheel(
    wheelsize=128,
    circ_size=50,
    interp_belt=20 / 180 * np.pi,
    sat_factor=1,
    discretize=True,
)
plt.show()

### Render 3D orientation with 2 channels (in-plane orientation and out-of-plane tilt)

# in-plane orientation
from matplotlib.colors import hsv_to_rgb

I_hsv = np.transpose(
    np.array(
        [
            (azimuth[0]) % np.pi / np.pi,
            np.ones_like(retardance_pr_PT[0]),
            (
                np.clip(
                    np.abs(retardance_pr_PT[0]), ret_min_color, ret_max_color
                )
                - ret_min_color
            )
            / (ret_max_color - ret_min_color),
        ]
    ),
    (3, 1, 2, 0),
)
in_plane_orientation = hsv_to_rgb(I_hsv.copy())

plt.figure(figsize=(10, 10))
plt.imshow(in_plane_orientation[z_layer], origin="lower")
plt.figure(figsize=(10, 10))
plt.imshow(in_plane_orientation[:, y_layer], origin="lower", aspect=psz / ps)
plt.figure(figsize=(3, 3))
visual.orientation_2D_colorwheel()
plt.show()

# out-of-plane tilt

threshold_inc = np.pi / 90

I_hsv = np.transpose(
    np.array(
        [
            (
                -np.maximum(0, np.abs(theta[0] - np.pi / 2) - threshold_inc)
                + np.pi / 2
                + threshold_inc
            )
            / np.pi,
            np.ones_like(retardance_pr_PT[0]),
            (
                np.clip(
                    np.abs(retardance_pr_PT[0]), ret_min_color, ret_max_color
                )
                - ret_min_color
            )
            / (ret_max_color - ret_min_color),
        ]
    ),
    (3, 1, 2, 0),
)
out_of_plane_tilt = hsv_to_rgb(I_hsv.copy())

plt.figure(figsize=(10, 10))
plt.imshow(out_of_plane_tilt[z_layer], origin="lower")
plt.figure(figsize=(10, 10))
plt.imshow(out_of_plane_tilt[:, y_layer], origin="lower", aspect=psz / ps)
plt.show()

## Evaluation of reconstructed 3D orientation
### 3D orientation overlaid as lines on retardance image.
vect = np.zeros((3,) + azimuth.shape)
vect[0] = np.sin(theta) * np.cos(azimuth)
vect[1] = np.sin(theta) * np.sin(azimuth)
vect[2] = np.cos(theta)

theta_x = np.arccos(vect[0])
azimuth_x = np.arctan2(vect[2], vect[1])

theta_y = np.arccos(vect[1])
azimuth_y = np.arctan2(vect[2], vect[0])

z_step = psz

### 3D data parameter ###
# spacing = 4
# z_layer =  25
# x_layer = 100
# y_layer = 100
# linelength_scale = 10
#########################

### select slices to plot ###
spacing = 4
z_layer = L // 2
x_layer = N // 2
y_layer = M // 2
linelength_scale = 20
#########################

fig, ax = plt.subplots(2, 2, figsize=(10, 10))
visual.plot3DVectorField(
    np.abs(retardance_pr_PT[0, :, :, z_layer]),
    azimuth[0, :, :, z_layer],
    theta[0, :, :, z_layer],
    anisotropy=0.4 * np.abs(retardance_pr_PT[0, :, :, z_layer]),
    cmapImage="gray",
    clim=[ret_min, ret_max],
    aspect=1,
    spacing=spacing,
    window=spacing,
    linelength=spacing * linelength_scale,
    linewidth=1,
    cmapAzimuth="hsv",
    alpha=0.4,
    subplot_ax=ax[0, 0],
)
ax[0, 0].set_title(f"XY section (z= {z_layer})")

visual.plot3DVectorField(
    np.transpose(np.abs(retardance_pr_PT[0, :, x_layer, :])),
    np.transpose(azimuth_x[0, :, x_layer, :]),
    np.transpose(theta_x[0, :, x_layer, :]),
    anisotropy=0.4 * np.transpose(np.abs(retardance_pr_PT[0, :, x_layer, :])),
    cmapImage="gray",
    clim=[ret_min, ret_max],
    aspect=z_step / ps,
    spacing=spacing,
    window=spacing,
    linelength=spacing * linelength_scale,
    linewidth=1,
    cmapAzimuth="hsv",
    alpha=0.4,
    subplot_ax=ax[0, 1],
)
ax[0, 1].set_title(f"YZ section (x = {x_layer})")

visual.plot3DVectorField(
    np.transpose(np.abs(retardance_pr_PT[0, y_layer, :, :])),
    np.transpose(azimuth_y[0, y_layer, :, :]),
    np.transpose(theta_y[0, y_layer, :, :]),
    anisotropy=0.4 * np.transpose(np.abs(retardance_pr_PT[0, y_layer, :, :])),
    cmapImage="gray",
    clim=[ret_min, ret_max],
    aspect=z_step / ps,
    spacing=spacing,
    window=spacing,
    linelength=spacing * linelength_scale,
    linewidth=1,
    cmapAzimuth="hsv",
    alpha=0.4,
    subplot_ax=ax[1, 0],
)
ax[1, 0].set_title(f"XZ section (y = {y_layer})")

ax[1, 1].remove()
plt.show()

### Angular histogram of computed 3D orientation

ret_mask = np.abs(retardance_pr_PT[0]).copy()

## threshold parameters for 3D dataset ##
# ret_mask[ret_mask<0.15]=0
######################################

## threshold parameters for 2D dataset ##
ret_mask[ret_mask < 0.00125] = 0
######################################

plt.figure(figsize=(10, 10))
plt.imshow(ret_mask[:, :, z_layer], cmap="gray", origin="lower")
visual.orientation_3D_hist(
    azimuth[0].flatten(),
    theta[0].flatten(),
    ret_mask.flatten(),
    bins=36,
    num_col=1,
    size=10,
    contour_level=100,
    hist_cmap="gnuplot2",
    top_hemi=True,
)
plt.show()


from datetime import datetime, timedelta

class Player():
    """
    Representation of player as stored in the in-memory database (a single row from the player table).
    """

    def __init__(self, user_name, last_ping, in_game=None, ready=False):
        self.user_name = user_name
        self.last_ping = last_ping
        self.in_game = None if not in_game or in_game == '' else in_game
        self.ready = False if not ready else True

class PlayerDB():

    def __init__(self, connection):
        self.connection = connection

    def build_table(self):
        self.connection.cursor().execute('''
            create table player (
                user_name text not null primary key,
                in_game text null,
                last_ping text not null,
                ready boolean null)
        ''')

        self.connection.commit()

    def create(self, user_name):
        query = 'insert into player (user_name, last_ping)'
        query += ' values (:user_name, :last_ping)'

        self.connection.cursor().execute(query, {
            'last_ping': datetime.now(),
            'user_name': user_name
        })

        self.connection.commit()

    def exists(self, user_name):
        query = 'select count(*) from player where user_name = :user_name'

        cursor = self.connection.cursor()

        cursor.execute(query, {
            'user_name': user_name
        })
        count = cursor.fetchone()[0]

        return count > 0

    def in_game(self, game_name):
        """ Returns a list of players in a given game. """
        query = 'select user_name, last_ping, ready from player where in_game = :game_name'

        players_in_game = []

        cursor = self.connection.cursor()
        cursor.execute(query, {'game_name': game_name})
        for result in cursor.fetchall():
            players_in_game.append(Player(
                user_name=result[0],
                last_ping=result[1],
                ready=result[2],
                in_game=game_name
            ))

        return players_in_game

    def join_game(self, user_name, game_name):
        query = 'update player set in_game = :game_name where user_name = :user_name'

        self.connection.cursor().execute(query, {
            'user_name': user_name,
            'game_name': game_name
        })

        self.connection.commit()

    def kick_idle_players(self):
        query = 'update player set in_game = null where user_name != :reserved_game'
        query += ' and last_ping < :timeout_threshold'

        cursor = self.connection.cursor()

        cursor.execute(query, {
            'reserved_game': 'reserved_game',
            'timeout_threshold': datetime.now() - timedelta(seconds=10)
        })

        self.connection.commit()

    def leave_game(self, user_name):
        query = 'update player set in_game = null where user_name = :user_name'

        self.connection.cursor().execute(query, {
            'user_name': user_name
        })

        self.connection.commit()

    def ping(self, user_name):
        query = 'update player set last_ping = :last_ping where user_name = :user_name'

        self.connection.cursor().execute(query, {
            'last_ping': datetime.now(),
            'user_name': user_name
        })

        self.connection.commit()

    def set_ready_status(self, user_name, ready):
        query = 'update player set ready = :ready where user_name = :user_name'

        self.connection.cursor().execute(query, {
            'ready': ready,
            'user_name': user_name
        })


# flake8: noqa
# pylint: disable=invalid-name
# pylint: disable=line-too-long
# pylint: disable=missing-docstring
# pylint: disable=too-many-arguments
# pylint: disable=too-many-locals
# pylint: disable=too-many-statements

import networkx as nx

from ._read_records import read_records
from .nx_apply_cdlib_algorithm import nx_apply_cdlib_algorithm
from .nx_compute_edge_width_from_edge_weight import nx_compute_edge_width_from_edge_weight
from .nx_compute_node_degree import nx_compute_node_degree
from .nx_compute_node_size_from_node_degree import nx_compute_node_size_from_node_degree
from .nx_compute_spring_layout import nx_compute_spring_layout
from .nx_compute_textfont_opacity_from_node_degree import (
    nx_compute_textfont_opacity_from_node_degree,
)
from .nx_compute_textfont_size_from_node_degree import nx_compute_textfont_size_from_node_degree
from .nx_compute_textposition_from_graph import nx_compute_textposition_from_graph
from .nx_set_edge_color_to_constant import nx_set_edge_color_to_constant
from .nx_set_node_color_from_group_attr import nx_set_node_color_from_group_attr

def nx_create_co_citation_graph(
    #
    # FUNCTION PARAMS:
    unit_of_analysis,
    #
    # COLUMN PARAMS:
    top_n=None,
    citations_min=None,
    custom_items=None,
    #
    # NETWORK CLUSTERING:
    algorithm_or_dict="louvain",
    #
    # LAYOUT:
    nx_k=None,
    nx_iterations=30,
    nx_random_state=0,
    #
    # NODES:
    node_size_min=30,
    node_size_max=70,
    textfont_size_min=10,
    textfont_size_max=20,
    textfont_opacity_min=0.35,
    textfont_opacity_max=1.00,
    #
    # EDGES:
    edge_color="#7793a5",
    edge_width_min=0.8,
    edge_width_max=3.0,
    #
    # AXES:
    # xaxes_range=None,
    # yaxes_range=None,
    # show_axes=False,
    #
    # DATABASE PARAMS:
    root_dir="./",
    database="main",
    year_filter=(None, None),
    cited_by_filter=(None, None),
    **filters,
):
    #
    # Create the networkx graph
    nx_graph = nx.Graph()

    nx_graph = __add_weighted_edges_from(
        nx_graph=nx_graph,
        unit_of_analysis=unit_of_analysis,
        #
        # COLUMN PARAMS:
        top_n=top_n,
        citations_min=citations_min,
        custom_items=custom_items,
        #
        # DATABASE PARAMS:
        root_dir=root_dir,
        database=database,
        year_filter=year_filter,
        cited_by_filter=cited_by_filter,
        **filters,
    )

    if unit_of_analysis == "cited_references":
        for node in nx_graph.nodes():
            nx_graph.nodes[node]["text"] = ", ".join(node.split(", ")[:2])

    else:
        for node in nx_graph.nodes():
            nx_graph.nodes[node]["text"] = node

    #
    # Cluster the networkx graph
    if isinstance(algorithm_or_dict, str):
        nx_graph = nx_apply_cdlib_algorithm(nx_graph, algorithm_or_dict)
    if isinstance(algorithm_or_dict, dict):
        nx_graph = __assign_group_from_dict(nx_graph, algorithm_or_dict)

    #
    # Sets the layout
    nx_graph = nx_compute_spring_layout(nx_graph, nx_k, nx_iterations, nx_random_state)

    #
    # Sets the node attributes
    nx_graph = nx_compute_node_degree(nx_graph)

    nx_graph = nx_set_node_color_from_group_attr(nx_graph)

    nx_graph = nx_compute_node_size_from_node_degree(nx_graph, node_size_min, node_size_max)
    nx_graph = nx_compute_textfont_size_from_node_degree(
        nx_graph, textfont_size_min, textfont_size_max
    )
    nx_graph = nx_compute_textfont_opacity_from_node_degree(
        nx_graph, textfont_opacity_min, textfont_opacity_max
    )

    #
    # Sets the edge attributes
    nx_graph = nx_compute_edge_width_from_edge_weight(nx_graph, edge_width_min, edge_width_max)
    nx_graph = nx_compute_textposition_from_graph(nx_graph)
    nx_graph = nx_set_edge_color_to_constant(nx_graph, edge_color)

    return nx_graph

def __add_weighted_edges_from(
    nx_graph,
    unit_of_analysis,
    #
    # COLUMN PARAMS:
    top_n=None,
    citations_min=None,
    custom_items=None,
    #
    # DATABASE PARAMS:
    root_dir="./",
    database="main",
    year_filter=(None, None),
    cited_by_filter=(None, None),
    **filters,
):
    records = read_records(
        #
        # DATABASE PARAMS:
        root_dir=root_dir,
        database=database,
        year_filter=year_filter,
        cited_by_filter=cited_by_filter,
        **filters,
    )

    matrix_list = records[["global_references"]].dropna().copy()
    matrix_list = matrix_list.rename(columns={"global_references": "column"})
    matrix_list = matrix_list.assign(row=records[["global_references"]])

    for name in ["column", "row"]:
        matrix_list[name] = matrix_list[name].str.split(";")
        matrix_list = matrix_list.explode(name)
        matrix_list[name] = matrix_list[name].str.strip()

    if unit_of_analysis == "cited_authors":
        matrix_list["row"] = matrix_list["row"].str.split(", ").map(lambda x: x[0])
        matrix_list["column"] = matrix_list["column"].str.split(", ").map(lambda x: x[0])

    elif unit_of_analysis == "cited_sources":
        matrix_list["row"] = matrix_list["row"].str.split(", ").map(lambda x: x[2])
        matrix_list["column"] = matrix_list["column"].str.split(", ").map(lambda x: x[2])

    elif unit_of_analysis == "cited_references":
        matrix_list["row"] = matrix_list["row"].str.split(", ").map(lambda x: x[:3]).str.join(", ")
        matrix_list["column"] = (
            matrix_list["column"].str.split(", ").map(lambda x: x[:3]).str.join(", ")
        )

    else:
        raise ValueError("Bad unit_of_analysis")

    matrix_list = matrix_list.loc[matrix_list.apply(lambda x: x.row != x.column, axis=1), :]

    matrix_list["OCC"] = 1
    matrix_list = matrix_list.groupby(["row", "column"], as_index=False).aggregate("sum")

    matrix_list = matrix_list.sort_values(["OCC", "row", "column"], ascending=[False, True, True])

    #
    # Filter the data
    valid_items = __compute_valid_items(
        unit_of_analysis=unit_of_analysis,
        records=records,
        top_n=top_n,
        citations_min=citations_min,
        custom_items=custom_items,
    )

    matrix_list = matrix_list.loc[matrix_list.row.isin(valid_items), :]
    matrix_list = matrix_list.loc[matrix_list.column.isin(valid_items), :]

    #
    # Adds the data to the network:
    for _, x in matrix_list.iterrows():
        nx_graph.add_weighted_edges_from(
            ebunch_to_add=[(x.row, x.column, x.size)],
            dash="solid",
        )

    return nx_graph

def __assign_group_from_dict(nx_graph, group_dict):
    #
    # The group is assigned using and external algorithm. It is designed
    # to provide analysis capabilities to the system when other types of
    # analysis are conducted, for example, factor analysis.
    for node, group in group_dict.items():
        nx_graph.nodes[node]["group"] = group
    return nx_graph

def __compute_valid_items(
    unit_of_analysis,
    records,
    top_n,
    citations_min,
    custom_items,
):
    if custom_items is not None:
        return custom_items

    global_references = records["global_references"].dropna().copy()

    global_references = global_references.str.split(";")
    global_references = global_references.explode()
    global_references = global_references.str.strip()

    if unit_of_analysis == "cited_authors":
        global_references = global_references.str.split(", ").map(lambda x: x[0])

    elif unit_of_analysis == "cited_sources":
        global_references = global_references.str.split(", ").map(lambda x: x[2])

    elif unit_of_analysis == "cited_references":
        global_references = global_references.str.split(", ").map(lambda x: x[:3]).str.join(", ")

    else:
        raise ValueError("Bad unit_of_analysis")

    valid_items = global_references.value_counts().to_frame()
    valid_items.columns = ["citations"]
    valid_items = valid_items.reset_index()
    valid_items.sort_values(["citations", "index"], ascending=[False, True], inplace=True)

    if top_n is not None:
        valid_items = valid_items.head(top_n)

    if citations_min is not None:
        valid_items = valid_items.loc[valid_items.citations >= citations_min]

    return valid_items["index"].tolist()


import dash_bootstrap_components as dbc
import dash_core_components as dcc
import dash_html_components as html
from datetime import datetime
import plotly.express as px
import plotly.graph_objects as go
from . import config
from . import chart_contour_years
from . import chart_series_total
from . import chart_indicators

def layout():
    def create_card(*args, **kw):
        return dbc.Card(dbc.CardBody(*args, className="text-center"), **kw)
    # indicators
    def create_indicator(figure, id):
        return create_card([
            chart_indicators.indicator_title[id[10:]],
            dcc.Graph(figure=figure, id=id, style={'height': '4rem'})
        ])
    indicatorKmTotal = create_indicator(chart_indicators.plot('km-total'), 'indicator-km-total')
    indicatorDaysTotal = create_indicator(chart_indicators.plot('days-total'), 'indicator-days-total')
    equatorsTotal = create_indicator(chart_indicators.plot('equators-total', '.4f'), 'indicator-equators-total')
    # plots
    traceYear = dcc.Graph(figure=chart_series_total.total(), id='trace-year')
    contourYears = dcc.Graph(figure=chart_contour_years.years(), id='contour-years', style={'height': '80vh'})
    # this month
    now = config.data.this_month()
    this_month = datetime.strptime('%04d-%02d-01' % now, '%Y-%m-%d')
    return html.Div([
        # header
        dbc.Row([
            dbc.Col([
                html.H2('Jura jede')
            ], className="text-center")
        ], className="contentHeader"),
        # total series
        traceYear,
        dbc.Tooltip(
            "Historie měsíčních kilometrů od roku 1992.",
            target="trace-year"
        ),
        # indicators
        dbc.CardDeck([
            indicatorKmTotal,
            equatorsTotal,
            indicatorDaysTotal
        ]),
        # contour
        contourYears,
        dbc.Tooltip(
            "Historie měsíčních kilometrů podél let.",
            target="contour-years"
        ),
        dbc.Row([
            dbc.Col([
                dcc.RangeSlider(
                    id='series-type-slider',
                    min=config.total.min(),
                    max=config.total.max(),
                    value=[config.total.default(index=True)],
                    pushable=2,
                    marks={i:label for i,label in enumerate(config.total.labels())},
                    persistence = False
                )
            ], width=9),#, style={'width': '60%', 'display': 'inline-block', 'paddingTop': '3vh', 'paddingBottom': '3vh'}),
            dbc.Col([
                config.format_month_year(this_month)
            ], width=3, style={'textAlign': 'right'})
        ], style={'paddingTop': '2vh', 'paddingBottom': '2vh'})#,
    ])


#!/usr/bin/python
# vim: fileencoding=utf-8

from ansible.module_utils.basic import AnsibleModule
import urllib.request
import json
import os.path

ANSIBLE_METADATA = {
    'metadata_version': '1.1',
    'status': ['preview'],
    'supported_by': 'Masumi Kanai'
}

DOCUMENTATION = '''
---
module: github_asset

short_description: Download GitHub's release asset

version_added: "0.1"

description:
    - "Download GitHub's release asset"

options:
    token:
        description:
            - GitHub access token
        required: false
    repository:
        description:
            - repository name: e.g. `masu-mi/ansible-playbook`
        required: true
    asset_name:
        description:
            - target asset file name
        required: true
    tag:
        description:
            - Control to demo if the result of this module is changed or not
        required: false
    base_url:
        description:
            - base_url: e.g. `https://api.github.com/repos`
        required: false
    dest:
        description:
            - dest path in controlled hosts
        required: true
requirements:
  python >= 3.5

extends_documentation_fragment:

author:
    - Who am I?
'''

EXAMPLES = '''
# Pass in a message
- name: Download nidd-dra
  github_assets:
    token: "{{ lookup('env', 'GITHUB_TOKEN') }}"
    repository: "masu-mi/test-repo"
    tag: "v0.1.15"
    asset_name: "nidd-dra_0.1.15_linux_amd64.tar.gz"
    dest: "/tmp/"
'''

RETURN = '''
asset_url:
    description: The asset url to downlaod the target asset
    type: str
dest_file_path:
    description: The downloaded file path
    type: str
'''

def run_module():
    module_args = dict(
        token=dict(type='str', required=False, default=None),
        repository=dict(type='str', required=True),
        asset_name=dict(type='str', required=True),
        tag=dict(type='str', required=False, default='latest'),
        base_url=dict(type='str', required=False,
                      default='https://api.github.com/repos'),
        dest=dict(type='str', required=True)
    )
    result = dict(changed=False, asset_url='', dest_file_path='')

    module = AnsibleModule(
        argument_spec=module_args,
        supports_check_mode=True
    )
    repo = GitHubReleases(module.params['base_url'],
                          module.params['token'],
                          module.params['repository'])

    (result['asset_url'],
     result['dest_file_path'],
     failed) = repo.download_asset(
        module.params['dest'],
        module.params['tag'],
        module.params['asset_name'])
    if failed:
        module.fail_json(msg='You requested this to fail', **result)
    result['changed'] = True

    module.exit_json(**result)

class GitHubReleases():

    def __init__(self, base_url, token, repository):
        self.base_url = base_url
        self.token = token
        self.repository = repository

    def download_asset(self, dest, tag, asset_name):
        """
        download_asset
        """
        headers = {'Accept': 'application/octet-stream'}
        if self.token is not None:
            headers['Authorization'] = 'token {}'.format(self.token)
        (asset_url, failed) = self.get_asset_url(tag, asset_name)
        if failed:
            return ("", "", True)

        req = urllib.request.Request(asset_url, headers=headers)
        file_path = dest_file_path(dest, asset_name)

        with urllib.request.urlopen(req) as res:
            with open(file_path, mode='wb') as dest_file:
                dest_file.write(res.read())
        return (asset_url, file_path, False)

    def get_asset_url(self, tag, asset_name):
        url = '{}/{}/releases/tags/{}'.format(
            self.base_url, self.repository, tag)
        headers = {'Accept': 'application/vnd.github.v3+json'}
        if self.token is not None:
            headers['Authorization'] = 'token {}'.format(self.token)
        req = urllib.request.Request(url, headers=headers)
        with urllib.request.urlopen(req) as res:
            info = json.loads(res.read())
            asset_urls = [item['url']
                          for item in info["assets"]
                          if item["name"] == asset_name]
            if len(asset_urls) == 0:
                return ("", True)
            return (asset_urls[0], False)

def dest_file_path(dest, asset_name):
    if os.path.isdir(dest):
        dest = dest + asset_name
    return dest

def main():
    run_module()

if __name__ == '__main__':
    main()


import numpy as np
from utils.lidar_postprocessing import postprocess_lidar
from PIL import Image
from tqdm import tqdm
from utils.carla_to_cityscapes_segmentation_mapping import map_original, seg_to_cityscape_rgb
from copy import deepcopy
import sys
if '/opt/ros/kinetic/lib/python2.7/dist-packages' in sys.path:
    sys.path.remove('/opt/ros/kinetic/lib/python2.7/dist-packages')
import cv2
from utils.img_utils import decode_labels

from utils.lidar_pgm import PGM

counter = 0
directions = ["straight", "right", "left"]
for x in directions:
    direction = x
    path = '/media/user/research_data1/CARLA/LiDAR/lidar_data_combined_with_PGM/' + str(direction) + '/'
    # path = '/media/user/research_data1/CARLA/LiDAR/rich_pix2pix_data_no_pedestrians/straight4'
    saving_path = '/media/user/research_data1/CARLA/LiDAR/lidar_data_combined_with_PGM/alldata_sep/'
    # saving_path = '/media/user/research_data1/CARLA/LiDAR/rich_pix2pix_data_no_pedestrians/'
    saving_path_rgb = saving_path + 'rgb/'+x+'/'
    saving_path_seg = saving_path + 'seg/'+x+'/'
    saving_path_lidar = saving_path + 'lidar/'+x+'/'
    saving_path_pgm = saving_path+'PGM/'+x+'/'
    # saving_path_depth = path+'depth_samples/'
    rgb = np.load(path + 'rgb.npy')
    # depth = np.load(path + 'depth.npy')
    seg = np.load(path + 'seg.npy')
    lidar = np.load(path + 'LiDAR.npy')
    pgm = np.load(path + 'PGM.npy')
    for i in tqdm(range(seg.shape[0])):
        idx = i
        label_colours_carla_13 = [(192, 0, 0),
                                  (0, 128, 0),
                                  (192, 128, 128),
                                  (0, 0, 128),
                                  (64, 0, 128),
                                  (128, 64, 0),
                                  (128, 192, 0),
                                  (0, 64, 128),
                                  (0, 0, 0),
                                  (255, 255, 0),
                                  (0, 255, 255),
                                  (0, 255, 0),
                                  (255, 0, 255)]

        label_colours_scala_3 = [(0, 0, 0),
                                 (250, 250, 0),
                                 (69, 179, 231)]

        def seg_color(mask, num_classes):
            """Decode batch of segmentation masks.

            Args:
              mask: result of inference after taking argmax.
              num_images: number of images to decode from the batch.
              num_classes: number of classes to predict (including background).

            Returns:
              A batch with num_images RGB images of the same size as the input.
            """
            colours = None
            if num_classes == 13:
                colours = label_colours_carla_13
            elif num_classes == 3:
                colours = label_colours_scala_3
            else:
                print("ERROR this number of classes don't have a defined colours")
                exit(-1)

            n, h, w = mask.shape
            outputs = np.zeros((n, h, w, 3), dtype=np.uint8)
            for i in range(n):
                img = Image.new('RGB', (len(mask[i, 0]), len(mask[i])))
                pixels = img.load()
                for j_, j in enumerate(mask[i, :, :]):
                    for k_, k in enumerate(j):
                        if k < num_classes:
                            pixels[k_, j_] = colours[k]
                outputs[i] = np.array(img)
            return outputs

        from scipy import misc

        ### RGB
        misc.imsave(saving_path_rgb + str(idx) + '.png', rgb[i])
        # ### DEPTH
        # depth_img = depth[idx]
        # upper = depth_img > 0.1
        # lower = depth_img <= 0.1
        # depth_img[upper] = 1.0
        # depth_img[lower] *= 10.0
        # # print(depth_img.shape)
        # misc.imsave(saving_path_depth + str(idx) + '.png', depth_img)
        # ### seg
        seg_img = seg[i].astype(np.uint8)
        # seg = decode_labels(np.expand_dims(seg,0),3)
        # Original segmentation conversion for Carla
        seg_img = seg_color(np.expand_dims(seg_img, 0), 13)
        misc.imsave(saving_path_seg + str(idx) + '.png', seg_img[0])
        # exit(0)
        # Carla to cityscapes conversion
        # seg_img_copy = deepcopy(seg_to_cityscape_rgb(map_original(seg_img)))
        # misc.imsave(saving_path_seg + str(idx) + '.png', seg_img_copy)
        # seg_img = None

        # print(seg_img.shape)
        # ### Lidar
        # LiDAR_to_RGB_hlimit = 1750  # parser.add_argument('--hlimit', type=int, default=1750)
        # LiDAR_to_RGB_vlimit = 4000  # parser.add_argument('--vlimit', type=int, default=4000)
        # LiDAR_to_RGB_clip_back = True  # parser.add_argument('--clip_back', type=bool, default=True)
        # LiDAR_to_RGB_resolution = 0.05  # parser.add_argument('--resolution', type=float, default=0.15)
        # LiDAR_to_RGB_thickness = 0  # par
        lidar_img = lidar[i]
        lidar_img = cv2.resize(lidar_img, (200, 88), interpolation=cv2.INTER_NEAREST)
        misc.imsave(saving_path_lidar + str(idx) + '.png', lidar_img * 255.0)
        #
        ### PGM
        pgm_out = pgm[idx]
        # _
        misc.imsave(saving_path_pgm + str(idx) + '.png',pgm_out)
        # print('Done')
        counter += 1


import cv2
import numpy as np
from matplotlib import pyplot as plt
import itertools as it
from shapely.geometry import Polygon

# image = plt.imread("a.jpg")
#plt.imshow(image)

#plt.imshow(image)

missing = (134, 54)
X = np.array([60, 83, 68, 87, 110, 90, 109, 134])
Y = np.array([62, 76, 43, 53, 64, 35,  44, 54])
#X = np.array([31, 48, 68, 37, 53, 72, 42, 57])
#Y = np.array([69, 82, 97, 96, 110, 124, 118, 132])
def find_outlier(corners):
    un = unique_necklaces(corners.tolist())
    #Polygon with maximal area is not crossing itself
    config = un[np.argmax([Polygon(i).area for i in un])]

    min_dist = 100000
    points = []

    dist_list = []
    for i in range(len(config)):
        b = list(zip(config[i], config[(i+1) % 4]))

        plt.plot(b[0], b[1], c = "red")
        dist = np.linalg.norm(np.array([b[0][0], b[1][0]]) - np.array([b[0][1],b[1][1]]))

        dist_list.append((dist, (b[0], b[1])))
        if dist < min_dist:
            min_dist = dist
            points = (b[0], b[1])
            #print(dist, points)
    dist_list = sorted(dist_list)

    a = dist_list[0][1]
    b = dist_list[1][1]

    a = list(zip(a[0], a[1]))
    b = list(zip(b[0], b[1]))
    print(a)
    print(b)
    point = [0,0]
    ref0 = []
    ref1 = []
    if a[0] in b:
        point = a[0]
        ref0 = a[1]
        ref1 = b[(b.index(point)+1)%2]
    elif a[1] in b:
        point = a[1]
        ref0 = a[0]
        ref1 = b[(b.index(point)+1)%2]
    corners = corners.tolist()
    other_point = []
    for c in corners:
        if point[0] != c[0] and point[0] != ref0[0] and point[0] != ref1[0]:
            other_point = c
            break
    print([point,ref0,ref1,other_point])
    print(corners)

    return point, dist_list
def unique_necklaces(points):
    L = [points[0], points[1], points[2], points[3]]
    B = it.combinations(L,2)
    swaplist = [e for e in B]
    unique_necklaces = []
    for pair in swaplist:
        necklace = list(L)
        e1 = pair[0]
        e2 = pair[1]
        indexe1 = L.index(e1)
        indexe2 = L.index(e2)
        necklace[indexe1],necklace[indexe2] = necklace[indexe2], necklace[indexe1]
        unique_necklaces.append(necklace)
    return unique_necklaces

def order_points(pts):
  rect = np.zeros((4, 2), dtype = "float32")

  s = pts.sum(axis = 1)
  rect[0] = pts[np.argmin(s)]
  rect[2] = pts[np.argmax(s)]

  diff = np.diff(pts, axis = 1)
  rect[1] = pts[np.argmin(diff)]
  rect[3] = pts[np.argmax(diff)]

  if len(set([tuple(i) for i in rect.tolist()])) < 4:
  	rect[0] = pts[np.argmin(pts[:,0])]
  	rect[2] = pts[np.argmax(pts[:,0])]

  	rect[1] = pts[np.argmax(pts[:,1])]
  	rect[3] = pts[np.argmin(pts[:,1])]
  return rect

def four_point_transform (transform, pts):
    rect = order_points(pts)
    (tl, tr, br, bl) = rect

    widthA = np.sqrt(((br[0] - bl[0]) **2) + ((br[1] - bl[1]) **2))
    widthB = np.sqrt(((tr[0] - tl[0]) **2) + ((tr[1] - tl[1]) **2))
    maxWidth = max(int(widthA), int(widthB))

    heightA = np.sqrt(((tr[0] - br[0]) **2) + ((tr[1] - br[1]) **2))
    heightB = np.sqrt(((tl[0] - bl[0]) **2) + ((tl[1] - bl[1]) **2))
    maxHeight = max(int(heightA), int(heightB))

    dst = np.array([
                    [0, 0],
                    [maxWidth -1, 0],
                    [maxWidth -1, maxHeight -1],
                    [0, maxHeight -1]
    ], dtype = "float32")

    M = cv2.getPerspectiveTransform(rect, dst)
    warped = cv2.warpPerspective (transform, M, (maxWidth, maxHeight))

    return warped
    # cv2_imshow(fileName)
    # cv2_imshow(warped)

def get_slope(points):
    return (points[1][1] - points[1][0]) / (points[0][1] - points[0][0])
point_list = np.array(list(zip(X,Y)))
# warped = four_point_transform(image, point_list)

corners = order_points(np.array(list(zip(X,Y))))
print(corners)

outlier = find_outlier(corners)
test = outlier[1]
for i in test:
	print(get_slope(i[1]))
print(outlier[0])

plt.scatter(corners[:,0], corners[:,1], c="blue")
plt.scatter([outlier[0][0]],[outlier[0][1]])
plt.show()


import scipy.stats as sps
import numpy as np
from collections import Counter
import matplotlib.pyplot as plt
import tensorflow as tf

# plt.rcParams.update({'font.size': 14, "figure.figsize": (10,6)})

def profileplot(x, y, xlabel='', ylabel='', bins=100, mode=0):
    means_result = sps.binned_statistic(x, [y, y**2], bins=bins, statistic='mean')
    means, means2 = means_result.statistic
    standard_deviations = np.sqrt(means2 - means**2)
    bin_edges = means_result.bin_edges
    bin_centers = (bin_edges[:-1] + bin_edges[1:])/2.
    # remove NaNs and single count bins
    nan_idx = np.argwhere(np.isnan(means) ).flatten()
    zero_idx = np.argwhere(standard_deviations == 0)
    to_remove = np.union1d(nan_idx, zero_idx)
    means = np.delete(means, to_remove, None)
    bin_centers = np.delete(bin_centers, to_remove, None)
    standard_deviations = np.delete(standard_deviations, to_remove, None)
    count = Counter(means_result.binnumber)
    to_remove_set = set(to_remove)
    N = []
    for i in range(1,bins+1):
        if i-1 in to_remove_set:
            continue
        if i in count:
            N.append(count[i])
    # print(to_remove.shape)
    # print(bin_centers.shape, means.shape)
    yerr = standard_deviations/np.sqrt(N)
    # yerr = standard_deviations
    # fitting
    # print(bin_centers, means, yerr)
    plt.figure()
    plt.errorbar(x=bin_centers, y=means, yerr=yerr, linestyle='none', marker='.', capsize=2)
    if mode == 1:
        fit, cov = np.polyfit(bin_centers, means, 1, w=1/yerr, cov=True)
        p = np.poly1d(fit)
        print(f"Fit params: {fit[0]}, {fit[1]}")
        print(f"Diag of cov: {cov[0][0]} , {cov[1][1]}")
        plt.plot(bin_centers, p(bin_centers), label=f'gradient:{fit[0]:.2f}\nintercept:{fit[1]:.3f}')
        return fit, cov
    plt.grid()
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.legend()
    plt.tight_layout()

def profileplot_plain(x, y, xlabel='', ylabel='', bins=100, plot_range=None):
    means_result = sps.binned_statistic(x, [y, y**2], bins=bins, statistic='mean', range=plot_range)
    means, means2 = means_result.statistic
    standard_deviations = np.sqrt(means2 - means**2)
    bin_edges = means_result.bin_edges
    bin_centers = (bin_edges[:-1] + bin_edges[1:])/2.
    plt.figure()
    plt.errorbar(x=bin_centers, y=means, yerr=standard_deviations, linestyle='none', marker='.', capsize=2)
    plt.grid()
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.tight_layout()

class TensorBoardExtended(tf.keras.callbacks.TensorBoard):
    """
    Extended Tensorboard log that allows to add text

    By default logs:
    - host
    - gpus available

    Parameters
    -------------
    text_dict_to_log : dict
        Dictionary with key, value string that will be logged with Tensorboard
    kwargs : dict
        All the other parameters that are fed to Tensorboard
    """
    def __init__(self, text_dict_to_log=None, **kwargs):
        super().__init__(**kwargs)
        self.text_dict_to_log = text_dict_to_log

    def on_train_begin(self, logs=None):
        # pylint: disable= E1101
        super().on_train_begin(logs=logs)

        try:
            writer = self._get_writer('train')
        except AttributeError: # this is due to differences between tf21, 22 and 23
            writer = self._train_writer

        with writer.as_default():
            for key, value in self.text_dict_to_log.items():
                tf.summary.text(key, tf.convert_to_tensor(value), step=0)


import torch
from torch import nn
import warnings
#from eca_module import eca_layer
import torch.nn.functional as F
import math

warnings.filterwarnings(action='ignore')

####################################################################################################
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False, dilation=dilation)

def conv1x1(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

class eca_layer(nn.Module):
    """Constructs a ECA module.

    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """
    def __init__(self, channel, k_size=3):
        super(eca_layer, self).__init__() # super类的作用是继承的时候，调用含super的哥哥的基类__init__函数。
        self.avg_pool = nn.AdaptiveAvgPool2d(1) # 全局平均池化
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size,
                              padding=(k_size - 1) // 2, bias=False) # 一维卷积
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # x: input features with shape [b, c, h, w]
        b, c, h, w = x.size() # b代表b个样本，c为通道数，h为高度，w为宽度

        # feature descriptor on the global spatial information
        y = self.avg_pool(x)

        # Two different branches of ECA module
        # torch.squeeze()这个函数主要对数据的维度进行压缩,torch.unsqueeze()这个函数 主要是对数据维度进行扩充
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)

        # Multi-scale information fusion多尺度信息融合
        y = self.sigmoid(y)
        # 原网络中克罗内克积，也叫张量积，为两个任意大小矩阵间的运算
        return x * y.expand_as(x)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None,k_size=3):
        super(BasicBlock, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        # self.eca = eca_layer(planes, k_size)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        # out = self.eca(out)
        # print(out.shape)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None, k_size=3):
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.)) * groups
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.eca = eca_layer(planes * 4, k_size)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)
        out = self.eca(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out

class ResNet(nn.Module):
    def __init__(self, block, layers, num_classes=1000, groups=1, width_per_group=64,
                 replace_stride_with_dilation=None,norm_layer=None):
        super(ResNet, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer

        self.inplanes = 64
        self.dilation = 1
        if replace_stride_with_dilation is None:
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError("replace_stride_with_dilation should be None "
                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
        self.groups = groups
        self.base_width = width_per_group
        self.conv1 = nn.Conv2d(6, self.inplanes, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = norm_layer(self.inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
                                       dilate=replace_stride_with_dilation[0])
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
                                       dilate=replace_stride_with_dilation[1])
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
                                       dilate=replace_stride_with_dilation[2])
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                norm_layer(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
                            self.base_width, previous_dilation, norm_layer))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes, groups=self.groups,
                                base_width=self.base_width, dilation=self.dilation,
                                norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

def _resnet(arch, block, layers, pretrained, progress, **kwargs):
    model = ResNet(block, layers, **kwargs)
    return model

def _resnet18(pretrained=False, progress=True, **kwargs):
    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
                   **kwargs)

def _resnet34(pretrained=False, progress=True, **kwargs):
    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
                   **kwargs)

def _resnet50(pretrained=False, progress=True, **kwargs):
    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
                   **kwargs)

def _resnet101(pretrained=False, progress=True, **kwargs):
    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,
                   **kwargs)

####################################################################################################

class resnet18(torch.nn.Module):
    def __init__(self, pretrained=False):
        super().__init__()
        self.features = _resnet18(pretrained=pretrained)
        self.conv1 = self.features.conv1
        self.bn1 = self.features.bn1
        self.relu = self.features.relu
        self.maxpool1 = self.features.maxpool
        self.layer1 = self.features.layer1
        self.layer2 = self.features.layer2
        self.layer3 = self.features.layer3
        self.layer4 = self.features.layer4

    def forward(self, input):
        x = self.conv1(input)
        x = self.relu(self.bn1(x))
        x = self.maxpool1(x)
        feature1 = self.layer1(x)  # 1 / 4
        feature2 = self.layer2(feature1)  # 1 / 8
        feature3 = self.layer3(feature2)  # 1 / 16
        feature4 = self.layer4(feature3)  # 1 / 32
        # global average pooling to build tail
        tail = torch.mean(feature4, 3, keepdim=True)
        tail = torch.mean(tail, 2, keepdim=True)
        return feature3, feature4, tail

class resnet34(torch.nn.Module):
    def __init__(self, pretrained=False):
        super().__init__()
        self.features = _resnet34(pretrained=pretrained)
        self.conv1 = self.features.conv1
        self.bn1 = self.features.bn1
        self.relu = self.features.relu
        self.maxpool1 = self.features.maxpool
        self.layer1 = self.features.layer1
        self.layer2 = self.features.layer2
        self.layer3 = self.features.layer3
        self.layer4 = self.features.layer4

    def forward(self, input):
        x = self.conv1(input)
        x = self.relu(self.bn1(x))
        x = self.maxpool1(x)
        feature1 = self.layer1(x)  # 1 / 4
        feature2 = self.layer2(feature1)  # 1 / 8
        feature3 = self.layer3(feature2)  # 1 / 16
        feature4 = self.layer4(feature3)  # 1 / 32
        # global average pooling to build tail
        tail = torch.mean(feature4, 3, keepdim=True)
        tail = torch.mean(tail, 2, keepdim=True)
        return feature3, feature4, tail

class resnet50(torch.nn.Module):
    def __init__(self, pretrained=False):
        super().__init__()
        self.features = _resnet50()
        self.conv1 = self.features.conv1
        self.bn1 = self.features.bn1
        self.relu = self.features.relu
        self.maxpool1 = self.features.maxpool
        self.layer1 = self.features.layer1
        self.layer2 = self.features.layer2
        self.layer3 = self.features.layer3
        self.layer4 = self.features.layer4

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(self.bn1(x))
        x = self.maxpool1(x)
        feature1 = self.layer1(x)  # 1 / 4
        feature2 = self.layer2(feature1)  # 1 / 8
        feature3 = self.layer3(feature2)  # 1 / 16
        feature4 = self.layer4(feature3)  # 1 / 32
        # global average pooling to build tail
        tail = torch.mean(feature4, 3, keepdim=True)
        print(tail.shape)
        tail = torch.mean(tail, 2, keepdim=True)
        return feature3, feature4, tail

class resnet101(torch.nn.Module):
    def __init__(self, pretrained=False):
        super().__init__()
        self.features = _resnet101(pretrained=pretrained)
        self.conv1 = self.features.conv1
        self.bn1 = self.features.bn1
        self.relu = self.features.relu
        self.maxpool1 = self.features.maxpool
        self.layer1 = self.features.layer1
        self.layer2 = self.features.layer2
        self.layer3 = self.features.layer3
        self.layer4 = self.features.layer4

    def forward(self, input):
        x = self.conv1(input)
        x = self.relu(self.bn1(x))
        x = self.maxpool1(x)
        feature1 = self.layer1(x)  # 1 / 4
        feature2 = self.layer2(feature1)  # 1 / 8
        feature3 = self.layer3(feature2)  # 1 / 16
        feature4 = self.layer4(feature3)  # 1 / 32
        # global average pooling to build tail
        tail = torch.mean(feature4, 3, keepdim=True)
        tail = torch.mean(tail, 2, keepdim=True)
        return feature3, feature4, tail

def build_contextpath(name='resnet50'):
    model = {
        'resnet34': resnet34(pretrained=False),
        'resnet18': resnet18(pretrained=False),
        'resnet50': resnet50(pretrained=False),
        'resnet101': resnet101(pretrained=False)
    }
    return model[name]

####################################################################################################

class ConvBlock(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=2, padding=1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU()
        self.init_weight()

    def forward(self, input):
        x = self.conv1(input)
        return self.relu(self.bn(x))

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

class Spatial_path(torch.nn.Module):
    """Assimilation Network"""
    def __init__(self):
        super().__init__()
        self.convblock1 = ConvBlock(in_channels=6, out_channels=128)
        self.convblock2 = ConvBlock(in_channels=128, out_channels=256)
        self.convblock3 = ConvBlock(in_channels=256, out_channels=512)
        self.init_weight()

    def forward(self, input):
        x = self.convblock1(input)
        x = self.convblock2(x)
        x = self.convblock3(x)
        ##TODO: the channel number of x
        # print("assimilation",x.shape)
        return x

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

class Differentiation_path(torch.nn.Module):
    """Differece Network"""
    def __init__(self):
        super().__init__()
        self.convblock1 = ConvBlock(in_channels=3, out_channels=64)
        self.convblock2 = ConvBlock(in_channels=64, out_channels=128)
        self.convblock3 = ConvBlock(in_channels=128, out_channels=256)
        self.init_weight()

    def forward(self, input):
        x = self.convblock1(input)
        x = self.convblock2(x)
        x = self.convblock3(x)
        ##TODO: the channel number of x
        # print("difference",x.shape)
        return x

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

class AttentionRefinementModule(torch.nn.Module):
    """Fusion Network"""
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
        self.bn = nn.BatchNorm2d(out_channels)
        self.sigmoid = nn.Sigmoid()
        self.in_channels = in_channels
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
        self.init_weight()

    def forward(self, input):
        # global average pooling
        x = self.avgpool(input)
        assert self.in_channels == x.size(1), 'in_channels and out_channels should all be {}'.format(x.size(1))
        x = self.conv(x)
        # x = self.sigmoid(self.bn(x))
        x = self.sigmoid(x)
        # channels of input and x should be same
        x = torch.mul(input, x)
        return x

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

class BGALayer4(nn.Module):#4倍
    """Feature Aggregation Network"""
    def __init__(self):
        super(BGALayer4, self).__init__()
        self.left1 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=1,padding=1, groups=512, bias=False),
            nn.BatchNorm2d(512),
            nn.Conv2d(512, 512, kernel_size=1, stride=1, padding=0, bias=False),
                                )
        self.left2 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=2,padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.AvgPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=False)
                                )
        self.right1 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=1,padding=1, bias=False),
            nn.BatchNorm2d(512),
                                )
        self.right2 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=1,padding=1, groups=512, bias=False),
            nn.BatchNorm2d(512),
            nn.Conv2d(512, 512, kernel_size=1, stride=1, padding=0, bias=False),
                              )

        ##TODO: does this really has no relu?

        self.conv = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=1,padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True), # not shown in paper
                                )
        self.init_weight()

    def forward(self, dx, cx2):#x_d.Size([2, 128, 64, 64]).x_s.Size([2, 128, 16, 16])  x_different.Size([2, 128, 64, 64])

        dsize = dx.size()[2:]
        left1 = self.left1(dx)#torch.Size([2, 128, 64, 64])
        left2 = self.left2(dx)#torch.Size([2, 128, 16, 16])
        right1 = self.right1(cx2)#torch.Size([2, 128, 16, 16])
        right2 = self.right2(cx2)#torch.Size([2, 128, 16, 16])
        right1 = F.interpolate(right1, size=dsize, mode='bilinear', align_corners=True)
        #右1上采用4倍才能和左1相加torch.Size([2, 128, 64, 64])
        left = left1 * torch.sigmoid(right1)#torch.Size([2, 128, 64, 64])
        right = left2 * torch.sigmoid(right2)#torch.Size([2, 128, 16, 16])
        right = F.interpolate(right, size=dsize, mode='bilinear', align_corners=True)
        #左2和右2合并上采用4倍然后输出torch.Size([2, 128, 64, 64])
        out = self.conv(left + right )#torch.Size([2, 128, 64, 64])
        return out

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

class BGALayer2(nn.Module):#2倍

    def __init__(self):
        super(BGALayer2, self).__init__()
        self.left1 = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, stride=1,padding=1, groups=256, bias=False),
            nn.BatchNorm2d(256),
            nn.Conv2d(256, 256, kernel_size=1, stride=1,padding=0, bias=False),
        )
        self.left2 = nn.Sequential(
            nn.Conv2d( 256, 256, kernel_size=3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(256),
            # nn.AvgPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=False)
        )
        self.right1 = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, stride=1,padding=1, bias=False),
            nn.BatchNorm2d(256),
        )
        self.right2 = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, stride=1,padding=1, groups=256, bias=False),
            nn.BatchNorm2d(256),
            nn.Conv2d(256, 256, kernel_size=1, stride=1,padding=0, bias=False),
        )

        ##TODO: does this really has no relu?

        self.conv = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, stride=1,padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True), # not shown in paper
        )
        self.init_weight()

    def forward(self, sx, cx1):#x_d.Size([2, 128, 64, 64]).x_s.Size([2, 128, 16, 16])  x_different.Size([2, 128, 64, 64])

        dsize = sx.size()[2:]
        left1 = self.left1(sx)#torch.Size([2, 128, 64, 64])
        left2 = self.left2(sx)#torch.Size([2, 128, 16, 16])
        right1 = self.right1(cx1)#torch.Size([2, 128, 16, 16])
        right2 = self.right2(cx1)#torch.Size([2, 128, 16, 16])
        right1 = F.interpolate(
            right1, size=dsize, mode='bilinear', align_corners=True)#右1上采用4倍才能和左1相加torch.Size([2, 128, 64, 64])
        left = left1 * torch.sigmoid(right1)#torch.Size([2, 128, 64, 64])
        right = left2 * torch.sigmoid(right2)#torch.Size([2, 128, 16, 16])
        right = F.interpolate(
            right, size=dsize, mode='bilinear', align_corners=True)#左2和右2合并上采用4倍然后输出torch.Size([2, 128, 64, 64])

        out = self.conv(left + right )#torch.Size([2, 128, 64, 64])

        return out

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

class FeatureFusionModule(torch.nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        # self.in_channels = input_1.channels + input_2.channels
        # self.in_channels = 3584
        # self.in_channels = 3328
        self.in_channels = 768
        self.convblock = ConvBlock(in_channels=self.in_channels, out_channels=num_classes, stride=1)
        self.conv1 = nn.Conv2d(num_classes, num_classes, kernel_size=1)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(num_classes, num_classes, kernel_size=1)
        self.sigmoid = nn.Sigmoid()
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
        self.init_weight()

    def forward(self, input_1, input_2):
        x = torch.cat((input_1, input_2), dim=1)#torch.Size([1, 1280, 64, 64])
        feature = self.convblock(x)#torch.Size([1, 2, 64, 64])
        x = self.avgpool(feature)

        x = self.relu(self.conv1(x))
        x = self.sigmoid(self.conv2(x))
        x = torch.mul(feature, x)#torch.Size([1, 2, 64, 64])
        x = torch.add(x, feature)#torch.Size([1, 2, 64, 64])
        return x

    def init_weight(self):
        for ly in self.children():
            if isinstance(ly, nn.Conv2d):
                nn.init.kaiming_normal_(ly.weight, a=1)
                if not ly.bias is None: nn.init.constant_(ly.bias, 0)

class MFGANnet(torch.nn.Module):
    def __init__(self, num_classes=2, context_path='resnet34'):
        super().__init__()
        # build spatial path
        self.saptial_path = Spatial_path()#(1,256,64,64)
        self.differentiation_path = Differentiation_path()#(1,256,64,64)

        # build context path
        self.context_path = build_contextpath(name=context_path)

        # build attention refinement module
        self.attention_refinement_module1 = AttentionRefinementModule(256, 256) # resnet 18 / 34
        self.attention_refinement_module2 = AttentionRefinementModule(512, 512) # resnet 18 / 34
        # self.attention_refinement_module1 = AttentionRefinementModule(1024, 1024)  # resnet 50 / 101
        # self.attention_refinement_module2 = AttentionRefinementModule(2048, 2048)  # resnet 50 / 101
        self.bga2 = BGALayer2()
        self.bga4 = BGALayer4()

        # build feature fusion module
        self.feature_fusion_module = FeatureFusionModule(num_classes)

        # build final convolution
        self.conv = nn.Conv2d(in_channels=num_classes, out_channels=num_classes, kernel_size=1)

        # supervision block
        self.supervision1 = nn.Conv2d(in_channels=256, out_channels=num_classes, kernel_size=1) #resnet 18 / 34
        self.supervision2 = nn.Conv2d(in_channels=512, out_channels=num_classes, kernel_size=1) # resnet 18 / 34
        # self.supervision1 = nn.Conv2d(in_channels=1024, out_channels=num_classes, kernel_size=1)  # resnet 50 / 101
        # self.supervision2 = nn.Conv2d(in_channels=2048, out_channels=num_classes, kernel_size=1)  # resnet 50 / 101

    def forward(self, x1, x2):
        input = torch.cat([x1, x2], dim=1)
        # output of spatial path
        sx = self.saptial_path(input) # (1, 256, input/8, input/8)([1, 512, 64, 64])
        # print('sx',sx.shape)
        dx1 = self.differentiation_path(x1)#([1, 256, 64, 64])
        # print('dx1',dx1.shape)
        dx2 = self.differentiation_path(x2)#([1, 256, 64, 64])
        dx = dx1 - dx2##([1, 256, 64, 64])
        # print(dx1.shape)

        # output of context path
        cx1, cx2, tail = self.context_path(input) # ([1, 256, 32, 32])  tail([1, 512, 1, 1])
        # print(cx1.shape)
        # print(cx2.shape)
        # print('tail',tail.shape) #[1, 512, 1, 1]
        cx1 = self.attention_refinement_module1(cx1)#[1, 256, 32, 32])
        cx2 = self.attention_refinement_module2(cx2) # [1, 512, 16, 16])
        cx2 = torch.mul(cx2, tail)#[1, 512, 16, 16])
        # print(cx1.shape)
        # print(cx2.shape)
        # print(tail.shape)

        feat4 = self.bga4(sx, cx2)#torch.Size([1, 512, 64, 64])
        # print('a',feat4.shape)
        feat2 = self.bga2(dx, cx1)#torch.Size([1, 256, 64, 64])
        # print('b',feat2.shape)

        # upsampling
        cx3 = torch.nn.functional.interpolate(cx1, scale_factor=2, mode='bilinear') # (1, 1024, input/8, input/8)([1, 256, 64, 64])
        cx4 = torch.nn.functional.interpolate(cx2, scale_factor=4, mode='bilinear') # (1, 2048, input/8, input/8)([1, 512, 64, 64])

        # cx = torch.cat((feat2, feat4), dim=1) # (1, 3072, input/8, input/8)([1, 768, 64, 64])

        if self.training == True:
            cx3_sup = self.supervision1(cx3)
            cx4_sup = self.supervision2(cx4)
            cx3_sup = torch.nn.functional.interpolate(cx3_sup, scale_factor=8, mode='bilinear')
            cx4_sup = torch.nn.functional.interpolate(cx4_sup, scale_factor=8, mode='bilinear')

        # output of feature fusion module
        result = self.feature_fusion_module(feat2, feat4)#([1, 2, 64, 64])

        # upsampling
        result = torch.nn.functional.interpolate(result, scale_factor=8, mode='bilinear')
        result = self.conv(result)

        if self.training == True:
            return result, cx3_sup, cx4_sup

        return result

if __name__ == '__main__':
    model = MFGANnet(2, 'resnet34')
    x1 = torch.rand(1, 3, 512, 512)
    x2 = torch.rand(1, 3, 512, 512)
    result = model(x1, x2)
    print(result.size())
    print(result(type))
    # y1, y2, y3 = model(x1, x2)
    # print(y1.size())
    # print(y2.size())
    # print(y3.size())
    # x = torch.randn(1, 3, 512, 512)
    # net1 = Differentiation_path()
    # x1 = net1(x)
    # print(x1.shape)


from django.urls import path

from . import views
app_name = 'clothesapp'

urlpatterns = [
    path('', views.init, name='init'),
    path('create', views.create , name='create'),
    path('update/<int:pk>', views.update , name='update'),
    path('delete/<int:pk>', views.delete , name='delete')
]