from __future__ import division

import logging


import matplotlib as mpl
import SimpleITK as sitk
import numpy as np


def convert_discrete_colormap(data, cm_name='custom', color_names=None):
    '''Generates a matplotlib continuous colormap on [0, 1] from a discrete 
    colormap at N evenly spaced points.

    Parameters
    ----------
    data : list of list
        Sublists are [r, g, b]. 
    
    Returns
    -------
    matplotlib.colors.LinearSegmentedColormap
        Gamma is 1. Output space is 3 X [0, 1]
    
    '''

    if color_names is None:
        color_names = ['red', 'green', 'blue']

    data = np.array(data)
    npoints = data.shape[0]

    domain = np.linspace(0, 1.0, npoints)
    color_arrays = {}

    for col, name in enumerate(color_names):
        color_array = np.zeros((npoints, 3))
        
        color_array[:, 0] = domain
        color_array[:, 1] = minmax_norm(data[:, col])
        color_array[:, 2] = color_array[:, 1]

        color_arrays[name] = color_array

    return mpl.colors.LinearSegmentedColormap(cm_name, color_arrays, npoints, gamma=1.0)



def minmax_norm(data):

    rng = np.amax(data) - np.amin(data)
    if rng == 0:
        return data 

    return (data - np.amin(data)) / rng



def sitk_safe_ln(data, minimum=10**-10):

    logging.info('thresholding below at {0}'.format(minimum))
    minimum = float(minimum)
    data = sitk.Threshold(data, minimum, np.inf, minimum)

    logging.info('taking natural log')
    return sitk.Log(data)


def normalize_intensity(data, in_min, in_max, out_min=0.0, out_max=0.0):

    logging.info('setting input range: [{0:2.3f}, {1:2.3f}]'.format(in_min, in_max))
    data = sitk.ShiftScale(data, -in_min, 1.0 / (in_max - in_min))
    data = sitk.Threshold(data, 0.0, np.inf, 0.0)
    data = sitk.Threshold(data, 0.0, 1.0, 1.0)

    logging.info('setting output range: [{0:2.3f}, {1:2.3f}]'.format(out_min, out_max))
    data = sitk.ShiftScale(data, 0.0, out_max - out_min) # want to scale first
    return sitk.ShiftScale(data, out_min, 1) # then shift


def blend(image_stack, weight_stack):
    '''

    Parameters
    ----------
    image_stack :: list of np.ndarray
        The images to be blended. Shapes cannot differ
    weight_stack :: list of np.ndarray
        The weight of each image at each pixel. Will be normalized.

    '''

    image_stack = np.array(image_stack)
    weight_stack = np.array(weight_stack)

    weight_stack = weight_stack - np.amin(weight_stack, axis=0) / (np.amax(weight_stack, axis=0) - np.amin(weight_stack, axis=0))
    weight_stack[np.isnan(weight_stack)] = 0.5
    weight_stack[np.isinf(weight_stack)] = 0.5

    return np.multiply(image_stack, weight_stack).sum(axis=0)
