import numpy as np
from random import randrange


SET_CHANNELS = [0,1,2]

def sliding_window(image, stride=10, window_size=(20,20)):
    """Extract patches according to a sliding window.
    Args:
        image (numpy array): The image to be processed.
        stride (int, optional): The sliding window stride (defaults to 10px).
        window_size(int, int, optional): The patch size (defaults to (20,20)).
    Returns:
        list: list of patches with window_size dimensions
    """
    patches = []
    # slide a window across the image
    for x in range(0, image.shape[1], stride):
        for y in range(0, image.shape[2], stride):
            new_patch = image[:, x:x + window_size[0], y:y + window_size[1]]
            if new_patch.shape[1:] == window_size:
                patches.append(new_patch)
    return patches


def transform(patch, flip=False, mirror=False):
    """Perform data augmentation on a patch.
    Args:
        patch (numpy array): The patch to be processed.
        flip (bool, optional): Up/down symetry.
        mirror (bool, optional): left/right symetry.
    Returns:
        array list: list of augmented patches
    """
    transformed_patches = [patch]
    if flip:
        transformed_patches.append(np.flipud(patch))
    if mirror:
        transformed_patches.append(np.fliplr(patch))
    return transformed_patches


def transform_mask(mask, img_name='normal'):
    if 'SCL' in img_name:
         mask[np.where(mask == 1) or np.where(mask == 2) or np.where(mask == 4) or np.where(mask == 5) or np.where(mask == 6) or np.where(mask == 7) or np.where(mask == 11)] = 0  # Land Water (Snow)
         mask[np.where(mask == 3)] = 1  # Shadow
         mask[np.where(mask == 8) or np.where(mask == 9) or np.where(mask == 10)] = 2  # Cloud
         mask[np.where(mask == 0)] = 3  # No observation
    else:
        if 'fmask' in img_name:
            mask[np.where(mask == 0) or np.where(mask == 1) or np.where(mask == 3)] = 0  # Land Water (Snow)
            mask[np.where(mask == 2)] = 1  # Shadow
            mask[np.where(mask == 4)] = 2  # Cloud
            mask[np.where(mask == 255)] = 3  # No observation
        else:
            mask[np.where(mask == 255)] = 3  # No observation
    return mask




def transform_img(img, NUM_CHANNELS=6, img_name='normal'):
    if 'czdata' in img_name or 'sensor.05' in img_name:
        print 'gisat'
        return transform_img_gisat(img, NUM_CHANNELS)
    else:
        return transform_img_normal(img, NUM_CHANNELS)


def transform_img_normal(img, NUM_CHANNELS=6):
    return img[SET_CHANNELS, :, :]

def transform_img_gisat(img, NUM_CHANNELS=6):
    img = img.astype('float32')
    img /= 1e4
    img *= 255
    img = img.astype('uint8')
    return img[SET_CHANNELS, :, :]


def slice_imgs(img, mask, img_width=128, img_height=128, num_slices=1000):
    H = img.shape[1]
    W = img.shape[2]

    #num_slices=int(H*W / IMAGE_WIDTH * IMAGE_HEIGHT)
    slices = []
    HALF_HEIGHT = img_height / 2
    HALF_WIDTH = img_width / 2
    for i in range(0, num_slices):
        hrand = randrange(HALF_HEIGHT + 1, H - HALF_HEIGHT - 1)
        wrand = randrange(HALF_WIDTH + 1, W - HALF_WIDTH - 1)

        islice = img[:, hrand-HALF_HEIGHT:hrand+HALF_HEIGHT, wrand - HALF_WIDTH:wrand+HALF_WIDTH]

        slices.append((islice, mask[hrand, wrand]))
    return slices
