import torch
from typing import Tuple


def intersectionAndUnionGPU(
    output: torch.Tensor,
    target: torch.Tensor,
    K: int,
    ignore_index: int=255,
    cuda_available: bool=True
    ) -> Tuple[torch.Tensor,torch.Tensor,torch.Tensor]:
    """
    Note output and target sizes are N or N * L or N * H * W 
        Args:
        -   output: Pytorch tensor represeting predicted label map,
                each value in range 0 to K - 1.
        -   target: Pytorch tensor representing ground truth label map,
                each value in range 0 to K - 1.
        -   K: integer number of possible classes
        -   ignore_index: integer representing class index to ignore
        -   cuda_available: CUDA is available to Pytorch to use

        Returns:
        -   area_intersection: 1d Pytorch tensor of length (K,) with counts
                for each of K classes, where pred & target matched
        -   area_union: 1d Pytorch tensor of length (K,) with counts
        -   area_target: 1d Pytorch tensor of length (K,) with bin counts
                for each of K classes, present in this GT label map.
    """
    assert (output.dim() in [1, 2, 3])
    assert output.shape == target.shape
    output = output.view(-1)
    target = target.view(-1)
    output[target == ignore_index] = ignore_index
    intersection = output[output == target]
    # https://github.com/pytorch/pytorch/issues/1382
    area_intersection = torch.histc(intersection.float().cpu(), bins=K, min=0, max=K-1)
    area_output = torch.histc(output.float().cpu(), bins=K, min=0, max=K-1)
    area_target = torch.histc(target.float().cpu(), bins=K, min=0, max=K-1)
    area_union = area_output + area_target - area_intersection
    if cuda_available:
        return area_intersection.cuda(), area_union.cuda(), area_target.cuda()
    else:
        return area_intersection, area_union, area_target