import os
from matplotlib import pyplot as plt
import matplotlib
import time
import numpy as np
from matplotlib.colors import LinearSegmentedColormap
from nilearn.plotting import plot_anat
from matplotlib import cm
from skimage.transform import resize
import io


def plot_batch(batch, seg: bool = False, slice: int = 32, batch_size: int=4):

    def unnorm(data, epsilon=1e-8):
        non_zero = data[data > 0.0]
        mean = non_zero.mean()
        std = non_zero.std() + epsilon
        out = data * std + mean
        out[data == 0] = 0
        return out

    plt.figure(figsize=(10, 3.5))

    for i, volume in enumerate(batch):
        plt.subplot(1, batch_size + 1, i + 1)

        img = volume[:, slice, :].T if seg else volume[0, :, slice, :].T

        npimg = img.cpu().detach().numpy()
        img = npimg if seg else unnorm(npimg)
        plt.imshow(img, cmap="gray")
        plt.axis("off")


    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    plt.close()
    return buf



def plot_3_view(modal: str, vol: np.ndarray, s: int=100, discrete: bool=False,
                color_map: str="gray", save: bool=True):

    views = [vol[s, :, :], vol[:, s, :], vol[:, :, s]]
    fig = plt.figure(figsize=(10, 3.5))

    for position in range(1, len(views) + 1):
        plt.subplot(1, len(views), position)
        plt.imshow(views[position - 1].T, cmap=color_map)
        plt.axis("off")
        if discrete:
            plt.clim(0, 4)

    if discrete:
        plt.colorbar(ticks=range(5))
    else:
        plt.colorbar()

    if save:
        fig.savefig(f'plot_{modal}_{time.time()}.png')
    else:
        plt.show()


def plot_3_view_uncertainty(modal: str, vol: np.ndarray, s: int=100, color_map: str="gray", save: bool=True):

    views = [vol[s, :, :], vol[:, s, :], vol[:, :, s]]
    fig = plt.figure(figsize=(10, 3.5))

    for position in range(1, len(views) + 1):
        plt.subplot(1, len(views), position)
        plt.imshow(views[position - 1].T, cmap=color_map)
        plt.axis("off")
        plt.clim(0, 100)

    plt.colorbar()

    if save:
        fig.savefig(f'plot_unc_{modal}_{time.time()}.png')
    else:
        plt.show()


def plot_axis_overlayed(modalities: dict, segmentation_mask: str, subject: int, axis: str = 'x', save: bool=False):
    """Save or show figure of provided axis"""
    fig, axes = plt.subplots(len(modalities), 1)

    for i, (modality_name, modality_path) in enumerate(modalities.items()):
        display = plot_anat(modality_path, draw_cross=False, display_mode=axis, axes=axes[i], figure=fig, title=modality_name)
        display.add_overlay(segmentation_mask)

    if save:
        fig.savefig(f'results/patient_{subject}.png')
    else:
        matplotlib.use('TkAgg')
        plt.show()


def plot_brain_batch_per_patient(patient_ids, data, save=True):
    for patient in patient_ids:
        patient = data[patient.item()]
        patient_modalities = list(map(lambda x: os.path.join(patient.data_path, patient.patch_name, x), [patient.flair, patient.t2, patient.t1, patient.t1ce]))
        patient_modalities = {"flair": patient_modalities[0],"t2": patient_modalities[1],"t1": patient_modalities[2],"t1ce": patient_modalities[3] }
        patient_seg = os.path.join(patient.data_path, patient.patch_name, patient.seg)
        plot_axis_overlayed(patient_modalities, patient_seg, patient.patch_name, axis='x', save=save)


def plot_batch_slice(images, gt, slice = 10, save=True):
    """Plot, for a given batch, different types of visualizations.
    If paths: plot overlayed axis plot
    If paths=None: plot slice of volume
    """
    for element_index in range(0, len(images)):
        for i, mod_id in enumerate(images):
            patient_mod = images[i][element_index]

            plot_3_view(f"batch_element_{i}", patient_mod, slice, save=save)

        patient_seg = gt[element_index]
        plot_3_view('seg', patient_seg, slice, save=save)


def plot_batch_cubes(patient_ids, batch_volumes, batch_gt, patches=1, img_size=30):
    for batch_pos, patient in enumerate(patient_ids[:patches]):
        patient = patient.item()
        modality = batch_volumes[batch_pos][0]
        gt =  batch_gt[batch_pos][0]
        resized_modality = resize(modality,(img_size, img_size, img_size), mode='constant')
        resized_gt = resize(gt, (img_size, img_size, img_size), mode='constant')

        fig = plot_cube(resized_modality, img_size)
        fig.savefig(f'results/3Dplot_{patient}_{batch_pos}.png')

        fig_seg = plot_cube(resized_gt, img_size)
        fig_seg.savefig(f'results/3Dplot_{patient}_gt_{batch_pos}.png')


def plot_cube(cube, dim, gt=False, angle=320):
    def normalize(arr):
        arr_min = np.min(arr)
        return (arr - arr_min) / (np.max(arr) - arr_min)

    def explode(data):
        shape_arr = np.array(data.shape)
        size = shape_arr[:3] * 2 - 1
        exploded = np.zeros(np.concatenate([size, shape_arr[3:]]), dtype=data.dtype)
        exploded[::2, ::2, ::2] = data
        return exploded

    def expand_coordinates(indices):
        x, y, z = indices
        x[1::2, :, :] += 1
        y[:, 1::2, :] += 1
        z[:, :, 1::2] += 1
        return x, y, z

    if gt:
        colors = [(1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 1)]
        my_cmap = LinearSegmentedColormap.from_list('my_cmap', colors, 4)
        plt.register_cmap(cmap=my_cmap)
        cmap = plt.get_cmap('my_cmap')
        cube = np.around(cube)
    else:
        cmap = cm.viridis

    cube = normalize(cube)
    facecolors = cmap(cube)
    facecolors[:, :, :, -1] = cube
    facecolors = explode(facecolors)

    filled = facecolors[:, :, :, -1] != 0
    x, y, z = expand_coordinates(np.indices(np.array(filled.shape) + 1))

    fig = plt.figure(figsize=(30 / 2.54, 30 / 2.54))
    ax = fig.gca(projection='3d')
    ax.view_init(30, angle)
    ax.set_xlim(right=dim * 2)
    ax.set_ylim(top=dim * 2)
    ax.set_zlim(top=dim * 2)

    ax.voxels(x, y, z, filled, facecolors=facecolors, shade=False)
    return fig