# New matplotlib colormaps by Nathaniel J. Smith, Stefan van der Walt,
# and (in the case of viridis) Eric Firing.
#
# This file and the colormaps in it are released under the CC0 license /
# public domain dedication. We would appreciate credit if you use or
# redistribute these colormaps, but do not impose any legal restrictions.
#
# To the extent possible under law, the persons who associated CC0 with
# mpl-colormaps have waived all copyright and related or neighboring rights
# to mpl-colormaps.
#
# You should have received a copy of the CC0 legalcode along with this
# work.  If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.

__all__ = ['apply_colormap', 'apply_colormap_single']

cdef float _magma_data[16][3]
_magma_data[0][:] = [0.001462, 0.000466, 0.013866]
_magma_data[1][:] = [0.039608, 0.031090, 0.133515]
_magma_data[2][:] = [0.113094, 0.065492, 0.276784]
_magma_data[3][:] = [0.211718, 0.061992, 0.418647]
_magma_data[4][:] = [0.316654, 0.071690, 0.485380]
_magma_data[5][:] = [0.414709, 0.110431, 0.504662]
_magma_data[6][:] = [0.512831, 0.148179, 0.507648]
_magma_data[7][:] = [0.613617, 0.181811, 0.498536]
_magma_data[8][:] = [0.716387, 0.214982, 0.475290]
_magma_data[9][:] = [0.816914, 0.255895, 0.436461]
_magma_data[10][:] = [0.904281, 0.319610, 0.388137]
_magma_data[11][:] = [0.960949, 0.418323, 0.359630]
_magma_data[12][:] = [0.986700, 0.535582, 0.382210]
_magma_data[13][:] = [0.996096, 0.653659, 0.446213]
_magma_data[14][:] = [0.996898, 0.769591, 0.534892]
_magma_data[15][:] = [0.992440, 0.884330, 0.640099]


cdef float _inferno_data[16][3]
_inferno_data[0][:] = [0.001462, 0.000466, 0.013866]
_inferno_data[1][:] = [0.042253, 0.028139, 0.141141]
_inferno_data[2][:] = [0.129285, 0.047293, 0.290788]
_inferno_data[3][:] = [0.238273, 0.036621, 0.396353]
_inferno_data[4][:] = [0.341500, 0.062325, 0.429425]
_inferno_data[5][:] = [0.441207, 0.099338, 0.431594]
_inferno_data[6][:] = [0.540920, 0.134729, 0.415123]
_inferno_data[7][:] = [0.640135, 0.171438, 0.381065]
_inferno_data[8][:] = [0.735683, 0.215906, 0.330245]
_inferno_data[9][:] = [0.822386, 0.275197, 0.266085]
_inferno_data[10][:] = [0.894305, 0.353399, 0.193584]
_inferno_data[11][:] = [0.946965, 0.449191, 0.115272]
_inferno_data[12][:] = [0.978422, 0.557937, 0.034931]
_inferno_data[13][:] = [0.987874, 0.675267, 0.065257]
_inferno_data[14][:] = [0.974638, 0.797692, 0.206332]
_inferno_data[15][:] = [0.947594, 0.917399, 0.410665]

cdef float _plasma_data[16][3]
_plasma_data[0][:] = [0.050383, 0.029803, 0.527975]
_plasma_data[1][:] = [0.193374, 0.018354, 0.590330]
_plasma_data[2][:] = [0.299855, 0.009561, 0.631624]
_plasma_data[3][:] = [0.399411, 0.000859, 0.656133]
_plasma_data[4][:] = [0.494877, 0.011990, 0.657865]
_plasma_data[5][:] = [0.584391, 0.068579, 0.632812]
_plasma_data[6][:] = [0.665129, 0.138566, 0.585582]
_plasma_data[7][:] = [0.736019, 0.209439, 0.527908]
_plasma_data[8][:] = [0.798216, 0.280197, 0.469538]
_plasma_data[9][:] = [0.853319, 0.351553, 0.413734]
_plasma_data[10][:] = [0.901807, 0.425087, 0.359688]
_plasma_data[11][:] = [0.942598, 0.502639, 0.305816]
_plasma_data[12][:] = [0.973416, 0.585761, 0.251540]
_plasma_data[13][:] = [0.991365, 0.675355, 0.198453]
_plasma_data[14][:] = [0.993033, 0.771720, 0.154808]
_plasma_data[15][:] = [0.974443, 0.874622, 0.144061]

cdef float _viridis_data[16][3]
_viridis_data[0][:] = [0.267004, 0.004874, 0.329415]
_viridis_data[1][:] = [0.282327, 0.094955, 0.417331]
_viridis_data[2][:] = [0.278826, 0.175490, 0.483397]
_viridis_data[3][:] = [0.258965, 0.251537, 0.524736]
_viridis_data[4][:] = [0.229739, 0.322361, 0.545706]
_viridis_data[5][:] = [0.199430, 0.387607, 0.554642]
_viridis_data[6][:] = [0.172719, 0.448791, 0.557885]
_viridis_data[7][:] = [0.149039, 0.508051, 0.557250]
_viridis_data[8][:] = [0.127568, 0.566949, 0.550556]
_viridis_data[9][:] = [0.120638, 0.625828, 0.533488]
_viridis_data[10][:] = [0.157851, 0.683765, 0.501686]
_viridis_data[11][:] = [0.246070, 0.738910, 0.452024]
_viridis_data[12][:] = [0.369214, 0.788888, 0.382914]
_viridis_data[13][:] = [0.515992, 0.831158, 0.294279]
_viridis_data[14][:] = [0.678489, 0.863742, 0.189503]
_viridis_data[15][:] = [0.845561, 0.887322, 0.099702]

import numpy as np
cimport cython

cmaps = {}
for (name, data) in (
        ('magma', _magma_data),
        ('inferno', _inferno_data),
        ('plasma', _plasma_data),
        ('viridis', _viridis_data)
):

    cmaps[name] = np.array(data, dtype=np.float32)*255

magma = cmaps['magma']
inferno = cmaps['inferno']
plasma = cmaps['plasma']
viridis = cmaps['viridis']


@cython.boundscheck(False)
@cython.wraparound(False)
cdef map_many(int[:, :] idxs, float[:, :, :] dst, float[:,:] cmap):
    cdef int height = idxs.shape[0]
    cdef int width = idxs.shape[1]
    for i in range(height):
        for j in range(width):
            dst[i][j][0] = cmap[idxs[i][j]][0]
            dst[i][j][1] = cmap[idxs[i][j]][1]
            dst[i][j][2] = cmap[idxs[i][j]][2]


def apply_colormap(input_: np.ndarray, colormap: str, min_: float=None, max_: float=None):

    if max_ is None:
        max_ = np.max(input_)
    if min_ is None:
        min_ = np.min(input_)
    range_ = max_ - min_
    if range_ == 0:
        idxs = np.zeros((input_.shape[0], input_.shape[1]), dtype=np.int32)
    else:
        normalized = 15*(input_ - min_) / range_
        idxs = normalized.astype(np.int32)
        np.clip(idxs, 0, 15, out=idxs)
    colors = np.zeros(shape=(input_.shape[0], input_.shape[1], 3), dtype=np.float32)
    map_many(idxs, colors, cmaps[colormap])
    return colors

def apply_colormap_single(input_: np.ndarray, colormap: str, min_: float=None, max_: float=None):
    range_ = max_ - min_
    if range_ == 0:
        idx = 0
    else:
        idx = int(15*(input_ - min_) / range_)
    idx = np.clip(idx, 0, 15)
    color = tuple(cmaps[colormap][idx])
    return color