
"""
Module that deals with color ramps

This product includes color specifications and designs developed by
Cynthia Brewer (http://colorbrewer.org/).

This file uses the minify function, available from
https://github.com/getify/JSON.minify

"""
# This file is part of 'TuiView' - a simple Raster viewer
# Copyright (C) 2012  Sam Gillingham
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import os
import numpy
import json

from . import viewererrors

# environment variable we use to specify file
# to load extra ramps from
EXTRA_RAMP_VAR = 'TUIVIEW_EXTRA_RAMP'
HAVE_LOADED_EXTRA_RAMPS = False

# init our dictionary of data
# longest list of colors for each name
# from http://colorbrewer.org/
# generated by colorbrewer2py.py

RAMP = {}
RAMP['Blues'] = {'author' : 'Cynthia Brewer', 'comments' : 'Colours from www.colorbrewer.org by Cynthia A. Brewer, Geography, Pennsylvania State University.', 'type' : 'Sequential'}
RAMP['Blues']['description'] = {}
RAMP['Blues']['description']['red'] = '247 222 198 158 107 66 33 8 8'
RAMP['Blues']['description']['green'] = '251 235 219 202 174 146 113 81 48'
RAMP['Blues']['description']['blue'] = '255 247 239 225 214 198 181 156 107'
RAMP['Greys'] = {'author' : 'Cynthia Brewer', 'comments' : 'Colours from www.colorbrewer.org by Cynthia A. Brewer, Geography, Pennsylvania State University.', 'type' : 'Sequential'}
RAMP['Greys']['description'] = {}
RAMP['Greys']['description']['red'] = '255 240 217 189 150 115 82 37 0'
RAMP['Greys']['description']['green'] = '255 240 217 189 150 115 82 37 0'
RAMP['Greys']['description']['blue'] = '255 240 217 189 150 115 82 37 0'
RAMP['RdBu'] = {'author' : 'Cynthia Brewer', 'comments' : 'Colours from www.colorbrewer.org by Cynthia A. Brewer, Geography, Pennsylvania State University.', 'type' : 'Diverging'}
RAMP['RdBu']['description'] = {}
RAMP['RdBu']['description']['red'] = '103 178 214 244 253 247 209 146 67 33 5'
RAMP['RdBu']['description']['green'] = '0 24 96 165 219 247 229 197 147 102 48'
RAMP['RdBu']['description']['blue'] = '31 43 77 130 199 247 240 222 195 172 97'
RAMP['RdYlBu'] = {'author' : 'Cynthia Brewer', 'comments' : 'Colours from www.colorbrewer.org by Cynthia A. Brewer, Geography, Pennsylvania State University.', 'type' : 'Diverging'}
RAMP['RdYlBu']['description'] = {}
RAMP['RdYlBu']['description']['red'] = '165 215 244 253 254 255 224 171 116 69 49'
RAMP['RdYlBu']['description']['green'] = '0 48 109 174 224 255 243 217 173 117 54'
RAMP['RdYlBu']['description']['blue'] = '38 39 67 97 144 191 248 233 209 180 149'
RAMP['RdYlGn'] = {'author' : 'Cynthia Brewer', 'comments' : 'Colours from www.colorbrewer.org by Cynthia A. Brewer, Geography, Pennsylvania State University.', 'type' : 'Diverging'}
RAMP['RdYlGn']['description'] = {}
RAMP['RdYlGn']['description']['red'] = '165 215 244 253 254 255 217 166 102 26 0'
RAMP['RdYlGn']['description']['green'] = '0 48 109 174 224 255 239 217 189 152 104'
RAMP['RdYlGn']['description']['blue'] = '38 39 67 97 139 191 139 106 99 80 55'
RAMP['Reds'] = {'author' : 'Cynthia Brewer', 'comments' : 'Colours from www.colorbrewer.org by Cynthia A. Brewer, Geography, Pennsylvania State University.', 'type' : 'Sequential'}
RAMP['Reds']['description'] = {}
RAMP['Reds']['description']['red'] = '255 254 252 252 251 239 203 165 103'
RAMP['Reds']['description']['green'] = '245 224 187 146 106 59 24 15 0'
RAMP['Reds']['description']['blue'] = '240 210 161 114 74 44 29 21 13'
RAMP['Spectral'] = {'author' : 'Cynthia Brewer', 'comments' : 'Colours from www.colorbrewer.org by Cynthia A. Brewer, Geography, Pennsylvania State University.', 'type' : 'Diverging'}
RAMP['Spectral']['description'] = {}
RAMP['Spectral']['description']['red'] = '158 213 244 253 254 255 230 171 102 50 94'
RAMP['Spectral']['description']['green'] = '1 62 109 174 224 255 245 221 194 136 79'
RAMP['Spectral']['description']['blue'] = '66 79 67 97 139 191 152 164 165 189 162'
RAMP['YlGnBu'] = {'author' : 'Cynthia Brewer', 'comments' : 'Colours from www.colorbrewer.org by Cynthia A. Brewer, Geography, Pennsylvania State University.', 'type' : 'Sequential'}
RAMP['YlGnBu']['description'] = {}
RAMP['YlGnBu']['description']['red'] = '255 237 199 127 65 29 34 37 8'
RAMP['YlGnBu']['description']['green'] = '255 248 233 205 182 145 94 52 29'
RAMP['YlGnBu']['description']['blue'] = '217 177 180 187 196 192 168 148 88'
RAMP['YlOrRd'] = {'author' : 'Cynthia Brewer', 'comments' : 'Colours from www.colorbrewer.org by Cynthia A. Brewer, Geography, Pennsylvania State University.', 'type' : 'Sequential'}
RAMP['YlOrRd']['description'] = {}
RAMP['YlOrRd']['description']['red'] = '255 255 254 254 253 252 227 189 128'
RAMP['YlOrRd']['description']['green'] = '255 237 217 178 141 78 26 0 0'
RAMP['YlOrRd']['description']['blue'] = '204 160 118 76 60 42 28 38 38'
# Viridis palettes
RAMP['viridis'] = {'author' : ' Stefan van der Walt and Nathaniel Smith', 'comments' : '', 'type' : 'Diverging'}
RAMP['viridis']['description'] = {}
RAMP['viridis']['description']['red'] = '68 72 67 56 45 37 30 43 81 133 194 253'
RAMP['viridis']['description']['green'] = '1 33 62 89 112 133 155 176 197 213 223 231'
RAMP['viridis']['description']['blue'] = '84 115 133 140 142 142 138 127 106 74 35 37'
RAMP['magma'] = {'author' : ' Stefan van der Walt and Nathaniel Smith', 'comments' : '', 'type' : 'Diverging'}
RAMP['magma']['description'] = {}
RAMP['magma']['description']['red'] = '0 18 51 90 125 163 200 233 249 254 254 252'
RAMP['magma']['description']['green'] = '0 13 16 22 36 48 62 85 124 168 211 253'
RAMP['magma']['description']['blue'] = '4 50 104 126 130 126 115 98 93 115 149 191'
RAMP['plasma'] = {'author' : ' Stefan van der Walt and Nathaniel Smith', 'comments' : '', 'type' : 'Diverging'}
RAMP['plasma']['description'] = {}
RAMP['plasma']['description']['red'] = '13 62 99 135 166 192 213 231 245 253 252 240'
RAMP['plasma']['description']['green'] = '8 4 0 7 32 58 84 111 140 173 210 249'
RAMP['plasma']['description']['blue'] = '135 156 167 166 152 131 110 90 70 50 37 33'
RAMP['inferno'] = {'author' : ' Stefan van der Walt and Nathaniel Smith', 'comments' : '', 'type' : 'Diverging'}
RAMP['inferno']['description'] = {}
RAMP['inferno']['description']['red'] = '0 20 58 96 133 169 203 230 247 252 245 252'
RAMP['inferno']['description']['green'] = '0 11 9 19 33 46 65 93 131 173 219 255'
RAMP['inferno']['description']['blue'] = '4 53 99 110 107 94 73 47 17 18 75 164'
RAMP['cividis'] = {'author' : ' Stefan van der Walt and Nathaniel Smith', 'comments' : '', 'type' : 'Diverging'}
RAMP['cividis']['description'] = {}
RAMP['cividis']['description']['red'] = '0 0 42 72 94 114 135 158 182 208 234 255'
RAMP['cividis']['description']['green'] = '32 48 64 82 98 115 132 150 169 190 211 234'
RAMP['cividis']['description']['blue'] = '77 111 108 107 110 116 121 119 113 103 87 70'

def HEXtoRGB(colorstring):
    """
    Converts #RRGGBB to an (R, G, B) tuple,
    from http://code.activestate.com/recipes/266466-html-colors-tofrom-rgb-tuples/
    """
    colorstring = colorstring.strip()
    if colorstring[0] == '#': colorstring = colorstring[1:]
    if len(colorstring) != 6:
        msg = "Colour #%s is not in #RRGGBB format" % colorstring
        raise viewererrors.ColorRampException(msg)
    r, g, b = colorstring[:2], colorstring[2:4], colorstring[4:]
    r, g, b = [int(n, 16) for n in (r, g, b)]
    return (r, g, b)

def getRampsFromFile(fname):
    """
    Read extra color ramps into our global RAMPS dictionary
    from specified file
    """
    from minify_json import json_minify
    # Read palette file
    palettesFobj = open(fname, "r")
    # Minify file contents
    palsMinified = json_minify(palettesFobj.read())
    # Loads palettes in a dict
    pals = json.loads(palsMinified)

    # For each palette that's been detected
    for pal in pals:
        # Check name is present and unique
        if "name" in pal:
            cur_name = pal["name"]
        else:
            # Quit - invalid colour scheme structure
            msg = 'Invalid colour ramp structure'
            raise viewererrors.ColorRampException(msg)

        if cur_name in RAMP.keys():
            # Quit - invalid colour scheme name
            msg = 'Duplicated colour ramp name'
            raise viewererrors.ColorRampException(msg)

        # Check red, green and blue fields are present
        if "description" in pal:
            # Parsing description of the colour ramp
            descriptionKeys = pal["description"].keys()
            # Is it RGB coded?
            isRgb = set(descriptionKeys).issubset(set(['red', 'green', 'blue']))
            # Is it hex coded?
            isHex = descriptionKeys == ['hex']
            # If RGB is given it has priority
            if isRgb:
                reds = pal["description"]["red"]
                greens = pal["description"]["green"]
                blues = pal["description"]["blue"]
            else:
                # Else look for HEX
                if isHex:
                    # Convert HEX to RGB
                    colrgb = numpy.array([HEXtoRGB(col) for col in pal["description"]["hex"].split()])
                    # Add RGB values to palette decription
                    reds = " ".join(map(str,colrgb[:,0]))
                    greens = " ".join(map(str,colrgb[:,1]))
                    blues = " ".join(map(str,colrgb[:,2]))
                else:
                    # Quit - invalid colour scheme
                    msg = 'Invalid colour ramp structure'
                    raise viewererrors.ColorRampException(msg)
        else:
            # Quit - invalid colour scheme
            msg = 'Invalid colour ramp structure'
            raise viewererrors.ColorRampException(msg)

        # Other fields optional
        # not sure what we'll be doing with this yet
        if "author" in pal:
            cur_author = pal["author"]
        else:
            cur_author = ''

        if "comments" in pal:
            cur_comments = pal["comments"]
        else:
            cur_comments = ''

        if "type" in pal:
            cur_type = pal["type"]
        else:
            cur_type = ''

        # Assembling dictionary entry
        RAMP[cur_name] = {'author' : cur_author, 'comments' : cur_comments, 'type' : cur_type, 'description' : {}}
        # Add decsription fields
        RAMP[cur_name]["description"]["red"] = reds
        RAMP[cur_name]["description"]["green"] = greens
        RAMP[cur_name]["description"]["blue"] = blues

def loadExtraRamps():
    """
    Try to load extra colour ramps
    from file spcified by EXTRA_RAMP_VAR env variable.
    Checks that we haven't done it before so we don't have
    duplicates
    """
    global HAVE_LOADED_EXTRA_RAMPS
    if HAVE_LOADED_EXTRA_RAMPS:
        return

    palettesFile = os.getenv(EXTRA_RAMP_VAR)

    # check that it is set
    if palettesFile is not None:
        try:
            # process palettes and append RAMP dict
            getRampsFromFile(palettesFile)
            HAVE_LOADED_EXTRA_RAMPS = True
        except IOError:
            # wasn't able to open file, set error
            msg = 'Unable to open %s' % palettesFile
            raise viewererrors.ColorRampException(msg)

def getRampsForDisplay():
    """
    Returns a list of (name, displayname) tuples for
    populating Combo boxes. Sorted by type
    """
    # sort by stype
    rampDict = {}
    for key in RAMP:
        rampType = RAMP[key]['type']
        if rampType in rampDict:
            rampDict[rampType].append(key)
        else:
            rampDict[rampType] = [key]

    display = []
    for rampType in rampDict:
        for name in rampDict[rampType]:
            title = "%s (%s)" % (name, rampType)
            display.append((name, title))
    return display

def getLUTForRamp(code, name, lutsize):
    """
    Returns an LUT for the specified color and name
    """
    # get the data as string
    colstr = RAMP[name]['description'][code]
    # turn it into a list of floats
    # numpy.interp() needs floats
    colList = [float(x) for x in colstr.split()]
    # the x-values of the observations
    # evenly spaced 0-255 with len(colList) obs
    xobs = numpy.linspace(0, 255, len(colList))
    # create an array from our list
    yobs = numpy.array(colList)
    # values to interpolate at 0-255
    # same size as the lut
    xinterp = numpy.linspace(0, 255, lutsize)
    # do the interp
    yinterp = numpy.interp(xinterp, xobs, yobs)
    # return as 8 bit int
    return yinterp.astype(numpy.uint8)
