from Segmentation.plotting.voxels import plot_volume, plot_to_image
import tensorflow as tf
import numpy as np

# TODO: Replace all tf.slice operations with more Pythonic expressions
# TODO: colour maps should be consistent with the ones for 2D

colour_maps = {
    1: [tf.constant([1, 1, 1], dtype=tf.float32), tf.constant([[[[255, 255, 0]]]], dtype=tf.float32)],  # background / black
    2: [tf.constant([2, 2, 2], dtype=tf.float32), tf.constant([[[[0, 255, 255]]]], dtype=tf.float32)],
    3: [tf.constant([3, 3, 3], dtype=tf.float32), tf.constant([[[[255, 0, 255]]]], dtype=tf.float32)],
    4: [tf.constant([4, 4, 4], dtype=tf.float32), tf.constant([[[[255, 255, 255]]]], dtype=tf.float32)],
    5: [tf.constant([5, 5, 5], dtype=tf.float32), tf.constant([[[[120, 120, 120]]]], dtype=tf.float32)],
    6: [tf.constant([6, 6, 6], dtype=tf.float32), tf.constant([[[[255, 165, 0]]]], dtype=tf.float32)],
}

def replace_vector(img, search, replace):
    condition = tf.equal(img, search)
    condition = tf.reduce_all(condition, axis=-1)
    condition = tf.stack((condition,) * img.shape[-1], axis=-1)
    replace_tiled = tf.tile(replace, img.shape[:-1])
    replace_tiled = tf.reshape(replace_tiled, img.shape)
    return tf.where(condition, replace_tiled, img)

def get_mid_slice(x, y, pred, multi_class):
    mid = tf.cast(tf.divide(tf.shape(y)[1], 2), tf.int32)
    start = [0, mid, 0, 0, 0]
    size = [1, 1, -1, -1, -1]
    x_slice = tf.slice(x, start, size)
    y_slice = tf.slice(y, start, size)
    pred_slice = tf.slice(pred, start, size)
    if multi_class:
        x_slice = tf.squeeze(x_slice, axis=-1)
        x_slice = tf.stack((x_slice,) * 3, axis=-1)
        y_slice = tf.argmax(y_slice, axis=-1)
        y_slice = tf.stack((y_slice,) * 3, axis=-1)
        y_slice = tf.cast(y_slice, tf.float32)
        pred_slice = tf.argmax(pred_slice, axis=-1)
        pred_slice = tf.stack((pred_slice,) * 3, axis=-1)
        pred_slice = tf.cast(pred_slice, tf.float32)
        for c in colour_maps:
            y_slice = replace_vector(y_slice, colour_maps[c][0], colour_maps[c][1])
            pred_slice = replace_vector(pred_slice, colour_maps[c][0], colour_maps[c][1])
    else:
        pred_slice = tf.math.round(pred_slice)

    img_pad = tf.ones((pred_slice.shape[0], pred_slice.shape[1], pred_slice.shape[2], 3, pred_slice.shape[4]))
    img = tf.concat((x_slice, img_pad, y_slice, img_pad, pred_slice), axis=-2)

    return tf.reshape(img, (img.shape[1:]))

def get_mid_vol(y, pred, multi_class, rad=12, check_empty=False):
    y_shape = tf.shape(y)
    slice_start = [0, (y_shape[1] // 2) - rad, (y_shape[2] // 2) - rad, (y_shape[3] // 2) - rad, 0]
    slice_size = [1, rad * 2, rad * 2, rad * 2, -1]
    y_subvol = tf.slice(y, slice_start, slice_size)

    if (check_empty) and (tf.math.reduce_sum(y_subvol) < 25):
        return None

    if multi_class:
        y_subvol = tf.argmax(y_subvol, axis=-1)
        y_subvol = tf.cast(y_subvol, tf.float32)
    else:
        y_subvol = tf.reshape(y_subvol, (y_subvol.shape[1:4]))
    y_subvol = tf.stack((y_subvol,) * 3, axis=-1)
    pred_subvol = tf.slice(pred, slice_start, slice_size)
    if multi_class:
        pred_subvol = tf.argmax(pred_subvol, axis=-1)
        pred_subvol = tf.cast(pred_subvol, tf.float32)
    else:
        pred_subvol = tf.math.round(pred_subvol)  # new
        pred_subvol = tf.reshape(pred_subvol, (pred_subvol.shape[1:4]))
    pred_subvol = tf.stack((pred_subvol,) * 3, axis=-1)
    if multi_class:
        for c in colour_maps:
            y_subvol = replace_vector(y_subvol, colour_maps[c][0], tf.divide(colour_maps[c][1], 255))
            pred_subvol = replace_vector(pred_subvol, colour_maps[c][0], tf.divide(colour_maps[c][1], 255))
        y_subvol = tf.squeeze(y_subvol, axis=0)
        pred_subvol = tf.squeeze(pred_subvol, axis=0)
    fig = plot_volume(y_subvol, show=False)
    y_img = plot_to_image(fig)
    del y_subvol
    fig = plot_volume(pred_subvol, show=False)
    del pred_subvol
    pred_img = plot_to_image(fig)

    img = tf.concat((y_img, pred_img), axis=-2)
    return img

def plot_through_slices(batch_idx, x_crop, y_crop, mean_pred, writer, multi_class=False):
    imgs = []
    slice_size = [1, 1, -1, -1, -1]
    for i in range(160):
        slice_start = [batch_idx, i, 0, 0, 0]
        x_slice = tf.slice(x_crop, slice_start, slice_size)
        y_slice = tf.slice(y_crop, slice_start, slice_size)
        m_slice = tf.slice(mean_pred, slice_start, slice_size)
        if multi_class:
            y_slice = tf.argmax(y_slice, axis=-1)
            y_slice = tf.cast(y_slice, tf.float32)
            y_slice = tf.stack((y_slice,) * 3, axis=-1)
            m_slice = tf.argmax(m_slice, axis=-1)
            m_slice = tf.cast(m_slice, tf.float32)
            m_slice = tf.stack((m_slice,) * 3, axis=-1)
            if multi_class:
                for c in colour_maps:
                    y_slice = replace_vector(y_slice, colour_maps[c][0], tf.divide(colour_maps[c][1], 255))
                    m_slice = replace_vector(m_slice, colour_maps[c][0], tf.divide(colour_maps[c][1], 255))
        else:
            m_slice = tf.math.round(m_slice)
        if multi_class:
            x_slice = tf.squeeze(x_slice, axis=-1)
            x_slice = tf.stack((x_slice,) * 3, axis=-1)        
        if not multi_class:
            x_min = tf.reduce_min(x_slice)
            x_max = tf.reduce_max(x_slice)
            x_slice = (x_slice - x_min) / (x_max - x_min)
        img = tf.concat((x_slice, y_slice, m_slice), axis=-2)
        img = tf.reshape(img, (img.shape[1:]))
        with writer.as_default():
            tf.summary.image(f"Whole Validation - All Slices", img, step=i)
        img = tf.reshape(img, (img.shape[1:]))
        if not multi_class:
            img = tf.reshape(img, (img.shape[:-1]))
            img = img * 255
        imgs.append(img.numpy().astype(np.uint8))
    return imgs