import functools
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

from perception.custom_datasets import ComparisonDataset
from perception.utils.visualization import plot_segmentation, plot_image, display_images_horizontally
from perception.utils.segmentation_labels import DEFAULT_CLASSES


def mean_jaccard_index(target, predictions):
    """
    Semantic segmentation metric. Calculates mean intersection over union over all classes.
    Works for batches of data.
    """

    intersection = torch.logical_and(target, predictions)
    union = torch.logical_or(target, predictions)

    intersection_sums = torch.sum(intersection, dim=(-2, -1))
    union_sums = torch.sum(union, dim=(-2,-1))

    class_exists_mask = union_sums != 0

    # union_sums will contain 0's if a class is not present in an image, which will give division by zero
    iou_scores_classwise = intersection_sums / (union_sums + 0.00000000001)

    iou_scores_imagewise_sum = iou_scores_classwise.sum(dim=1)
    class_exists_mask_sum = class_exists_mask.sum(dim=1)
    iou_scores_imagewise_mean = iou_scores_imagewise_sum / class_exists_mask_sum

    iou_score_batch_mean = torch.mean(iou_scores_imagewise_mean)

    return iou_score_batch_mean.numpy(), iou_scores_classwise.numpy()


def weighted_jaccard_index(target, predictions):
    """
    Semantic segmentation metric. Calculates mean intersection over union over all classes, weighted by class.
    Works for batches of data.
    """

    class_counts = torch.sum(target, dim=(-2, -1), dtype=torch.int32)
    n_pixels = torch.sum(class_counts, dim=(-1))[0]
    class_frequencies = class_counts / n_pixels

    intersection = torch.logical_and(target, predictions)
    union = torch.logical_or(target, predictions)

    intersection_sums = torch.sum(intersection, dim=(-2, -1))
    union_sums = torch.sum(union, dim=(-2,-1))

    # union_sums will contain 0's if a class is not present in an image, which will give division by zero
    iou_scores_classwise = intersection_sums / (union_sums + 0.00000000001)

    iou_scores_weighted = torch.sum(iou_scores_classwise * class_frequencies, dim=(-1))

    iou_score_batch_weighted_mean = torch.mean(iou_scores_weighted)

    return iou_score_batch_weighted_mean.numpy()


def rmse(targets, predictions):
    """
    Depth estimation evaluation method. Average Root Mean Squared Error for each pixel in an image.
    Should work for batches and single images.
    """
    return np.sqrt(np.average((targets-predictions)**2))


def accuracy_within_threshold(targets, predictions, threshold=1.25):
    """
    Depth estimation evaluation method. Calculates a delta value for each pixel in an image, and
    then checks which percentage of pixels are within a certain threshold.
    Should work for batches and single images.
    """

    targets_over_preds = targets / predictions
    preds_over_targets = predictions / targets

    deltas = np.maximum(targets_over_preds, preds_over_targets)

    within_threshold_matrix = deltas < threshold

    uniques, counts = np.unique(within_threshold_matrix, return_counts=True)

    if len(counts) > 1:
        accuracy = counts[1] / (counts[0] + counts[1])  # this will work as long as there are both True and False values
    else:
        if True in uniques:
            accuracy = 1.
            # print("Accuracy within threshold warning: Accuracy is 1. uniques:", uniques)# TODO uncomment for real eval
        else:
            accuracy = 0.
            print("Accuracy within threshold warning: Accuracy is 0. uniques:", uniques)
    return accuracy


def compare_models(data_folder, segmentation_models, depth_models, batch_size=1, max_n_instances=None,
                   n_classes=len(DEFAULT_CLASSES)+1):
    targets = ComparisonDataset(data_folder, segmentation_models, depth_models, max_n_instances=max_n_instances)

    dataloader = DataLoader(targets, batch_size=batch_size, shuffle=False, num_workers=0,
                            pin_memory=True)

    # semantic segmentation metrics
    mean_intersection_over_union_accumulated = defaultdict(int)
    weighted_mean_intersection_over_union_accumulated = defaultdict(int)

    # depth estimation metrics
    accuracy_with_threshold_accumulated = defaultdict(int)
    accuracy_with_threshold2_accumulated = defaultdict(int)
    accuracy_with_threshold3_accumulated = defaultdict(int)
    rmse_accumulated = defaultdict(int)

    classwise_iou_accumulated = defaultdict(functools.partial(np.zeros, n_classes))
    classwise_iou_class_counts = defaultdict(functools.partial(np.zeros, n_classes))

    for rgb_targets, segmentation_targets, depth_targets, segmentation_preds, depth_preds in tqdm(dataloader):
        #print("SEMANTIC SEGMENTATION:")
        #pepe = depth_targets[0].numpy().transpose(1, 2, 0)
        #plot_image(depth_targets[0].numpy().transpose(1, 2, 0), title="ground truth")
        #plot_image(depth_targets[0].numpy().transpose(1, 2, 0), title="ground truth gray", cmap="gray")
        for model in segmentation_preds:
            mean_iou, batch_classwise_iou = mean_jaccard_index(segmentation_targets, segmentation_preds[model])
            mean_intersection_over_union_accumulated[model] += mean_iou
            for img_classwise_iou in batch_classwise_iou:
                classwise_iou_accumulated[model] += img_classwise_iou

                # count if class actually is in img, to get correct averages
                for i_class in range(len(img_classwise_iou)):
                    if img_classwise_iou[i_class] > 0:
                        classwise_iou_class_counts[model][i_class] += 1

            weighted_mean_intersection_over_union_accumulated[model] \
                += weighted_jaccard_index(segmentation_targets, segmentation_preds[model])

            #img = segmentation_preds[model].numpy()[0].transpose(1, 2, 0)
            #plot_segmentation(img, title=model)

        #print("\nDEPTH ESTIMATION")
        for model in depth_preds:
            accuracy_with_threshold_accumulated[model] += accuracy_within_threshold(depth_targets, depth_preds[model],
                                                                                    threshold=1.25)
            accuracy_with_threshold2_accumulated[model] += accuracy_within_threshold(depth_targets, depth_preds[model],
                                                                                    threshold=1.25**2)
            accuracy_with_threshold3_accumulated[model] += accuracy_within_threshold(depth_targets, depth_preds[model],
                                                                                    threshold=1.25**3)
            rmse_accumulated[model] += rmse(depth_targets, depth_preds[model])

            #img = depth_preds[model].numpy()[0].transpose(1, 2, 0)
            #plot_image(img, title=model, cmap="gray")

    n_batches = np.ceil(len(targets) / batch_size)

    # calculate average over batches, semantic segmentation
    mean_intersection_over_union_avg = {}
    weighted_mean_intersection_over_union_avg = {}
    class_intersection_over_union_avg = defaultdict(functools.partial(np.zeros, n_classes))
    for model in segmentation_models:
        model_name = model[0]
        mean_intersection_over_union_avg[model_name] = mean_intersection_over_union_accumulated[model_name] / n_batches
        weighted_mean_intersection_over_union_avg[model_name] = weighted_mean_intersection_over_union_accumulated[model_name] / n_batches

        for i_class in range(len(classwise_iou_accumulated[model_name])):
            class_intersection_over_union_avg[model_name][i_class] = classwise_iou_accumulated[model_name][i_class] / (classwise_iou_class_counts[model_name][i_class]+0.0000000001)

        print("---")
        print("Model:", model_name, "has mean jaccard index avg:", mean_intersection_over_union_avg[model_name])
        print("Model:", model_name, "has weighted jaccard index avg:", weighted_mean_intersection_over_union_avg[model_name])
        print("Model:", model_name, "has classwise iou's:", [i for i in class_intersection_over_union_avg[model_name]])
        print("---")

    # calculate average over batches, depth estimation
    accuracy_within_threshold_avg = {}
    accuracy_within_threshold2_avg = {}
    accuracy_within_threshold3_avg = {}
    rmse_avg = {}
    for model in depth_models:
        model_name = model[0]
        accuracy_within_threshold_avg[model_name] = accuracy_with_threshold_accumulated[model_name] / n_batches
        accuracy_within_threshold2_avg[model_name] = accuracy_with_threshold2_accumulated[model_name] / n_batches
        accuracy_within_threshold3_avg[model_name] = accuracy_with_threshold3_accumulated[model_name] / n_batches
        rmse_avg[model_name] = rmse_accumulated[model_name] / n_batches
        print("---")
        print("Model:", model_name, "has accuracy within threshold avg:", accuracy_within_threshold_avg[model_name])
        print("Model:", model_name, "has accuracy within threshold2 avg:", accuracy_within_threshold2_avg[model_name])
        print("Model:", model_name, "has accuracy within threshold3 avg:", accuracy_within_threshold3_avg[model_name])
        print("Model:", model_name, "has rmse avg:", rmse_avg[model_name])
        print("---")


if __name__ == "__main__":

    test = "test2"
    # location of where to find training, test1, test2
    data_folder = Path("data/perception") / test
    predictions_folder = Path("data/perception/predictions")

    # lagres på formatet (Navn, lokasjon)
    segmentation_models = [("unet_resnet50", predictions_folder / "semseg/unet_resnet50" / test),
                           ("unet_resnet50_weighted_2.5", predictions_folder / "semseg/unet_resnet50_weighted_2.5" / test),
                           ("unet_resnet50_weighted_5", predictions_folder / "semseg/unet_resnet50_weighted_5" / test),
                           ("fcn_resnet101", predictions_folder / "semseg/fcn_resnet101" / test),
                           ("deeplabv3-mobilenet", predictions_folder / "semseg/deeplabv3_mobilenet" / test),
                           ("deeplabv3-resnet50", predictions_folder / "semseg/deeplabv3_resnet50" / test),
                           ("deeplabv3-resnet101", predictions_folder / "semseg/deeplabv3_resnet101" / test),
                           ]
                           #("semantic-test1 (ground truf)", predictions_folder / "semantic_test1"),
                           #("semantic-test2 (ground truf)", predictions_folder / "semantic_test2")]

    # lagres på formatet (Navn, lokasjon, invert_pixels_in_loading)
    # ("test1-depth", data_folder / "depth", False)
    depth_models = [("midas-small", predictions_folder / "depth/midas_small" / test, True),
                    ("midas-large", predictions_folder / "depth/midas_large" / test, True),
                    ("UNet", predictions_folder / "depth/unet" / test, False),
                    ("UNet-resnet34", predictions_folder / "depth/unet_resnet34" / test, False)
                    ]
                    #("depth-test1", predictions_folder / "depth_test1", False),
                    #("depth-test2", predictions_folder / "depth_test2", False)
                    #]

    compare_models(data_folder, segmentation_models, depth_models, batch_size=20, max_n_instances=None)


import functools
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

from perception.custom_datasets import ComparisonDataset
from perception.utils.visualization import plot_segmentation, plot_image, display_images_horizontally
from perception.utils.segmentation_labels import DEFAULT_CLASSES


def mean_jaccard_index(target, predictions):
    """
    Semantic segmentation metric. Calculates mean intersection over union over all classes.
    Works for batches of data.
    """

    intersection = torch.logical_and(target, predictions)
    union = torch.logical_or(target, predictions)

    intersection_sums = torch.sum(intersection, dim=(-2, -1))
    union_sums = torch.sum(union, dim=(-2,-1))

    class_exists_mask = union_sums != 0

    # union_sums will contain 0's if a class is not present in an image, which will give division by zero
    iou_scores_classwise = intersection_sums / (union_sums + 0.00000000001)

    iou_scores_imagewise_sum = iou_scores_classwise.sum(dim=1)
    class_exists_mask_sum = class_exists_mask.sum(dim=1)
    iou_scores_imagewise_mean = iou_scores_imagewise_sum / class_exists_mask_sum

    iou_score_batch_mean = torch.mean(iou_scores_imagewise_mean)

    return iou_score_batch_mean.numpy(), iou_scores_classwise.numpy()


def weighted_jaccard_index(target, predictions):
    """
    Semantic segmentation metric. Calculates mean intersection over union over all classes, weighted by class.
    Works for batches of data.
    """

    class_counts = torch.sum(target, dim=(-2, -1), dtype=torch.int32)
    n_pixels = torch.sum(class_counts, dim=(-1))[0]
    class_frequencies = class_counts / n_pixels

    intersection = torch.logical_and(target, predictions)
    union = torch.logical_or(target, predictions)

    intersection_sums = torch.sum(intersection, dim=(-2, -1))
    union_sums = torch.sum(union, dim=(-2,-1))

    # union_sums will contain 0's if a class is not present in an image, which will give division by zero
    iou_scores_classwise = intersection_sums / (union_sums + 0.00000000001)

    iou_scores_weighted = torch.sum(iou_scores_classwise * class_frequencies, dim=(-1))

    iou_score_batch_weighted_mean = torch.mean(iou_scores_weighted)

    return iou_score_batch_weighted_mean.numpy()


def rmse(targets, predictions):
    """
    Depth estimation evaluation method. Average Root Mean Squared Error for each pixel in an image.
    Should work for batches and single images.
    """
    return np.sqrt(np.average((targets-predictions)**2))


def accuracy_within_threshold(targets, predictions, threshold=1.25):
    """
    Depth estimation evaluation method. Calculates a delta value for each pixel in an image, and
    then checks which percentage of pixels are within a certain threshold.
    Should work for batches and single images.
    """

    targets_over_preds = targets / predictions
    preds_over_targets = predictions / targets

    deltas = np.maximum(targets_over_preds, preds_over_targets)

    within_threshold_matrix = deltas < threshold

    uniques, counts = np.unique(within_threshold_matrix, return_counts=True)

    if len(counts) > 1:
        accuracy = counts[1] / (counts[0] + counts[1])  # this will work as long as there are both True and False values
    else:
        if True in uniques:
            accuracy = 1.
            # print("Accuracy within threshold warning: Accuracy is 1. uniques:", uniques)# TODO uncomment for real eval
        else:
            accuracy = 0.
            print("Accuracy within threshold warning: Accuracy is 0. uniques:", uniques)
    return accuracy


def compare_models(data_folder, segmentation_models, depth_models, batch_size=1, max_n_instances=None,
                   n_classes=len(DEFAULT_CLASSES)+1):
    targets = ComparisonDataset(data_folder, segmentation_models, depth_models, max_n_instances=max_n_instances)

    dataloader = DataLoader(targets, batch_size=batch_size, shuffle=False, num_workers=0,
                            pin_memory=True)

    # semantic segmentation metrics
    mean_intersection_over_union_accumulated = defaultdict(int)
    weighted_mean_intersection_over_union_accumulated = defaultdict(int)

    # depth estimation metrics
    accuracy_with_threshold_accumulated = defaultdict(int)
    accuracy_with_threshold2_accumulated = defaultdict(int)
    accuracy_with_threshold3_accumulated = defaultdict(int)
    rmse_accumulated = defaultdict(int)

    classwise_iou_accumulated = defaultdict(functools.partial(np.zeros, n_classes))
    classwise_iou_class_counts = defaultdict(functools.partial(np.zeros, n_classes))

    for rgb_targets, segmentation_targets, depth_targets, segmentation_preds, depth_preds in tqdm(dataloader):
        #print("SEMANTIC SEGMENTATION:")
        #pepe = depth_targets[0].numpy().transpose(1, 2, 0)
        #plot_image(depth_targets[0].numpy().transpose(1, 2, 0), title="ground truth")
        #plot_image(depth_targets[0].numpy().transpose(1, 2, 0), title="ground truth gray", cmap="gray")
        for model in segmentation_preds:
            mean_iou, batch_classwise_iou = mean_jaccard_index(segmentation_targets, segmentation_preds[model])
            mean_intersection_over_union_accumulated[model] += mean_iou
            for img_classwise_iou in batch_classwise_iou:
                classwise_iou_accumulated[model] += img_classwise_iou

                # count if class actually is in img, to get correct averages
                for i_class in range(len(img_classwise_iou)):
                    if img_classwise_iou[i_class] > 0:
                        classwise_iou_class_counts[model][i_class] += 1

            weighted_mean_intersection_over_union_accumulated[model] \
                += weighted_jaccard_index(segmentation_targets, segmentation_preds[model])

            #img = segmentation_preds[model].numpy()[0].transpose(1, 2, 0)
            #plot_segmentation(img, title=model)

        #print("\nDEPTH ESTIMATION")
        for model in depth_preds:
            accuracy_with_threshold_accumulated[model] += accuracy_within_threshold(depth_targets, depth_preds[model],
                                                                                    threshold=1.25)
            accuracy_with_threshold2_accumulated[model] += accuracy_within_threshold(depth_targets, depth_preds[model],
                                                                                    threshold=1.25**2)
            accuracy_with_threshold3_accumulated[model] += accuracy_within_threshold(depth_targets, depth_preds[model],
                                                                                    threshold=1.25**3)
            rmse_accumulated[model] += rmse(depth_targets, depth_preds[model])

            #img = depth_preds[model].numpy()[0].transpose(1, 2, 0)
            #plot_image(img, title=model, cmap="gray")

    n_batches = np.ceil(len(targets) / batch_size)

    # calculate average over batches, semantic segmentation
    mean_intersection_over_union_avg = {}
    weighted_mean_intersection_over_union_avg = {}
    class_intersection_over_union_avg = defaultdict(functools.partial(np.zeros, n_classes))
    for model in segmentation_models:
        model_name = model[0]
        mean_intersection_over_union_avg[model_name] = mean_intersection_over_union_accumulated[model_name] / n_batches
        weighted_mean_intersection_over_union_avg[model_name] = weighted_mean_intersection_over_union_accumulated[model_name] / n_batches

        for i_class in range(len(classwise_iou_accumulated[model_name])):
            class_intersection_over_union_avg[model_name][i_class] = classwise_iou_accumulated[model_name][i_class] / (classwise_iou_class_counts[model_name][i_class]+0.0000000001)

        print("---")
        print("Model:", model_name, "has mean jaccard index avg:", mean_intersection_over_union_avg[model_name])
        print("Model:", model_name, "has weighted jaccard index avg:", weighted_mean_intersection_over_union_avg[model_name])
        print("Model:", model_name, "has classwise iou's:", [i for i in class_intersection_over_union_avg[model_name]])
        print("---")

    # calculate average over batches, depth estimation
    accuracy_within_threshold_avg = {}
    accuracy_within_threshold2_avg = {}
    accuracy_within_threshold3_avg = {}
    rmse_avg = {}
    for model in depth_models:
        model_name = model[0]
        accuracy_within_threshold_avg[model_name] = accuracy_with_threshold_accumulated[model_name] / n_batches
        accuracy_within_threshold2_avg[model_name] = accuracy_with_threshold2_accumulated[model_name] / n_batches
        accuracy_within_threshold3_avg[model_name] = accuracy_with_threshold3_accumulated[model_name] / n_batches
        rmse_avg[model_name] = rmse_accumulated[model_name] / n_batches
        print("---")
        print("Model:", model_name, "has accuracy within threshold avg:", accuracy_within_threshold_avg[model_name])
        print("Model:", model_name, "has accuracy within threshold2 avg:", accuracy_within_threshold2_avg[model_name])
        print("Model:", model_name, "has accuracy within threshold3 avg:", accuracy_within_threshold3_avg[model_name])
        print("Model:", model_name, "has rmse avg:", rmse_avg[model_name])
        print("---")


if __name__ == "__main__":

    test = "test2"
    # location of where to find training, test1, test2
    data_folder = Path("data/perception") / test
    predictions_folder = Path("data/perception/predictions")

    # lagres på formatet (Navn, lokasjon)
    segmentation_models = [("unet_resnet50", predictions_folder / "semseg/unet_resnet50" / test),
                           ("unet_resnet50_weighted_2.5", predictions_folder / "semseg/unet_resnet50_weighted_2.5" / test),
                           ("unet_resnet50_weighted_5", predictions_folder / "semseg/unet_resnet50_weighted_5" / test),
                           ("fcn_resnet101", predictions_folder / "semseg/fcn_resnet101" / test),
                           ("deeplabv3-mobilenet", predictions_folder / "semseg/deeplabv3_mobilenet" / test),
                           ("deeplabv3-resnet50", predictions_folder / "semseg/deeplabv3_resnet50" / test),
                           ("deeplabv3-resnet101", predictions_folder / "semseg/deeplabv3_resnet101" / test),
                           ]
                           #("semantic-test1 (ground truf)", predictions_folder / "semantic_test1"),
                           #("semantic-test2 (ground truf)", predictions_folder / "semantic_test2")]

    # lagres på formatet (Navn, lokasjon, invert_pixels_in_loading)
    # ("test1-depth", data_folder / "depth", False)
    depth_models = [("midas-small", predictions_folder / "depth/midas_small" / test, True),
                    ("midas-large", predictions_folder / "depth/midas_large" / test, True),
                    ("UNet", predictions_folder / "depth/unet" / test, False),
                    ("UNet-resnet34", predictions_folder / "depth/unet_resnet34" / test, False)
                    ]
                    #("depth-test1", predictions_folder / "depth_test1", False),
                    #("depth-test2", predictions_folder / "depth_test2", False)
                    #]

    compare_models(data_folder, segmentation_models, depth_models, batch_size=20, max_n_instances=None)


Semantic segmentation metric. Calculates mean intersection over union over all classes.     Works for batches of data. # union_sums will contain 0's if a class is not present in an image, which will give division by zero Semantic segmentation metric. Calculates mean intersection over union over all classes, weighted by class.     Works for batches of data. # union_sums will contain 0's if a class is not present in an image, which will give division by zero Depth estimation evaluation method. Average Root Mean Squared Error for each pixel in an image.     Should work for batches and single images. Depth estimation evaluation method. Calculates a delta value for each pixel in an image, and     then checks which percentage of pixels are within a certain threshold.     Should work for batches and single images. # this will work as long as there are both True and False values # print("Accuracy within threshold warning: Accuracy is 1. uniques:", uniques)# TODO uncomment for real eval # semantic segmentation metrics # depth estimation metrics #print("SEMANTIC SEGMENTATION:") #pepe = depth_targets[0].numpy().transpose(1, 2, 0) #plot_image(depth_targets[0].numpy().transpose(1, 2, 0), title="ground truth") #plot_image(depth_targets[0].numpy().transpose(1, 2, 0), title="ground truth gray", cmap="gray") # count if class actually is in img, to get correct averages #img = segmentation_preds[model].numpy()[0].transpose(1, 2, 0) #plot_segmentation(img, title=model) #print("\nDEPTH ESTIMATION") #img = depth_preds[model].numpy()[0].transpose(1, 2, 0) #plot_image(img, title=model, cmap="gray") # calculate average over batches, semantic segmentation # calculate average over batches, depth estimation # location of where to find training, test1, test2 # lagres på formatet (Navn, lokasjon) #("semantic-test1 (ground truf)", predictions_folder / "semantic_test1"), #("semantic-test2 (ground truf)", predictions_folder / "semantic_test2")] # lagres på formatet (Navn, lokasjon, invert_pixels_in_loading) # ("test1-depth", data_folder / "depth", False) #("depth-test1", predictions_folder / "depth_test1", False), #("depth-test2", predictions_folder / "depth_test2", False) #]