from torchvision import transforms

import pandas as pd

import numpy as np

import time

import pdb

import PIL.Image as Image

import h5py

from torch.utils.data import Dataset

import torch

from wsi_core.util_classes import Contour_Checking_fn, isInContourV1, isInContourV2, isInContourV3_Easy, isInContourV3_Hard

def default_transforms(mean = (0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):

    t = transforms.Compose(

                        [transforms.ToTensor(),

                         transforms.Normalize(mean = mean, std = std)])

    return t

def get_contour_check_fn(contour_fn='four_pt_hard', cont=None, ref_patch_size=None, center_shift=None):

    if contour_fn == 'four_pt_hard':

        cont_check_fn = isInContourV3_Hard(contour=cont, patch_size=ref_patch_size, center_shift=center_shift)

    elif contour_fn == 'four_pt_easy':

        cont_check_fn = isInContourV3_Easy(contour=cont, patch_size=ref_patch_size, center_shift=0.5)

    elif contour_fn == 'center':

        cont_check_fn = isInContourV2(contour=cont, patch_size=ref_patch_size)

    elif contour_fn == 'basic':

        cont_check_fn = isInContourV1(contour=cont)

    else:

        raise NotImplementedError

    return cont_check_fn

class Wsi_Region(Dataset):

    '''

    args:

        wsi_object: instance of WholeSlideImage wrapper over a WSI

        top_left: tuple of coordinates representing the top left corner of WSI region (Default: None)

        bot_right tuple of coordinates representing the bot right corner of WSI region (Default: None)

        level: downsample level at which to prcess the WSI region

        patch_size: tuple of width, height representing the patch size

        step_size: tuple of w_step, h_step representing the step size

        contour_fn (str): 

            contour checking fn to use

            choice of ['four_pt_hard', 'four_pt_easy', 'center', 'basic'] (Default: 'four_pt_hard')

        t: custom torchvision transformation to apply 

        custom_downsample (int): additional downscale factor to apply 

        use_center_shift: for 'four_pt_hard' contour check, how far out to shift the 4 points

    '''

    def __init__(self, wsi_object, top_left=None, bot_right=None, level=0, 

                 patch_size = (256, 256), step_size=(256, 256), 

                 contour_fn='four_pt_hard',

                 t=None, custom_downsample=1, use_center_shift=False):

        self.custom_downsample = custom_downsample

        # downscale factor in reference to level 0

        self.ref_downsample = wsi_object.level_downsamples[level]

        # patch size in reference to level 0

        self.ref_size = tuple((np.array(patch_size) * np.array(self.ref_downsample)).astype(int)) 

        if self.custom_downsample > 1:

            self.target_patch_size = patch_size

            patch_size = tuple((np.array(patch_size) * np.array(self.ref_downsample) * custom_downsample).astype(int))

            step_size = tuple((np.array(step_size) * custom_downsample).astype(int))

            self.ref_size = patch_size

        else:

            step_size = tuple((np.array(step_size)).astype(int))

            self.ref_size = tuple((np.array(patch_size) * np.array(self.ref_downsample)).astype(int)) 

        self.wsi = wsi_object.wsi

        self.level = level

        self.patch_size = patch_size

        if not use_center_shift:

            center_shift = 0.

        else:

            overlap = 1 - float(step_size[0] / patch_size[0])

            if overlap < 0.25:

                center_shift = 0.375

            elif overlap >= 0.25 and overlap < 0.75:

                center_shift = 0.5

            elif overlap >=0.75 and overlap < 0.95:

                center_shift = 0.5

            else:

                center_shift = 0.625

            #center_shift = 0.375 # 25% overlap

            #center_shift = 0.625 #50%, 75% overlap

            #center_shift = 1.0 #95% overlap

        filtered_coords = []

        #iterate through tissue contours for valid patch coordinates

        for cont_idx, contour in enumerate(wsi_object.contours_tissue): 

            print('processing {}/{} contours'.format(cont_idx, len(wsi_object.contours_tissue)))

            cont_check_fn = get_contour_check_fn(contour_fn, contour, self.ref_size[0], center_shift)

            coord_results, _ = wsi_object.process_contour(contour, wsi_object.holes_tissue[cont_idx], level, '', 

                            patch_size = patch_size[0], step_size = step_size[0], contour_fn=cont_check_fn,

                            use_padding=True, top_left = top_left, bot_right = bot_right)

            if len(coord_results) > 0:

                filtered_coords.append(coord_results['coords'])

        coords=np.vstack(filtered_coords)

        self.coords = coords

        print('filtered a total of {} coordinates'.format(len(self.coords)))

        # apply transformation

        if t is None:

            self.transforms = default_transforms()

        else:

            self.transforms = t

    def __len__(self):

        return len(self.coords)

    def __getitem__(self, idx):

        coord = self.coords[idx]

        patch = self.wsi.read_region(tuple(coord), self.level, self.patch_size).convert('RGB')

        if self.custom_downsample > 1:

            patch = patch.resize(self.target_patch_size)

        patch = self.transforms(patch).unsqueeze(0)

        return patch, coord