# Copyright 2019 Population Health Sciences and Image Analysis, German Center for Neurodegenerative Diseases(DZNE)
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.gridspec as gridspec
import matplotlib.cm as cm
from matplotlib.colors import LinearSegmentedColormap
import itertools

def get_colors(inp, colormap, vmin=None, vmax=None):
    """generate the normalize rgb values for matplolib
"""
    norm = plt.Normalize(vmin, vmax)
    return colormap(norm(inp))


def multiview_plotting(data,labels,control_point, savepath,classes=5,alpha=0.5,nbviews=3,plot_labels=True,plot_control_point=True):
    """Plot data and label in different views
    Args:
        data: Original 3D volume
        labels: Original labels for the 3d Volume
        control_point: select the center point where the different views are going to be created
        savepath:path where the image is going to be safe
        classes: number of classes in the labeles
        alpha: transparency of the labels on the original data
        nbviews: 1 only axial view,2 axial and frontal, 3 the three views
        plot_labels: True plot labels, False only plot data


    Returns:
        An images with the diffent views and the corresponding label
"""
    # Create the colormap for the labels
    dz=np.arange(classes)
    colors = get_colors(dz, plt.cm.jet)
    #replace first color for black
    colors[0, 0:3] = [0, 0, 0]
    my_cm=LinearSegmentedColormap.from_list('mylist',colors,classes)
    plt.ioff()
    if plot_labels:
        grid_size = [2, nbviews]
    else:
        grid_size = [1, nbviews]
    #fig = plt.figure(dpi=600)
    fig = plt.gcf()


    outer_grid = gridspec.GridSpec(grid_size[0], grid_size[1], wspace=0.05, hspace=0.0005)
    index = 0
    for i in range(nbviews):
        if i == 0 :
            ax = plt.subplot(outer_grid[i])
            ax.imshow(data[control_point[0], :, :], cmap=cm.gray)
            if plot_control_point:
                ax.scatter(y=control_point[1], x=control_point[2], c='r', s=2)
            ax.set_xticklabels([])
            ax.set_yticklabels([])
            if plot_labels:
                ax = plt.subplot(outer_grid[i+nbviews])
                ax.imshow(data[control_point[0], :, :], cmap=cm.gray)
                ax.imshow(labels[control_point[0], :, :],vmin=0,vmax=classes, cmap=my_cm,alpha=alpha)
                if plot_control_point:
                    ax.scatter(y=control_point[1], x=control_point[2], c='r', s=2)
                ax.set_xticklabels([])
                ax.set_yticklabels([])

        elif i == 1:
            ax = plt.subplot(outer_grid[i])
            ax.imshow(data[:, control_point[1], :], cmap=cm.gray,aspect=(data.shape[1]/data.shape[0]))
            if plot_control_point:
                ax.scatter(y=control_point[0], x=control_point[2], c='r', s=2)
            ax.set_xticklabels([])
            ax.set_yticklabels([])
            if plot_labels:
                ax = plt.subplot(outer_grid[i+nbviews])
                ax.imshow(data[:, control_point[1], :], cmap=cm.gray,aspect=(data.shape[1]/data.shape[0]))
                ax.imshow(labels[:, control_point[1], :],vmin=0,vmax=classes, cmap=my_cm,alpha=alpha,aspect=(data.shape[1]/data.shape[0]))
                if plot_control_point:
                    ax.scatter(y=control_point[0], x=control_point[2], c='r', s=2)
                ax.set_xticklabels([])
                ax.set_yticklabels([])

        elif i == 2:
            img=np.zeros((data.shape[0],data.shape[2]))
            diff_spacing=int((data.shape[2]-data.shape[1])/2)
            img[:,diff_spacing:data.shape[2]-diff_spacing]=data[:, :, control_point[2]]

            ax = plt.subplot(outer_grid[i])
            ax.imshow(img, cmap=cm.gray,aspect=(data.shape[1]/data.shape[0]))
            if plot_control_point:
                ax.scatter(y=control_point[0], x=control_point[2], c='r', s=2)
            ax.set_xticklabels([])
            ax.set_yticklabels([])
            if  plot_labels:
                img_label = np.zeros((data.shape[0], data.shape[2]))
                img_label[:, diff_spacing:data.shape[2] - diff_spacing] = labels[:, :, control_point[2]]
                ax = plt.subplot(outer_grid[i + nbviews])
                ax.imshow(img, cmap=cm.gray, aspect=(data.shape[1] / data.shape[0]))
                ax.imshow(img_label,vmin=0,vmax=classes, cmap=my_cm,alpha=alpha, aspect=(data.shape[1] / data.shape[0]))
                if plot_control_point:
                    ax.scatter(y=control_point[0], x=control_point[2], c='r', s=2)
                ax.set_xticklabels([])
                ax.set_yticklabels([])

    #fig.subplots_adjust(wspace=0.001, hspace=0.001)

    plt.subplots_adjust(0,0,1,1,0,0)
    for ax in fig.axes:
        ax.axis('off')
        ax.margins(0,0)
        ax.xaxis.set_major_locator(plt.NullLocator())
        ax.yaxis.set_major_locator(plt.NullLocator())

    plt.savefig(savepath, transparent=True,bbbox_inches='tight',pad_inches=0,dpi=300)
    plt.close(fig)