from PIL import Image
from typing import Tuple
import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage
import torch
from torchvision.transforms import functional as tvf

from pathlib import Path

def sliced_mean(x, slice_size):
    cs_y = np.cumsum(x, axis=0)
    cs_y = np.concatenate((np.zeros((1, cs_y.shape[1]), dtype=cs_y.dtype), cs_y), axis=0)
    slices_y = (cs_y[slice_size:] - cs_y[:-slice_size])/slice_size
    cs_xy = np.cumsum(slices_y, axis=1)
    cs_xy = np.concatenate((np.zeros((cs_xy.shape[0], 1), dtype=cs_xy.dtype), cs_xy), axis=1)
    slices_xy = (cs_xy[:,slice_size:] - cs_xy[:,:-slice_size])/slice_size
    return slices_xy

def sliced_var(x, slice_size):
    x = x.astype('float64')
    return sliced_mean(x**2, slice_size) - sliced_mean(x, slice_size)**2

def calculate_local_variance(img, var_window):
    """return local variance map with the same size as input image"""
    var = sliced_var(img, var_window)

    left_pad = var_window // 2 -1
    right_pad = var_window -1 - left_pad
    var_padded = np.pad(
        var,
        pad_width=(
            (left_pad,right_pad),
            (left_pad,right_pad)
        ))
    return var_padded

def get_crop_batch(img: np.ndarray, mask: np.ndarray, crop_size=96, crop_scales=np.geomspace(0.5, 2, 7), samples_per_scale=32, use_variance_threshold=False) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Generate a batch of cropped images from an input image and corresponding mask, at various scales and rotations.

    Parameters
    ----------
    img : np.ndarray
        The input image from which crops are generated.
    mask : np.ndarray
        The binary mask indicating the region of interest in the image.
    crop_size : int, optional
        The size of the square crop.
    crop_scales : np.ndarray, optional
        An array of scale factors to apply to the crop size.
    samples_per_scale : int, optional
        Number of samples to generate per scale factor.
    use_variance_threshold : bool, optional
        Flag to use variance thresholding for selecting crop locations.

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        A tuple containing the tensor of crops, their rotation angles, and scale factors.
    """
    
    # pad
    pad_size = int(np.ceil(0.5*crop_size*max(crop_scales)*(np.sqrt(2)-1)))
    img_padded = np.pad(img, pad_size)
    mask_padded = np.pad(mask, pad_size)

    # distance map
    distance_map_padded = ndimage.distance_transform_edt(mask_padded)
    # TODO: adjust scales and samples_per_scale
    
    if use_variance_threshold:
        variance_window = min(crop_size//2, min(img.shape))
        variance_map_padded = np.pad(calculate_local_variance(img, variance_window), pad_size)
        variance_median = np.ma.median(np.ma.masked_where(distance_map_padded<0.5*variance_window, variance_map_padded))
        variance_mask = variance_map_padded >= variance_median
    else:
        variance_mask = np.ones_like(mask_padded)
    
    # initilize output
    crops_granum = []
    angles_granum = []
    scales_granum = []
    # loop over scales
    for scale in crop_scales: 
        half_crop_size_scaled = int(np.floor(scale*0.5*crop_size)) # half of crop size after scaling
        crop_pad = int(np.ceil((np.sqrt(2) - 1)*half_crop_size_scaled)) # pad added in order to allow rotation
        half_crop_size_external = half_crop_size_scaled + crop_pad # size of "external crop" which will be rotated

        possible_indices = np.stack(np.where(variance_mask & (distance_map_padded >= 2*half_crop_size_scaled)), axis=1)
        if len(possible_indices) == 0:
            continue
        chosen_indices = np.random.choice(np.arange(len(possible_indices)), min(len(possible_indices), samples_per_scale), replace=False)

        crops = [
            img_padded[y-half_crop_size_external:y+half_crop_size_external, x-half_crop_size_external:x+half_crop_size_external] for y, x in possible_indices[chosen_indices]
        ]

        # rotate
        rotation_angles = np.random.rand(len(crops))*180 - 90
        crops = [
            ndimage.rotate(crop, angle, reshape=False)[crop_pad:-crop_pad,crop_pad:-crop_pad] for crop, angle in zip(crops, rotation_angles)
        ]
        # add to output
        crops_granum.append(tvf.resize(torch.tensor(np.array(crops)), (crop_size,crop_size),antialias=True)) # resize crops to crop_size
        angles_granum.extend(rotation_angles.tolist())
        scales_granum.extend([scale]*len(crops))

    if len(angles_granum) == 0:
        return [], [], []
    
    crops_granum = torch.concat(crops_granum)
    angles_granum = torch.tensor(angles_granum, dtype=torch.float)
    scales_granum = torch.tensor(scales_granum, dtype=torch.float)
    
    return crops_granum, angles_granum, scales_granum

def get_crop_batch_from_path(img_path, mask_path=None, use_variance_threshold=False):
    """
    Load an image and its mask from file paths and generate a batch of cropped images.

    Parameters
    ----------
    img_path : str
        Path to the input image.
    mask_path : str, optional
        Path to the binary mask image. If None, assumes mask path by replacing image extension with '.npy'.
    use_variance_threshold : bool, optional
        Flag to use variance thresholding for selecting crop locations.

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        A tuple containing the tensor of crops, their rotation angles, and scale factors, obtained from the specified image path.
    """
    if mask_path is None:
        mask_path = str(Path(img_path).with_suffix('.npy'))
    mask = np.load(mask_path)
    img = np.array(Image.open(img_path))[:,:,0]
    
    return get_crop_batch(img, mask, use_variance_threshold=use_variance_threshold)