import torch
import slideflow as sf
import rasterio
import numpy as np
import shapely.affinity as sa

from typing import Tuple, Union, Optional, List, Dict
from torchvision import transforms
from os.path import join, exists
from rich.progress import track
from shapely.ops import unary_union
from shapely.geometry import Polygon
from shapely.ops import unary_union
from slideflow.util import path_to_name

from .utils import topleft_pad_torch

# -----------------------------------------------------------------------------

class ThumbMaskDataset(torch.utils.data.Dataset):

    def __init__(
        self,
        dataset: "sf.Dataset",
        mpp: float,
        roi_labels: List[str],
        *,
        mode: str = 'binary',
    ) -> None:
        """Dataset that generates thumbnails and ROI masks.

        Args:
            dataset (sf.Dataset): The dataset to use.
            mpp (float): The target microns per pixel. The thumbnail will be
                scaled to this resolution.
            roi_labels (List[str]): The ROI labels to include in the mask.

        Keyword args:
            mode (str, optional): The mode to use for the mask. One of:
                'binary', 'multiclass', 'multilabel'. Defaults to 'binary'.

        """
        super().__init__()
        self.mpp = mpp
        self.mode = mode
        self.roi_labels = roi_labels

        # Subsample dataset to only include slides with ROIs.
        self.rois = dataset.rois()
        slides = set(map(path_to_name, dataset.slide_paths()))
        slides = slides.intersection(set(map(path_to_name, self.rois)))
        dataset = dataset.filter({'slide': list(slides)})

        # Prepare WSI objects (for slides with ROIs).
        self.paths = dataset.slide_paths()

    def __len__(self) -> int:
        return len(self.paths)

    def process(
        self,
        img: np.ndarray,
        mask: np.ndarray
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Process the image/mask and convert to a tensor."""
        img = torch.from_numpy(img)
        mask = torch.from_numpy(mask)
        return img, mask

    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:

        # Load the image and mask.
        path = self.paths[index]
        wsi = sf.WSI(path, 299, 512, rois=self.rois, roi_filter_method=0.1, verbose=False)
        output = get_thumb_and_mask(wsi, self.mpp, self.roi_labels, skip_missing=False)
        if output is None:
            return None
        img = output['image']               # CHW (np.ndarray)
        mask = output['mask'].astype(int)   # 1HW (np.ndarray)

        if self.mode == 'multiclass':
            mask = mask * np.arange(1, mask.shape[0]+1)[:, None, None]
            mask = mask.max(axis=0)
        elif self.mode == 'binary' and mask.ndim == 3:
            mask = np.any(mask, axis=0)[None, :, :].astype(int)

        # Process.
        img, mask = self.process(img, mask)

        return {
            'image': img,
            'mask': mask
        }


class RandomCropDataset(ThumbMaskDataset):

    def __init__(self, *args, size: int = 1024, **kwargs):
        """Dataset that generates thumbnails & ROI masks, with random crops.

        Thumbnails and masks and randomly cropped and rotated together to
        a square size of `size` pixels.

        Args:
            dataset (sf.Dataset): The dataset to use.
            mpp (float): The target microns per pixel. The thumbnail will be
                scaled to this resolution.
            roi_labels (List[str]): The ROI labels to include in the mask.
            size (int, optional): The size of the random crop. Defaults to 1024.

        Keyword Args:
            mode (str, optional): The mode to use for the mask. One of:
                'binary', 'multiclass', 'multilabel'. Defaults to 'binary'.

        """
        super().__init__(*args, **kwargs)
        self.size = size

    def process(self, img, mask):
        """Randomly crop/rotate the image and mask and convert to a tensor."""
        return random_crop_and_rotate(img, mask, size=self.size)

# -----------------------------------------------------------------------------
# Buffered datasets

class BufferedMaskDataset(torch.utils.data.Dataset):

    def __init__(self, dataset: "sf.Dataset", source: str, *, mode: str = 'binary'):
        """Dataset that loads buffered image and mask pairs.

        Args:
            dataset (sf.Dataset): The dataset to use.
            source (str): The directory containing the buffered image/mask pairs.

        Keyword Args:
            mode (str, optional): The mode to use for the mask. One of:
                'binary', 'multiclass', 'multilabel'. Defaults to 'binary'.

        """
        super().__init__()
        if mode not in ['binary', 'multiclass', 'multilabel']:
            raise ValueError("Invalid mode: {}. Expected one of: binary, "
                             "multiclass, multilabel".format(mode))
        self.dataset = dataset
        self.mode = mode
        self.paths = [
            join(source, s + '.pt') for s in dataset.slides()
            if exists(join(source, s + '.pt'))
        ]


    def __len__(self) -> int:
        return len(self.paths)

    def process(
        self,
        img: np.ndarray,
        mask: np.ndarray
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Process the image/mask and convert to a tensor."""
        img = torch.from_numpy(img)
        mask = torch.from_numpy(mask)
        return img, mask

    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
        # Load the image and mask.
        output = torch.load(self.paths[index])
        img = output['image']               # CHW (np.ndarray)
        mask = output['mask'].astype(int)   # 1HW (np.ndarray)

        if self.mode == 'multiclass':
            mask = mask * np.arange(1, mask.shape[0]+1)[:, None, None]
            mask = mask.max(axis=0)
        elif self.mode == 'binary' and mask.ndim == 3:
            mask = np.any(mask, axis=0)[None, :, :].astype(int)

        # Process.
        img, mask = self.process(img, mask)

        return {
            'image': img,
            'mask': mask
        }


class BufferedRandomCropDataset(BufferedMaskDataset):

    def __init__(self, *args, size: int = 1024, **kwargs):
        """Dataset that loads buffered image/mask pairs and randomly crops.

        Loaded thumbnails and masks and randomly cropped and rotated together to
        a square size of `size` pixels.

        Args:
            dataset (sf.Dataset): The dataset to use.
            source (str): The directory containing the buffered image/mask pairs.
            size (int, optional): The size of the random crop. Defaults to 1024.

        Keyword Args:
            mode (str, optional): The mode to use for the mask. One of:
                'binary', 'multiclass', 'multilabel'. Defaults to 'binary'.

        """
        super().__init__(*args, **kwargs)
        self.size = size

    def process(
        self,
        img: np.ndarray,
        mask: np.ndarray
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Randomly crop/rotate the image and mask and convert to a tensor."""
        return random_crop_and_rotate(img, mask, size=self.size)

# -----------------------------------------------------------------------------

class TileMaskDataset(torch.utils.data.Dataset):

    def __init__(
        self,
        dataset: "sf.Dataset",
        tile_px: int,
        tile_um: Union[int, str],
        *,
        roi_labels: Optional[List[str]] = None,
        stride_div: int = 2,
        crop_margin: int = 0,
        filter_method: str = 'otsu',
        mode: str = 'binary'
    ):
        """Dataset that generates tiles and ROI masks from slides.

        Args:
            dataset (sf.Dataset): The dataset to use.
            tile_px (int): The size of the tiles (in pixels).
            tile_um (Union[int, str]): The size of the tiles (in microns).

        Keyword args:
            stride_div (int, optional): The divisor for the stride.
                Defaults to 2.
            crop_margin (int, optional): The number of pixels to add to the
                tile size before random cropping to the target tile_px size.
                Defaults to 0 (no random cropping).
            filter_method (str, optional): The method to use for identifying
                tiles for training. If 'roi', selects only tiles that intersect
                with an ROI. If 'otsu', selects tiles based on an Otsu threshold
                of the slide. Defaults to 'roi'.

        """
        super().__init__()

        rois = dataset.rois()
        slides_with_rois = [path_to_name(r) for r in rois]
        slides = [s for s in dataset.slide_paths()
                  if path_to_name(s) in slides_with_rois]
        kw = dict(
            tile_px=tile_px + crop_margin,
            tile_um=tile_um,
            verbose=False,
            stride_div=stride_div
        )
        if roi_labels is None:
            roi_labels = []
        self.mode = mode
        self.roi_labels = roi_labels
        self.tile_px = tile_px
        self.coords = []
        self.all_wsi = dict()
        self.all_wsi_with_roi = dict()
        self.all_extract_px = dict()
        for slide in track(slides, description="Loading slides"):
            name = path_to_name(slide)
            wsi = sf.WSI(slide, **kw)
            try:
                wsi_with_rois = sf.WSI(slide, roi_filter_method=0.1, rois=rois, **kw)
            except Exception as e:
                sf.log.error("Failed to load slide {}: {}".format(slide, e))
                raise e

            # Filter ROIs to only include the specified labels.
            if self.roi_labels:
                wsi_with_rois.rois = [roi for roi in wsi_with_rois.rois if roi.label in self.roi_labels]
                wsi_with_rois.process_rois()

            if not len(wsi_with_rois.rois):
                continue
            if filter_method == 'roi':
                wsi_inner = sf.WSI(slide, roi_filter_method=0.9, rois=rois, **kw)
                if self.roi_labels:
                    wsi_inner.rois = [roi for roi in wsi_with_rois.rois if roi.label in self.roi_labels]
                    wsi_inner.process_rois()
                coords = np.argwhere(wsi_with_rois.grid & (~wsi_inner.grid)).tolist()
            elif filter_method == 'otsu':
                wsi.qc('otsu')
                coords = np.argwhere(wsi.grid).tolist()
                wsi.remove_qc()
            elif filter_method in ['all', 'none', None]:
                coords = np.argwhere(wsi_with_rois.grid).tolist()
            else:
                raise ValueError("Invalid filter method: {}. Expected one of: "
                                 "roi, otsu".format(filter_method))
            for c in coords:
                self.coords.append([name] + c)
            self.all_wsi[name] = wsi
            self.all_wsi_with_roi[name] = wsi_with_rois
            self.all_extract_px[name] = int(wsi.tile_um / wsi.mpp)

    def __len__(self):
        return len(self.coords)

    def get_scaled_and_intersecting_polys(
        self,
        polys: "Polygon",
        tile: "Polygon",
        scale: float,
        full_stride: int,
        grid_idx: Tuple[int, int]
    ):
        """Get scaled and intersecting polygons for a given tile."""
        gx, gy = grid_idx
        A = polys.intersection(tile)

        # Translate polygons so the intersection origin is at (0, 0)
        B = sa.translate(A, -(full_stride*gx), -(full_stride*gy))

        # Scale to the target tile size
        C = sa.scale(B, xfact=scale, yfact=scale, origin=(0, 0))
        return C

    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """Get an image and mask for a given index."""
        slide, gx, gy = self.coords[index]
        wsi = self.all_wsi[slide]
        wsi_with_roi = self.all_wsi_with_roi[slide]
        fe = self.all_extract_px[slide]
        fs = wsi.full_stride
        scale = wsi.tile_px / fe

        # Get the image.
        img = wsi[gx, gy].transpose(2, 0, 1)

        # Get a polygon for the tile, used for determining overlapping ROIs.
        tile = Polygon([
            [fs*gx, fs*gy],
            [fs*gx, (fs*gy)+fe],
            [(fs*gx)+fe, (fs*gy)+fe],
            [(fs*gx)+fe, fs*gy]
        ])

        # Compute the mask from ROIs.
        if len(wsi_with_roi.rois) == 0:
            if self.roi_labels:
                mask = np.zeros((len(self.roi_labels), wsi.tile_px, wsi.tile_px), dtype=int)
            else:
                mask = np.zeros((1, wsi.tile_px, wsi.tile_px), dtype=int)

        # Handle ROIs with labels (multilabel or multiclass)
        elif self.roi_labels:
            labeled_masks = []
            for i, label in enumerate(self.roi_labels):
                wsi_polys = [p.poly for p in wsi_with_roi.rois if p.label == label]
                if len(wsi_polys) == 0:
                    mask = np.zeros((wsi.tile_px, wsi.tile_px), dtype=int)
                    labeled_masks.append(mask)
                else:
                    all_polys = unary_union(wsi_polys)
                    polys = self.get_scaled_and_intersecting_polys(
                        all_polys, tile, scale, fs, (gx, gy)
                    )
                    if isinstance(polys, Polygon) and polys.is_empty:
                        mask = np.zeros((wsi.tile_px, wsi.tile_px), dtype=int)
                    else:
                        # Rasterize to an int mask.
                        mask = rasterio.features.rasterize([polys], out_shape=[wsi.tile_px, wsi.tile_px]).astype(int)
                    labeled_masks.append(mask)
            mask = np.stack(labeled_masks, axis=0)

        # Handle ROIs without labels (binary)
        else:
            # Determine the intersection at the given tile location.
            all_polys = unary_union([p.poly for p in wsi_with_roi.rois])
            polys = self.get_scaled_and_intersecting_polys(
                all_polys, tile, scale, fs, (gx, gy)
            )

            if isinstance(polys, Polygon) and polys.is_empty:
                mask = np.zeros((wsi.tile_px, wsi.tile_px), dtype=int)
            else:
                # Rasterize to an int mask.
                try:
                    mask = rasterio.features.rasterize([polys], out_shape=[wsi.tile_px, wsi.tile_px]).astype(bool).astype(np.int32)
                except ValueError:
                    mask = np.zeros((wsi.tile_px, wsi.tile_px), dtype=int)

            # Add a dummy channel dimension.
            mask = mask[None, :, :]

        # Process according to the mode.
        if self.mode == 'multiclass':
            mask = mask * np.arange(1, mask.shape[0]+1)[:, None, None]
            mask = mask.max(axis=0)
        elif self.mode == 'binary' and mask.ndim == 3:
            mask = np.any(mask, axis=0)[None, :, :].astype(int)

        # Process.
        img, mask = self.process(img, mask)

        return {
            'image': img,
            'mask': mask
        }

    def process(
        self,
        img: np.ndarray,
        mask: np.ndarray
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Randomly crop/rotate the image and mask and convert to a tensor."""
        return random_crop_and_rotate(img, mask, size=self.tile_px)

# -----------------------------------------------------------------------------

def random_crop_and_rotate(img, mask, size):
    if mask.ndim == 2:
        to_squeeze = True
        mask = mask[None, :, :]
    else:
        to_squeeze = False

    # Convert to tensor.
    img = torch.from_numpy(img).permute(1, 2, 0)
    mask = torch.from_numpy(mask).permute(1, 2, 0)

    # Pad to target size.
    img = topleft_pad_torch(img, size).permute(2, 0, 1)
    mask = topleft_pad_torch(mask, size).permute(2, 0, 1)

    # Random crop.
    i, j, h, w = transforms.RandomCrop.get_params(
        img, output_size=(size, size))
    img = transforms.functional.crop(img, i, j, h, w)
    mask = transforms.functional.crop(mask, i, j, h, w)

    # Random flip.
    if np.random.rand() > 0.5:
        img = transforms.functional.hflip(img)
        mask = transforms.functional.hflip(mask)
    if np.random.rand() > 0.5:
        img = transforms.functional.vflip(img)
        mask = transforms.functional.vflip(mask)

    # Random cardinal rotation.
    r = np.random.randint(4)
    img = transforms.functional.rotate(img, r * 90)
    mask = transforms.functional.rotate(mask, r * 90)

    if to_squeeze:
        mask = mask.squeeze(0)

    return img, mask

# -----------------------------------------------------------------------------

def get_thumb_and_mask(
    wsi: "sf.WSI",
    mpp: float,
    roi_labels: Optional[List[str]] = None,
    skip_missing: bool = False
) -> Dict[str, np.ndarray]:
    """Get a thumbnail and segmentation mask for a slide."""

    if len(wsi.rois) == 0 and skip_missing:
        return None

    # Sanity check.
    width = int((wsi.mpp * wsi.dimensions[0]) / mpp)
    ds = wsi.dimensions[0] / width
    level = wsi.slide.best_level_for_downsample(ds)
    level_dim = wsi.slide.level_dimensions[level]
    if any([d > 10000 for d in level_dim]):
        sf.log.warning("Large thumbnail found ({}) at level={} for {}".format(
            level_dim, level, wsi.path
        ))

    # Get the thumbnail.
    thumb = wsi.thumb(mpp=mpp).convert('RGB')
    img = np.array(thumb).transpose(2, 0, 1)
    xfact = thumb.size[1] / wsi.dimensions[1]
    yfact = thumb.size[0] / wsi.dimensions[0]

    if len(wsi.rois) == 0:
        if roi_labels:
            mask = np.zeros((len(roi_labels), thumb.size[1], thumb.size[0])).astype(bool)
        else:
            mask = np.zeros((1, thumb.size[1], thumb.size[0])).astype(bool)
    elif roi_labels:
        labeled_masks = []
        for i, label in enumerate(roi_labels):
            wsi_polys = [p.poly for p in wsi.rois if p.label == label]
            if len(wsi_polys) == 0:
                mask = np.zeros((thumb.size[1], thumb.size[0])).astype(bool)
                labeled_masks.append(mask)
            else:
                all_polys = unary_union(wsi_polys)
                # Scale ROIs to the thumbnail size.
                C = sa.scale(all_polys, xfact=xfact, yfact=yfact, origin=(0, 0))
                # Rasterize to an int mask.
                mask = rasterio.features.rasterize([C], out_shape=(thumb.size[1], thumb.size[0])).astype(bool).astype(np.int32)
                labeled_masks.append(mask)
        mask = np.stack(labeled_masks, axis=0)

    else:
        all_polys = unary_union([p.poly for p in wsi.rois])
        # Scale ROIs to the thumbnail size.
        C = sa.scale(all_polys, xfact=xfact, yfact=yfact, origin=(0, 0))
        # Rasterize to an int mask.
        mask = rasterio.features.rasterize([C], out_shape=(thumb.size[1], thumb.size[0])).astype(bool)
        # Add a dummy channel dimension.
        mask = mask[None, :, :]

    assert img.shape[1:] == mask.shape[1:], "Image and mask must have the same dimensions."
    assert mask.ndim == 3, "Mask must have 3 dimensions (C, H, W)."
    assert img.ndim == 3, "Image must have 3 dimensions (C, H, W)."

    return {
        'image': img,
        'mask': mask
    }