from multiprocessing.pool import Pool

import numpy as np
from skimage import measure


def calculate_cell_score_kaggle(y_true, y_pred, num_threads=32):
    yps = []
    for m in range(len(y_true)):
        yps.append((y_true[m].copy(), y_pred[m].copy()))
    pool = Pool(num_threads)
    results = pool.map(score_kaggle, yps)
    return np.mean(results)


def calculate_cell_score_selim(y_true, y_pred, num_threads=32, ids=None):
    yps = []
    for m in range(len(y_true)):
        yps.append((y_true[m].copy(), y_pred[m].copy()))
    pool = Pool(num_threads)
    results = pool.map(calculate_jaccard, yps)
    if ids:
        import pandas as pd
        s_iou = np.argsort(results)
        d = []
        for i in range(len(s_iou)):
            id = ids[s_iou[i]]
            res = results[s_iou[i]]
            d.append([id, res])
            pd.DataFrame(d, columns=["ID", "METRIC_SCORE"]).to_csv("gt_vs_oof.csv", index=False)

    return np.array(results).mean()


def get_cells(mask):
    return measure.label(mask, return_num=True)

def score_kaggle(yp):
    y, p = yp
    return calc_score(np.expand_dims(y, 0), np.expand_dims(p, 0))



def calc_score(labels, y_pred):
    true_objects = len(np.unique(labels))
    pred_objects = len(np.unique(y_pred))
    #    print("Number of true objects:", true_objects)
    #    print("Number of predicted objects:", pred_objects)
    # Compute intersection between all objects
    intersection = np.histogram2d(labels.flatten(), y_pred.flatten(), bins=(true_objects, pred_objects))[0]

    # Compute areas (needed for finding the union between all objects)
    area_true = np.histogram(labels, bins=true_objects)[0]
    area_pred = np.histogram(y_pred, bins=pred_objects)[0]
    area_true = np.expand_dims(area_true, -1)
    area_pred = np.expand_dims(area_pred, 0)

    # Compute union
    union = area_true + area_pred - intersection

    # Exclude background from the analysis
    intersection = intersection[1:, 1:]
    union = union[1:, 1:]
    union[union == 0] = 1e-9

    # Compute the intersection over union
    iou = intersection / union

    # Precision helper function
    def precision_at(threshold, iou):
        matches = iou > threshold
        true_positives = np.sum(matches, axis=1) == 1  # Correct objects
        false_positives = np.sum(matches, axis=0) == 0  # Missed objects
        false_negatives = np.sum(matches, axis=1) == 0  # Extra objects
        tp, fp, fn = np.sum(true_positives), np.sum(false_positives), np.sum(false_negatives)
        return tp, fp, fn

    # Loop over IoU thresholds
    prec = []
    #    print("Thresh\tTP\tFP\tFN\tPrec.")
    for t in np.arange(0.5, 1.0, 0.05):
        tp, fp, fn = precision_at(t, iou)
        p = tp / (tp + fp + fn)
        #        print("{:1.3f}\t{}\t{}\t{}\t{:1.3f}".format(t, tp, fp, fn, p))
        prec.append(p)
    #    print("AP\t-\t-\t-\t{:1.3f}".format(np.mean(prec)))
    return np.mean(prec)


def calculate_jaccard(yps):
    y, p = yps
    jaccards = []
    iou_thresholds = [0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95]
    for iou_threshold in iou_thresholds:
        tp = 0
        fp = 0
        fn = 0
        processed_gt = set()
        matched = set()
        size = p.shape[0], p.shape[1]
        mask_img = np.reshape(p, size)

        gt_mask_img = np.reshape(y, size)
        predicted_labels, predicted_count = get_cells(mask_img)
        gt_labels, gt_count = get_cells(gt_mask_img)

        gt_cells = [rp.coords for rp in measure.regionprops(gt_labels)]
        pred_cells = [rp.coords for rp in measure.regionprops(predicted_labels)]
        gt_cells = [to_point_set(b) for b in gt_cells]
        pred_cells = [to_point_set(b) for b in pred_cells]
        for j in range(predicted_count):
            match_found = False
            for i in range(gt_count):
                pred_ind = j + 1
                gt_ind = i + 1
                if match_found:
                    break
                if gt_ind in processed_gt:
                    continue
                pred_cell = pred_cells[j]
                gt_cell = gt_cells[i]
                intersection = len(pred_cell.intersection(gt_cell))
                union = len(pred_cell) + len(gt_cell) - intersection
                iou = intersection / union
                if iou > iou_threshold:
                    processed_gt.add(gt_ind)
                    matched.add(pred_ind)
                    match_found = True
                    tp += 1
            if not match_found:
                fp += 1
        fn += gt_count - len(processed_gt)
        jaccards.append(tp / (tp + fp + fn))
    return np.mean(jaccards)


def to_point_set(cell):
    return set([(row[0], row[1]) for row in cell])