import argparse

import os

import time

import numpy as np

import openslide

import cv2

from PIL import Image, ImageDraw

from shapely.affinity import scale

from shapely.geometry import Polygon, MultiPolygon

from shapely.ops import unary_union

from collections import defaultdict

import nmslib

from sklearn.metrics.pairwise import euclidean_distances, cosine_similarity

# Optional if stain deconvolution is used.

import histomicstk as htk #pip install histomicstk --find-links https://girder.github.io/large_image_wheels

import torch

from torch.utils.data import DataLoader, Dataset

from torchvision import transforms

from esvit.utils import bool_flag

# You can use your own encoder to extract features. Here's examples including EsVIT the one used in the publication. 

from encoders import load_encoder_esVIT, load_encoder_resnet

def get_args_parser():

    parser = argparse.ArgumentParser('Preprocessing script esvit', add_help=False)

    parser.add_argument(

        "--input_slide",

        type=str,

        help="Path to input WSI file",

    )

    parser.add_argument(

        "--output_dir",

        type=str,

        help="Directory to save output data",

    )

    parser.add_argument(

        "--checkpoint",

        type=str,

        help="Feature extractor weights checkpoint",

    )

    parser.add_argument(

        "--batch_size",

        type=int,

        default=512,

    )

    parser.add_argument(

        "--tile_size",

        help="Desired tile size in microns (should be the same value as used in feature extraction model).",

        type=int,

        required=True,

    )

    parser.add_argument(

        "--out_size",

        help="Resize the square tile to this output size (in pixels).",

        type=int,

        default=224,

    )

    parser.add_argument(

        "--method",

        help="Segmentation method, otsu or stain deconv",

        type=str,

        default='otsu',

    )

    parser.add_argument(

        "--dist_threshold",

        type=int,

        default=4,

        help="L2 norm distance when spatially merging pacthes.",

    )

    parser.add_argument(

        "--corr_threshold",

        type=float,

        default=0.6,

        help="Cosine similarity distance when semantically merging pacthes.",

    )

    parser.add_argument(

        "--workers",

        help="The number of workers to use for the data loader. Only relevant when using a GPU.",

        type=int,

        default=4,

    )

    parser.add_argument(

        '--cfg',

        help='experiment configure file name. See EsVIT repo.',

        type=str

    )

    parser.add_argument(

        '--arch', default='deit_small', type=str,

        choices=['cvt_tiny', 'swin_tiny','swin_small', 'swin_base', 'swin_large', 'swin', 'vil', 'vil_1281', 'vil_2262', 'deit_tiny', 'deit_small', 'vit_base'],

        help="""Name of architecture to train. For quick experiments with ViTs, we recommend using deit_tiny or deit_small. See EsVIT repo."""

    )

    parser.add_argument(

        '--n_last_blocks', 

        default=4, 

        type=int, 

        help="""Concatenate [CLS] tokens for the `n` last blocks. We use `n=4` when evaluating DeiT-Small and `n=1` with ViT-Base. See EsVIT repo."""

    )

    parser.add_argument(

        '--avgpool_patchtokens', 

        default=False, 

        type=bool_flag,

        help="""Whether ot not to concatenate the global average pooled features to the [CLS] token.

        We typically set this to False for DeiT-Small and to True with ViT-Base. See EsVIT repo."""

    )

    parser.add_argument(

        '--patch_size', 

        default=8, 

        type=int, 

        help='Patch resolution of the model. See EsVIT repo.'

    )

    parser.add_argument(

        'opts',

        help="Modify config options using the command-line. See EsVIT repo.",

        default=None,

        nargs=argparse.REMAINDER

    )

    parser.add_argument(

        "--rank", 

        default=0, 

        type=int, 

        help="Please ignore and do not set this argument.")

    return parser

def segment_tissue(img):

    img_hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)

    mthresh = 7

    img_med = cv2.medianBlur(img_hsv[:, :, 1], mthresh)

    _, img_prepped = cv2.threshold(img_med, 0, 255, cv2.THRESH_OTSU + cv2.THRESH_BINARY)

    close = 4

    kernel = np.ones((close, close), np.uint8)

    img_prepped = cv2.morphologyEx(img_prepped, cv2.MORPH_CLOSE, kernel)

    # Find and filter contours

    contours, hierarchy = cv2.findContours(

        img_prepped, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE

    )

    return contours, hierarchy

def segment_tissue_deconv_stain(img):

    """

    Method 2: Tissue segmentation using stain deconvolution. Alternative to Otsu thresholding. 

    """

    image = img.copy()

    image[image[...,-1]==0] = [255,255,255,0]

    image = Image.fromarray(image)

    image = np.asarray(image.convert('RGB'))

    I_0 = 255

    # Create stain to color map

    stain_color_map = htk.preprocessing.color_deconvolution.stain_color_map

    # Specify stains of input image

    stains = ['hematoxylin',  # nuclei stain

              'eosin']        # cytoplasm stain

    w_est = htk.preprocessing.color_deconvolution.rgb_separate_stains_macenko_pca(image, I_0)

    deconv_result = htk.preprocessing.color_deconvolution.color_deconvolution(image, w_est, I_0)

    final_mask = np.zeros(image.shape[0:2], np.uint8)

    for i in 0, 1: 

        channel = htk.preprocessing.color_deconvolution.find_stain_index(

            stain_color_map[stains[i]], w_est)

        img_for_thresholding = 255 - deconv_result.Stains[:, :, channel]

        _, img_prepped = cv2.threshold(

            img_for_thresholding, 0, 255, cv2.THRESH_OTSU+cv2.THRESH_BINARY)

        final_mask = cv2.bitwise_or(final_mask, img_prepped)

    for i in range(5):

        close = 3

        kernel = np.ones((close, close), np.uint8)

        final_mask = cv2.morphologyEx(final_mask, cv2.MORPH_OPEN, kernel)

        final_mask = cv2.morphologyEx(final_mask, cv2.MORPH_CLOSE, kernel)

    return final_mask

def mask_to_polygons(mask, min_area, min_area_holes=10., epsilon=10.):

    """Convert a mask ndarray (binarized image) to Multipolygons"""

    # first, find contours with cv2: it's much faster than shapely

    contours, hierarchy = cv2.findContours(mask,

                                  cv2.RETR_CCOMP,

                                  cv2.CHAIN_APPROX_NONE)

    if not contours:

        return MultiPolygon()

    cnt_children = defaultdict(list)

    child_contours = set()

    assert hierarchy.shape[0] == 1

    # http://docs.opencv.org/3.1.0/d9/d8b/tutorial_py_contours_hierarchy.html

    for idx, (_, _, _, parent_idx) in enumerate(hierarchy[0]):

        if parent_idx != -1:

            child_contours.add(idx)

            cnt_children[parent_idx].append(contours[idx])

    # create actual polygons filtering by area (removes artifacts)

    all_polygons = []

    for idx, cnt in enumerate(contours):

        if idx not in child_contours and cv2.contourArea(cnt) >= min_area:

            assert cnt.shape[1] == 1

            poly = Polygon(

                shell=cnt[:, 0, :],

                holes=[c[:, 0, :] for c in cnt_children.get(idx, [])

                       if cv2.contourArea(c) >= min_area_holes])

            if not poly.is_valid:

                # This is likely becausee the polygon is self-touching or self-crossing.

                # Try and 'correct' the polygon using the zero-length buffer() trick.

                # See https://shapely.readthedocs.io/en/stable/manual.html#object.buffer

                poly = poly.buffer(0)

            all_polygons.append(poly)

    if len(all_polygons) == 0:

        raise Exception("Raw tissue mask consists of 0 polygons")

    # if this raises an issue - try instead unary_union(all_polygons)        

    all_polygons = MultiPolygon(all_polygons)

    return all_polygons

def detect_foreground(contours, hierarchy):

    hierarchy = np.squeeze(hierarchy, axis=(0,))[:, 2:]

    # find foreground contours (parent == -1)

    hierarchy_1 = np.flatnonzero(hierarchy[:, 1] == -1)

    foreground_contours = [contours[cont_idx] for cont_idx in hierarchy_1]

    all_holes = []

    for cont_idx in hierarchy_1:

        all_holes.append(np.flatnonzero(hierarchy[:, 1] == cont_idx))

    hole_contours = []

    for hole_ids in all_holes:

        holes = [contours[idx] for idx in hole_ids]

        hole_contours.append(holes)

    return foreground_contours, hole_contours

def construct_polygon(foreground_contours, hole_contours, min_area):

    polys = []

    for foreground, holes in zip(foreground_contours, hole_contours):

        # We remove all contours that consist of fewer than 3 points, as these won't work with the Polygon constructor.

        if len(foreground) < 3:

            continue

        # remove redundant dimensions from the contour and convert to Shapely Polygon

        poly = Polygon(np.squeeze(foreground))

        # discard all polygons that are considered too small

        if poly.area < min_area:

            continue

        if not poly.is_valid:

            # This is likely becausee the polygon is self-touching or self-crossing.

            # Try and 'correct' the polygon using the zero-length buffer() trick.

            # See https://shapely.readthedocs.io/en/stable/manual.html#object.buffer

            poly = poly.buffer(0)

        # Punch the holes in the polygon

        for hole_contour in holes:

            if len(hole_contour) < 3:

                continue

            hole = Polygon(np.squeeze(hole_contour))

            if not hole.is_valid:

                continue

            # ignore all very small holes

            if hole.area < min_area:

                continue

            poly = poly.difference(hole)

        polys.append(poly)

    if len(polys) == 0:

        raise Exception("Raw tissue mask consists of 0 polygons")

    # If we have multiple polygons, we merge any overlap between them using unary_union().

    # This will result in a Polygon or MultiPolygon with most tissue masks.

    return unary_union(polys)

def generate_tiles(tile_width_pix, tile_height_pix, img_width, img_height, offsets=[(0, 0)]):

    # Generate tiles covering the entire image.

    # Provide an offset (x,y) to create a stride-like overlap effect.

    # Add an additional tile size to the range stop to prevent tiles being cut off at the edges.

    range_stop_width = int(np.ceil(img_width + tile_width_pix))

    range_stop_height = int(np.ceil(img_height + tile_height_pix))

    rects = []

    for xmin, ymin in offsets:

        cols = range(int(np.floor(xmin)), range_stop_width, tile_width_pix)

        rows = range(int(np.floor(ymin)), range_stop_height, tile_height_pix)

        for x in cols:

            for y in rows:

                rect = Polygon(

                    [

                        (x, y),

                        (x + tile_width_pix, y),

                        (x + tile_width_pix, y - tile_height_pix),

                        (x, y - tile_height_pix),

                    ]

                )

                rects.append(rect)

    return rects

def make_tile_QC_fig(tiles, slide, level, line_width_pix=1, extra_tiles=None):

    # Render the tiles on an image derived from the specified zoom level

    img = slide.read_region((0, 0), level, slide.level_dimensions[level])

    downsample = 1 / slide.level_downsamples[level]

    draw = ImageDraw.Draw(img, "RGBA")

    for tile in tiles:

        bbox = tuple(np.array(tile.bounds) * downsample)

        draw.rectangle(bbox, outline="lightgreen", width=line_width_pix)

    # allow to display other tiles, such as excluded or sampled

    if extra_tiles:

        for tile in extra_tiles:

            bbox = tuple(np.array(tile.bounds) * downsample)

            draw.rectangle(bbox, outline="blue", width=line_width_pix + 1)

    return img

def create_tissue_mask(wsi, seg_level, method='otsu'):

    # Determine the best level to determine the segmentation on

    level_dims = wsi.level_dimensions[seg_level]

    img = np.array(wsi.read_region((0, 0), seg_level, level_dims))

    # Get the total surface area of the slide level that was used

    level_area = level_dims[0] * level_dims[1]

    # Minimum surface area of tissue polygons (in pixels)

    # Note that this value should be sensible in the context of the chosen tile size

    min_area = level_area / 500

    if method=='stain_deconv':

        tissue_mask = segment_tissue_deconv_stain(img)

        tissue_mask = mask_to_polygons(tissue_mask, min_area)

    else:

        contours, hierarchy = segment_tissue(img)

        foreground_contours, hole_contours = detect_foreground(contours, hierarchy)

        tissue_mask = construct_polygon(foreground_contours, hole_contours, min_area)

    # Scale the tissue mask polygon to be in the coordinate space of the slide's level 0

    scale_factor = wsi.level_downsamples[seg_level]

    tissue_mask_scaled = scale(

        tissue_mask, xfact=scale_factor, yfact=scale_factor, zfact=1.0, origin=(0, 0)

    )

    return tissue_mask_scaled

def create_tissue_tiles(wsi, tissue_mask_scaled, tile_size_microns, offsets_micron=None):

    print(f"tile size is {tile_size_microns} um")

    # Compute the tile size in pixels from the desired tile size in microns and the image resolution

    assert (

        openslide.PROPERTY_NAME_MPP_X in wsi.properties

    ), "microns per pixel along X-dimension not available"

    assert (

        openslide.PROPERTY_NAME_MPP_Y in wsi.properties

    ), "microns per pixel along Y-dimension not available"

    mpp_x = float(wsi.properties[openslide.PROPERTY_NAME_MPP_X])

    mpp_y = float(wsi.properties[openslide.PROPERTY_NAME_MPP_Y])

    # For larger tiles in micron, NKI scanner outputs mppx slight different than mppy.

    # Force tiles to be squared.

    mpp_scale_factor = min(mpp_x, mpp_y)

    if mpp_x != mpp_y:

        print(

            f"mpp_x of {mpp_x} and mpp_y of {mpp_y} are not the same. Using smallest value: {mpp_scale_factor}"

        )

    tile_size_pix = round(tile_size_microns / mpp_scale_factor)

    # Use the tissue mask bounds as base offsets (+ a margin of a few tiles) to avoid wasting CPU power creating tiles that are never going

    # to be inside the tissue mask.

    tissue_margin_pix = tile_size_pix * 2

    minx, miny, maxx, maxy = tissue_mask_scaled.bounds

    min_offset_x = minx - tissue_margin_pix

    min_offset_y = miny - tissue_margin_pix

    offsets = [(min_offset_x, min_offset_y)]

    if offsets_micron is not None:

        assert (

            len(offsets_micron) > 0

        ), "offsets_micron needs to contain at least one value"

        # Compute the offsets in micron scale

        offset_pix = [round(o / mpp_scale_factor) for o in offsets_micron]

        offsets = [(o + min_offset_x, o + min_offset_y) for o in offset_pix]

    # Generate tiles covering the entire WSI

    all_tiles = generate_tiles(

        tile_size_pix,

        tile_size_pix,

        maxx + tissue_margin_pix,

        maxy + tissue_margin_pix,

        offsets=offsets,

    )

    # Retain only the tiles that sit within the tissue mask polygon

    filtered_tiles = [rect for rect in all_tiles if tissue_mask_scaled.intersects(rect)]

    return filtered_tiles

def tile_is_not_empty(tile, threshold_white=20):

    histogram = tile.histogram()

    # Take the median of each RGB channel. Alpha channel is not of interest.

    # If roughly each chanel median is below a threshold, i.e close to 0 till color value around 250 (white reference) then tile mostly white.

    whiteness_check = [0, 0, 0]

    for channel_id in (0, 1, 2):

        whiteness_check[channel_id] = np.median(

            histogram[256 * channel_id : 256 * (channel_id + 1)][100:200]

        )

    if all(c <= threshold_white for c in whiteness_check):

        # exclude tile

        return False

    # keep tile

    return True

def crop_rect_from_slide(slide, rect):

    minx, miny, maxx, maxy = rect.bounds

    # Note that the y-axis is flipped in the slide: the top of the shapely polygon is y = ymax,

    # but in the slide it is y = 0. Hence: miny instead of maxy.

    top_left_coords = (int(minx), int(miny))

    return slide.read_region(top_left_coords, 0, (int(maxx - minx), int(maxy - miny)))

class BagOfTiles(Dataset):

    def __init__(self, wsi, tiles, resize_to=224):

        self.wsi = wsi

        self.tiles = tiles

        self.roi_transforms = transforms.Compose(

            [

                # As we can't be sure that the input tile dimensions are all consistent, we resize

                # them to a commonly used size before feeding them to the model.

                # Note: assumes a square image.

                transforms.Resize(resize_to),

                # Turn the PIL image into a (C x H x W) float tensor in the range [0.0, 1.0]

                transforms.ToTensor(),

            ]

        )

    def __len__(self):

        return len(self.tiles)

    def __getitem__(self, idx):

        tile = self.tiles[idx]

        img = crop_rect_from_slide(self.wsi, tile)

        # RGB filtering - calling here speeds up computation since it requires crop_rect_from_slide function.

        #is_tile_kept = tile_is_not_empty(img, threshold_white=20)

        is_tile_kept = True

        # Ensure the img is RGB, as expected by the pretrained model.

        # See https://pytorch.org/docs/stable/torchvision/models.html

        img = img.convert("RGB")

        # Ensure we have a square tile in our hands.

        # We can't handle non-squares currently, as this would requiring changes to

        # the aspect ratio when resizing.

        width, height = img.size

        assert width == height, "input image is not a square"

        img = self.roi_transforms(img).unsqueeze(0)

        coord = tile.bounds

        return img, coord, is_tile_kept

def collate_features(batch):

    # Item 2 is the boolean value from tile filtering.

    img = torch.cat([item[0] for item in batch if item[2]], dim=0)

    coords = np.vstack([item[1] for item in batch if item[2]])

    return [img, coords]

def mergedpatch_gen(features, coords, dist_threshold=4, corr_threshold = 0.6):

    # Get patch distance in pixels with rendered segmentation level. Note that each patch is squared and therefore same distance.

    patch_dist = abs(coords[0,2] - coords[0,0]) 

    print(patch_dist)

    # Compute feature similarity (cosine) and nearby pacthes (L2 norm - only need the top left x,y coordinates)

    cosine_matrix = cosine_similarity(features, features)

    coordinate_matrix = euclidean_distances(coords[:,:2], coords[:,:2])

    # NOTE: random selection for the first patch for patch merging might be less biased towards tissue orientation and size. 

    indices_avail = np.arange(features.shape[0])

    np.random.seed(0)  

    np.random.shuffle(indices_avail)

    # Merging together nearby patches and similar within pre-defined threshold. 

    mergedfeatures = []

    indices_used = []

    for ref in indices_avail:

        # This has been merged already

        if ref not in indices_used:

            # Making sure they won't be selected once more

            if indices_used:

                coordinate_matrix[ref,indices_used] = [np.Inf]*len(indices_used)

                cosine_matrix[ref,indices_used] = [0.0]*len(indices_used)

            indices_dist = np.where(coordinate_matrix[ref] < patch_dist*dist_threshold, 1 , 0)

            indices_corr = np.where(cosine_matrix[ref] > corr_threshold, 1 , 0)

            final_indices = indices_dist * indices_corr

            # which includes already the ref patch

            indices_used.extend(list(np.where(final_indices == 1)[0]))

            mergedfeatures.append(tuple((features[final_indices==1,:], coords[final_indices==1,:])))

        else:

            continue

    assert len(indices_used)==features.shape[0], f'Probably issue in contruscting merged features for graph {len(indices_used)}!={features.shape[0]}'

    return mergedfeatures

class HNSW:

    def __init__(self, space):

        self.space = space

    def fit(self, X):

        # See https://nmslib.github.io/nmslib/quickstart.html

        index = nmslib.init(space=self.space, method='hnsw')

        index.addDataPointBatch(X)

        index.createIndex()

        self.index_ = index

        return self

    def query(self, vector, topn):

        indices, dist = self.index_.knnQuery(vector, k=topn)

        return indices, dist

@torch.no_grad()

def extract_features(model, device, wsi, filtered_tiles, workers, out_size, batch_size, n_last_blocks, avgpool_patchtokens, depths):

    # Use multiple workers if running on the GPU, otherwise we'll need all workers for evaluating the model.

    kwargs = (

        {"num_workers": workers, "pin_memory": True} if device.type == "cuda" else {}

    )

    loader = DataLoader(

        dataset=BagOfTiles(wsi, filtered_tiles, resize_to=out_size),

        batch_size=batch_size,

        collate_fn=collate_features,

        **kwargs,

    )

    features_ = []

    coords_ = []

    for batch, coords in loader:

        batch = batch.to(device, non_blocking=True)

        # NOTE: Example using EsVIT. You may want to call your own feature extractor otherwise. 

        features = model.forward_return_n_last_blocks(batch, n_last_blocks, avgpool_patchtokens, depths).cpu().numpy()

        features_.extend(features)

        coords_.extend(coords)

    return np.asarray(features_), np.asarray(coords_)

def extract_save_features(args):

    # Derive the slide ID from its name.

    slide_id, _ = os.path.splitext(os.path.basename(args.input_slide))

    wip_file_path = os.path.join(args.output_dir, slide_id + "_wip.h5")

    output_file_path = os.path.join(args.output_dir, slide_id + "_features.h5")

    os.makedirs(args.output_dir, exist_ok=True)

    # Check if the _features output file already exist. If so, we terminate to avoid

    # overwriting it by accident. This also simplifies resuming bulk batch jobs.

    if os.path.exists(output_file_path):

        raise Exception(f"{output_file_path} already exists")

    # Open the slide for reading.

    wsi = openslide.open_slide(args.input_slide)

    # Decide on which slide level we want to base the segmentation.

    seg_level = wsi.get_best_level_for_downsample(64)

    # Run the segmentation and  tiling procedure.

    start_time = time.time()

    tissue_mask_scaled = create_tissue_mask(wsi, seg_level, method=args.method)

    filtered_tiles = create_tissue_tiles(wsi, tissue_mask_scaled, args.tile_size)

    # Build a figure for quality control purposes, to check if the tiles are where we expect them.

    qc_img = make_tile_QC_fig(filtered_tiles, wsi, seg_level, 2)

    qc_img_target_width = 1920

    qc_img = qc_img.resize((qc_img_target_width, int(qc_img.height / (qc_img.width / qc_img_target_width))))

    qc_img_file_path = os.path.join(args.output_dir, f"{slide_id}_features_QC.png")

    qc_img.save(qc_img_file_path)

    print(f"Finished creating {len(filtered_tiles)} tissue tiles in {time.time() - start_time}s")

    # Save QC figure.

    qc_img_file_path = os.path.join(

        args.output_dir, f"{slide_id}_N{len(mergedpatches)}mergedpatches_distThreshold{args.dist_threshold}_corrThreshold{args.corr_threshold}.png"

    )

    # Extract the rectangles, and compute the feature vectors. Example using EsVIT. 

    device = torch.device("cuda") 

    model, _, depths = load_encoder_esVIT(args, device)

    features, coords = extract_features(

        model,

        device,

        wsi,

        filtered_tiles,

        args.workers,

        args.out_size,

        args.batch_size,

        n_last_blocks = args.n_last_blocks, 

        avgpool_patchtokens = args.avgpool_patchtokens,

        depths = depths,

    )

    print(f'Number of features N={len(features)}')

    # Merging nearby patches with similar semantic. 

    mergedpatches = mergedpatch_gen(features, coords, dist_threshold=args.dist_threshold, corr_threshold=args.corr_threshold)

    print(f'Merging step => N={len(mergedpatches)}')

    # Saving features.

    torch.save(mergedpatches, wip_file_path)

    # Rename the file containing the patches to ensure we can easily

    # distinguish incomplete bags of patches (due to e.g. errors) from complete ones in case a job fails.

    os.rename(wip_file_path, output_file_path)

    print('Done.')

if __name__ == '__main__':

    parser = argparse.ArgumentParser('Preprocessing script esvit', parents=[get_args_parser()])

    args = parser.parse_args()

    assert os.path.isfile(args.checkpoint), f'{args.checkpoint} does not exist'

    assert torch.cuda.is_available(), 'Need cuda for this job'

    assert os.path.isfile(args.input_slide), f'{args.input_slide} does not exist'

    extract_save_features(args)