import os
import numpy as np
import SimpleITK as sitk
import matplotlib.pyplot as plt
# set a dark background
plt.rcParams.update({
    "lines.color": "white",
    "patch.edgecolor": "white",
    "text.color": "black",
    "axes.facecolor": "white",
    "axes.edgecolor": "lightgray",
    "axes.labelcolor": "white",
    "xtick.color": "white",
    "ytick.color": "white",
    "grid.color": "lightgray",
    "figure.facecolor": "black",
    "figure.edgecolor": "black",
    "savefig.facecolor": "black",
    "savefig.edgecolor": "black"})

"""
Standalone reproduction of Nilearn's 'cold_hot' colormap.

Original reference:
Nilearn developers. Nilearn: Statistical learning for neuroimaging in Python.
URL: https://nilearn.github.io/
"""
from matplotlib.colors import LinearSegmentedColormap
# Segment data extracted from the actual Nilearn 'cold_hot'
_cold_hot_segments = {
    'red': [
        (0.0, 1.0, 1.0),
        (0.126984, 0.0, 0.0),
        (0.5, 0.0, 0.0),
        (0.5, 0.0416, 0.0416),
        (0.682540, 1.0, 1.0),
        (1.0, 1.0, 1.0),
    ],
    'green': [
        (0.0, 1.0, 1.0),
        (0.126984, 1.0, 1.0),
        (0.317461, 0.0, 0.0),
        (0.5, 0.0, 0.0),
        (0.5, 0.0, 0.0),
        (0.682540, 0.0, 0.0),
        (0.873016, 1.0, 1.0),
        (1.0, 1.0, 1.0),
    ],
    'blue': [
        (0.0, 1.0, 1.0),
        (0.317461, 1.0, 1.0),
        (0.5, 0.0416, 0.0416),
        (0.5, 0.0, 0.0),
        (0.873016, 0.0, 0.0),
        (1.0, 1.0, 1.0),
    ]
}
# Create the LinearSegmentedColormap
cold_hot = LinearSegmentedColormap('cold_hot', _cold_hot_segments, N=256)


def otsu_scaling(input_image):
    from skimage.filters import threshold_multiotsu
    
    # the input can be either a nifti file or an SITK image
    if isinstance(input_image, sitk.Image):
        img = input_image
    elif os.path.isfile(input_image):
        img = sitk.ReadImage(input_image)
    else:
        raise ValueError(
            "input_image must be a SimpleITK.Image instance or a path to an existing image file"
        )
    array = sitk.GetArrayFromImage(img)

    thresholds = threshold_multiotsu(array.astype(float).flatten(), classes=5, nbins=100)
    voxel_subset=array[array>thresholds[1]] # clip off the background using otsu thresholds
    # select a maximal value which encompasses 90% of the voxels
    voxel_subset.sort()
    vmax=voxel_subset[int(len(voxel_subset)*0.9)]

    scaled = array/vmax
    scaled_img=sitk.GetImageFromArray(scaled, isVector=False)
    scaled_img.CopyInformation(img)
    return scaled_img


def plot_3d(axes,sitk_img,fig,vmin=0,vmax=1,cmap='gray', alpha=1, cbar=False, threshold=None, planes=('sagittal', 'coronal', 'horizontal'), num_slices=4, slice_spacing=0.1):
    physical_dimensions = (np.array(sitk_img.GetSpacing())*np.array(sitk_img.GetSize()))[::-1] # invert because the array is inverted indices
    array=sitk.GetArrayFromImage(sitk_img)

    array[array==0]=None # set 0 values to be empty

    if not threshold is None:
        array[np.abs(array)<threshold]=None

    slice_0 = (1.0-((num_slices-1)*slice_spacing))/2
    slice_fractions=[slice_0]
    for i in range(1,num_slices):
        slice_fractions.append(slice_0+(i*slice_spacing))

    cbar_list = []
    
    ax_number=0
    if 'sagittal' in planes:
        ax=axes[ax_number]
        ax_number+=1
        empty_slice = np.array([np.nan]).repeat(array.shape[0])[:,np.newaxis]
        slices=empty_slice
        for s in slice_fractions:
            slice=array[::-1,:,int(array.shape[2]*s)]
            slices=np.concatenate((slices,slice,empty_slice),axis=1)
        pos = ax.imshow(slices, extent=[0,physical_dimensions[1]*num_slices,0,physical_dimensions[0]], vmin=vmin, vmax=vmax,cmap=cmap, alpha=alpha, interpolation='none')
        ax.axis('off')
        if cbar:
            cbar_list.append(fig.colorbar(pos, ax=ax))

    if 'coronal' in planes:
        ax=axes[ax_number]
        ax_number+=1
        empty_slice = np.array([np.nan]).repeat(array.shape[0])[:,np.newaxis]
        slices=empty_slice
        for s in slice_fractions:
            slice=array[::-1,int(array.shape[1]*s),:]
            slices=np.concatenate((slices,slice,empty_slice),axis=1)
        pos = ax.imshow(slices, extent=[0,physical_dimensions[2]*num_slices,0,physical_dimensions[0]], vmin=vmin, vmax=vmax,cmap=cmap, alpha=alpha, interpolation='none')
        ax.axis('off')
        if cbar:
            cbar_list.append(fig.colorbar(pos, ax=ax))

    if 'horizontal' in planes:
        ax=axes[ax_number]
        ax_number+=1
        empty_slice = np.array([np.nan]).repeat(array.shape[1])[:,np.newaxis]
        slices=empty_slice
        for s in slice_fractions:
            slice=array[int(array.shape[0]*s),::-1,:]
            slices=np.concatenate((slices,slice,empty_slice),axis=1)
        pos = ax.imshow(slices, extent=[0,physical_dimensions[2]*num_slices,0,physical_dimensions[1]], vmin=vmin, vmax=vmax,cmap=cmap, alpha=alpha, interpolation='none')
        ax.axis('off')
        if cbar:
            cbar_list.append(fig.colorbar(pos, ax=ax))
    return cbar_list
