import numpy as np


class NeuralVolume:
    def __init__(self, x_max, y_max, z_max, voxel_size, x_min=0, y_min=0, z_min=0):
        self.axes = tuple([np.arange(ax_min, ax_max + voxel_size, voxel_size) for ax_min, ax_max in
                           zip([x_min, y_min, z_min], [x_max, y_max, z_max])])
        self.voxel_size = voxel_size
        self.mesh = np.meshgrid(*self.axes)

        self.mesh_shape = self.mesh[0].shape
        self.number_of_voxels = np.prod(self.mesh_shape)

    def get_coordinates(self):
        X, Y, Z = self.mesh
        number_of_elements = self.number_of_voxels
        return np.vstack(
            (X.reshape((number_of_elements,)), Y.reshape((number_of_elements,)), Z.reshape((number_of_elements,))))


def get_distance_matrix(positions):
    number_of_positions = positions.shape[1]

    repeated_coordinates = np.repeat(np.expand_dims(positions, axis=2), number_of_positions,
                                     axis=2)
    return np.linalg.norm(repeated_coordinates - np.transpose(repeated_coordinates, axes=(0, 2, 1)),
                          axis=0)


def get_distance_from_to(pos_1, pos_2):
    if len(pos_1.shape) == 1:
        pos_1 = np.expand_dims(pos_1, axis=1)
    no_pos_1 = pos_1.shape[1]
    no_pos_2 = pos_2.shape[1]
    rep_pos_1 = np.repeat(np.expand_dims(pos_1, axis=2), no_pos_2,
                          axis=2)
    rep_pos_2 = np.repeat(np.expand_dims(pos_2, axis=2), no_pos_1,
                          axis=2)
    return np.linalg.norm(rep_pos_1 - np.transpose(rep_pos_2, axes=(0, 2, 1)), axis=0)

def get_distance_matrix(pos):
    return get_distance_from_to(pos, pos)

def get_second_order_distance_cube(pos):
    number_of_positions = pos.shape[1]
    distance_matrix = get_distance_matrix(pos)

    rep_distance_matrix = np.repeat(np.expand_dims(distance_matrix, axis=2), number_of_positions, axis=2)
    second_order_distances = rep_distance_matrix+np.transpose(rep_distance_matrix, axes=(2,1,0))
    return second_order_distances

def mask_invalid_entries(distance_cube):
    # mask all entries (source,target_1,target_2) where either source =target_1, target_1=target_2 or source=target_2
    cube_shape=distance_cube.shape
    helper_mask = np.eye(distance_cube.shape[0])
    rep_mask = np.repeat(np.expand_dims(helper_mask, axis=2), distance_cube.shape[0], axis=2)
    mask = rep_mask+np.transpose(rep_mask, axes=(1,2,0))+np.transpose(rep_mask, axes=(2,0,1))
    mask = mask>0
    return distance_cube[~mask].reshape((cube_shape[0], cube_shape[0]-1,-1))
