from PepperPepper.environment import threading, np, torch, F, cv2
from skimage import measure

__all__ = ['SegmentationMetricTPFNFP']

def get_miou_prec_recall_fscore(total_tp, total_fp, total_fn):
    miou = 1.0 * total_tp / (np.spacing(1) + total_tp + total_fp + total_fn)
    prec = 1.0 * total_tp / (np.spacing(1) + total_tp + total_fp)
    recall = 1.0 * total_tp / (np.spacing(1) + total_tp + total_fn)
    fscore = 2.0 * prec * recall / (np.spacing(1) + prec + recall)

    return miou, prec, recall, fscore

class SegmentationMetricTPFNFP(object):
    """Computes pixAcc and mIoU metric scroes
    """

    def __init__(self, nclass):
        self.nclass = nclass
        self.lock = threading.Lock()
        self.reset()

    def update(self, labels, preds):
        def evaluate_worker(self, label, pred):
            tp, fp, fn = batch_tp_fp_fn(pred, label, self.nclass)
            with self.lock:
                self.total_tp += tp
                self.total_fp += fp
                self.total_fn += fn
            return

        if isinstance(preds, torch.Tensor):
            preds = (preds.detach().numpy() > 0).astype('int64')  # P
            labels = labels.numpy().astype('int64')  # T
            evaluate_worker(self, labels, preds)
        elif isinstance(preds, (list, tuple)):
            threads = [threading.Thread(target=evaluate_worker,
                                        args=(self, label, pred),
                                        )
                       for (label, pred) in zip(labels, preds)]
            for thread in threads:
                thread.start()
            for thread in threads:
                thread.join()
        #elif preds.dtype == numpy.uint8:
        elif isinstance(preds, np.ndarray):
            preds = ((preds / np.max(preds)) > 0.5).astype('int64')  # P
            labels = (labels / np.max(labels)).astype('int64')  # T
            evaluate_worker(self, labels, preds)
        else:
            raise NotImplemented

    def get_all(self):
        return self.total_tp, self.total_fp, self.total_fn

    def get(self):
        return get_miou_prec_recall_fscore(self.total_tp, self.total_fp, self.total_fn)

    def reset(self):
        self.total_tp = 0
        self.total_fp = 0
        self.total_fn = 0
        return

def batch_tp_fp_fn(predict, target, nclass):
    """Batch Intersection of Union
    Args:
        predict: input 4D tensor
        target: label 3D tensor
        nclass: number of categories (int)
    """

    mini = 1
    maxi = nclass
    nbins = nclass

    # predict = (output.detach().numpy() > 0).astype('int64')  # P
    # target = target.numpy().astype('int64')  # T
    intersection = predict * (predict == target)  # TP

    # areas of intersection and union
    area_inter, _ = np.histogram(intersection, bins=nbins, range=(mini, maxi))
    area_pred, _ = np.histogram(predict, bins=nbins, range=(mini, maxi))
    area_lab, _ = np.histogram(target, bins=nbins, range=(mini, maxi))

    # areas of TN FP FN
    area_tp = area_inter[0]
    area_fp = area_pred[0] - area_inter[0]
    area_fn = area_lab[0] - area_inter[0]

    # area_union = area_pred + area_lab - area_inter
    assert area_tp <= (area_tp + area_fn + area_fp)
    return area_tp, area_fp, area_fn


class PD_FA():
    def __init__(self, nclass, bins, size):
        super(PD_FA, self).__init__()
        self.nclass = nclass
        self.bins = bins
        self.image_area_total = []
        self.image_area_match = []
        self.FA = np.zeros(self.bins + 1)
        self.PD = np.zeros(self.bins + 1)
        self.target = np.zeros(self.bins + 1)
        self.size = size

    def update(self, preds, labels):

        for iBin in range(self.bins + 1):
            score_thresh = iBin * (255 / self.bins)
            predits = np.array((preds > score_thresh).cpu()).astype('int64')

            predits = np.reshape(predits, (self.size, self.size))
            labelss = np.array((labels).cpu()).astype('int64')
            labelss = np.reshape(labelss, (self.size, self.size))

            image = measure.label(predits, connectivity=2)
            coord_image = measure.regionprops(image)
            label = measure.label(labelss, connectivity=2)
            coord_label = measure.regionprops(label)

            self.target[iBin] += len(coord_label)
            self.image_area_total = []
            self.image_area_match = []
            self.distance_match = []
            self.dismatch = []

            for K in range(len(coord_image)):
                area_image = np.array(coord_image[K].area)
                self.image_area_total.append(area_image)

            for i in range(len(coord_label)):
                centroid_label = np.array(list(coord_label[i].centroid))
                for m in range(len(coord_image)):
                    centroid_image = np.array(list(coord_image[m].centroid))
                    distance = np.linalg.norm(centroid_image - centroid_label)
                    area_image = np.array(coord_image[m].area)
                    if distance < 3:
                        self.distance_match.append(distance)
                        self.image_area_match.append(area_image)

                        del coord_image[m]
                        break

            self.dismatch = [x for x in self.image_area_total if x not in self.image_area_match]
            self.FA[iBin] += np.sum(self.dismatch)
            self.PD[iBin] += len(self.distance_match)

    def get(self, img_num):

        Final_FA = self.FA / ((self.size * self.size) * img_num)
        Final_PD = self.PD / self.target

        return Final_FA, Final_PD

    def reset(self):
        # self.image_area_total = []
        # self.image_area_match = []
        # self.FA = np.zeros([self.bins + 1])
        # self.PD = np.zeros([self.bins + 1])
        # self.target = np.zeros(self.bins + 1)

        self.nclass = self.nclass
        self.bins = self.bins
        self.image_area_total = []
        self.image_area_match = []
        self.FA = np.zeros(self.bins + 1)
        self.PD = np.zeros(self.bins + 1)
        self.target = np.zeros(self.bins + 1)
        self.size = self.size


    
















# class PD_FA():
#     def __init__(self, bins, size):
#         super(PD_FA, self).__init__()
#         self.bins = bins
#         self.size = size
#         self.image_area_total = []
#         self.image_area_match = []
#         self.dismatch_pixel = 0
#         self.all_pixel = 0
#         self.PD = 0
#         self.target= 0
#     def update(self, preds, labels):
#         predits  = np.array((preds).cpu()).astype('int64')
#         labelss = np.array((labels).cpu()).astype('int64') 

#         image = measure.label(predits, connectivity=2)
#         coord_image = measure.regionprops(image)
#         label = measure.label(labelss , connectivity=2)
#         coord_label = measure.regionprops(label)

#         self.target    += len(coord_label)
#         self.image_area_total = []
#         self.distance_match   = []
#         self.dismatch         = []

#         for K in range(len(coord_image)):
#             area_image = np.array(coord_image[K].area)
#             self.image_area_total.append(area_image)

#         true_img = np.zeros(predits.shape)
#         for i in range(len(coord_label)):
#             centroid_label = np.array(list(coord_label[i].centroid))
#             for m in range(len(coord_image)):
#                 centroid_image = np.array(list(coord_image[m].centroid))
#                 distance = np.linalg.norm(centroid_image - centroid_label)
#                 area_image = np.array(coord_image[m].area)
#                 if distance < 7:
#                     self.distance_match.append(distance)
#                     true_img[coord_image[m].coords[:,0], coord_image[m].coords[:,1]] = 1
#                     del coord_image[m]
#                     break

#         self.dismatch_pixel += (predits - true_img).sum()
#         # self.all_pixel +=size[0]*size[1]
#         self.all_pixel += self.size * self.size
#         self.PD +=len(self.distance_match)

#     def get(self):
#         Final_FA =  self.dismatch_pixel / self.all_pixel
#         Final_PD =  self.PD /self.target
#         return Final_PD, float(Final_FA.cpu().detach().numpy())

#     def reset(self):
#         # self.FA  = np.zeros([self.bins+1])
#         # self.PD  = np.zeros([self.bins+1])
#         # self.bins = bins
#         # self.size = size
#         self.image_area_total = []
#         self.image_area_match = []
#         self.dismatch_pixel = 0
#         self.all_pixel = 0
#         self.PD = 0
#         self.target= 0











