import random
import itertools
import numpy as np


def generate_general_indexs():
    indexes = [[0, 2],
              [1, 3],
              [2, 0],
              [3, 1],
              [4, 6],
              [5, 7],
              [6, 4],
              [7, 5]]
    return indexes

def generate_general_permute_keys():
    rotate_y = 0
    transpose = 0
    keys = [((rotate_y, 0), 0, 0, 0, transpose),
            ((rotate_y, 1), 0, 0, 0, transpose),
            ((rotate_y, 0), 0, 1, 1, transpose),
            ((rotate_y, 1), 0, 1, 1, transpose),
            ((rotate_y, 0), 1, 0, 0, transpose),
            ((rotate_y, 1), 1, 0, 0, transpose),
            ((rotate_y, 0), 1, 1, 1, transpose),
            ((rotate_y, 1), 1, 1, 1, transpose)]
    return keys

def random_permute_key():
    """
    Generate and randomly selects a permute key
    """
    return random.choice(list(generate_general_permute_keys()))

def generate_permutation_keys():
    """
    This function returns a set of "keys" that represent the 48 unique rotations &
    reflections of a 3D matrix.

    Each item of the set is a tuple:
    ((rotate_y, rotate_z), flip_x, flip_y, flip_z, transpose)

    As an example, ((0, 1), 0, 1, 0, 1) represents a permutation in which the data is
    rotated 90 degrees around the z-axis, then reversed on the y-axis, and then
    transposed.

    48 unique rotations & reflections:
    https://en.wikipedia.org/wiki/Octahedral_symmetry#The_isometries_of_the_cube
    """
    return set(itertools.product(
        itertools.combinations_with_replacement(range(2), 2), range(2), range(2), range(2), range(2)))

def random_permutation_key():
    """
    Generates and randomly selects a permutation key. See the documentation for the
    "generate_permutation_keys" function.
    """
    return random.choice(list(generate_permutation_keys()))

def augment_data(data):
    padded_data = -1000 * np.ones((256,256,256))
    padded_data[104:152, :, :] = data

    data = padded_data.transpose(1, 2, 0)

    return data

def permute_data(data, key):
    """
    Permutes the given data according to the specification of the given key. Input data
    must be of shape (n_modalities, x, y, z).

    Input key is a tuple: (rotate_y, rotate_z), flip_x, flip_y, flip_z, transpose)

    As an example, ((0, 1), 0, 1, 0, 1) represents a permutation in which the data is
    rotated 90 degrees around the z-axis, then reversed on the y-axis, and then
    transposed.
    """
    data = np.copy(data)
    #print(data.shape)
    (rotate_y, rotate_z), flip_x, flip_y, flip_z, transpose = key

    if rotate_y != 0:
        data = np.rot90(data, rotate_y, axes=(1, 3))
    if rotate_z != 0:
        data = np.rot90(data, rotate_z, axes=(2, 3))
    if flip_x:
        data = data[:, ::-1]
    if flip_y:
        data = data[:, :, ::-1]
    if flip_z:
        data = data[:, :, :, ::-1]
    if transpose:
        for i in range(data.shape[0]):
            data[i] = data[i].T
    return data

def random_permutation_data(data):
    key = random_permute_key()
    return permute_data(data, key)