import numpy
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plotter
import sys

print 'Python version:     %s' %sys.version
print 'matplotlib version: %s' %matplotlib.__version__
print 'numpy version:      %s' %numpy.__version__


# define a custom color map ...

bit_depth = 8
levels    = 2 ** bit_depth

reds = [
     0,   0,   0,   0,   0,   0,   0,   1,   1,   1,   2,   2,   3,   4,   4,   5,   5,   6,
     7,   8,   9,  10,  11,  12,  13,  14,  15,  17,  18,  19,  20,  22,  23,  25,  26,  28,
    29,  31,  33,  34,  36,  38,  40,  41,  43,  45,  47,  49,  50,  53,  55,  57,  59,  61,
    64,  66,  68,  70,  72,  74,  76,  79,  82,  84,  87,  89,  91,  94,  96,  99, 101, 104,
   106, 108, 110, 113, 116, 118, 121, 124, 127, 128, 131, 134, 137, 138, 141, 144, 147, 149,
   152, 155, 156, 160, 161, 164, 168, 169, 172, 174, 177, 179, 182, 184, 188, 189, 191, 195,
   196, 198, 202, 203, 205, 209, 211, 212, 214, 216, 220, 222, 224, 225, 227, 229, 231, 233,
   235, 237, 239, 241, 243, 245, 247, 249, 251, 253, 255, 255, 255, 255, 255, 255, 255, 255,
   255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
   255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
   255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
   255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
   255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
   255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
   255, 255, 255, 255]
greens = [
     0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
     0,   0,   0,   0,   0,   0,   0,   0,   1,   1,   1,   1,   1,   1,   1,   1,   1,   1,
     2,   2,   2,   2,   2,   2,   2,   2,   3,   3,   3,   3,   3,   3,   4,   4,   4,   4,
     4,   5,   5,   5,   5,   5,   5,   6,   6,   6,   6,   7,   7,   7,   7,   8,   8,   8,
     9,   9,   9,  10,  10,  10,  11,  11,  11,  11,  12,  12,  13,  13,  13,  14,  14,  15,
    15,  15,  16,  16,  17,  17,  18,  18,  18,  19,  19,  20,  20,  21,  22,  22,  23,  23,
    23,  24,  25,  25,  26,  27,  27,  28,  29,  29,  30,  31,  31,  32,  33,  33,  34,  35,
    36,  36,  37,  38,  39,  40,  40,  41,  42,  43,  44,  44,  45,  46,  47,  48,  49,  50,
    50,  51,  52,  53,  54,  56,  57,  58,  59,  60,  61,  62,  63,  64,  66,  67,  68,  69,
    70,  71,  72,  74,  75,  76,  77,  79,  81,  82,  83,  84,  85,  87,  89,  90,  91,  93,
    95,  96,  97,  99, 100, 102, 104, 105, 106, 109, 110, 112, 113, 116, 117, 118, 121, 122,
   124, 125, 128, 129, 131, 134, 135, 137, 138, 141, 143, 144, 147, 149, 150, 153, 155, 156,
   160, 161, 163, 166, 168, 169, 172, 174, 176, 179, 181, 182, 186, 188, 189, 193, 195, 196,
   200, 202, 205, 207, 209, 212, 214, 216, 220, 222, 224, 227, 229, 233, 235, 237, 241, 243,
   245, 249, 251, 255]
blues = [
     0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
     0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
     0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
     0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
     0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
     0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   1,   1,   1,   1,   1,   1,
     1,   1,   1,   1,   2,   2,   2,   2,   2,   2,   3,   3,   3,   3,   4,   4,   4,   4,
     5,   5,   5,   5,   6,   6,   6,   7,   7,   7,   8,   8,   9,   9,  10,  10,  11,  11,
    11,  12,  13,  13,  14,  15,  15,  16,  17,  17,  18,  19,  19,  20,  21,  22,  23,  23,
    25,  25,  26,  27,  28,  29,  30,  31,  33,  33,  35,  36,  37,  38,  40,  41,  42,  44,
    45,  46,  48,  50,  50,  52,  54,  56,  57,  59,  61,  63,  64,  66,  68,  70,  72,  74,
    75,  77,  80,  82,  84,  87,  89,  91,  93,  95,  97, 100, 102, 105, 108, 110, 113, 116,
   118, 121, 124, 127, 129, 132, 135, 138, 141, 144, 147, 150, 153, 156, 160, 163, 166, 169,
   172, 177, 181, 184, 188, 191, 195, 198, 202, 205, 209, 212, 218, 222, 225, 229, 233, 237,
   241, 245, 249, 255]

assert len(reds) == len(greens) == len(blues) == levels   

palette = [reds,greens,blues]
palette = numpy.array(palette)
assert palette.shape == (3,levels)
assert (palette[0] == reds).all()
assert (palette[1] == greens).all()
assert (palette[2] == blues).all()

colormap = matplotlib.colors.ListedColormap(palette,'custom-orange')


# use custom palette to colorize an image ...

image = numpy.random.randint(0,levels, (512,580))

rows_cnt, columns_cnt = image.shape
shape = rows_cnt, columns_cnt
# x,y coordinate definitions for the image map ...
rows    = numpy.arange(rows_cnt)
columns = numpy.arange(columns_cnt)
XI, YI  = numpy.meshgrid(rows,columns)

plotter.clf() # clear the figure -- presumed (but not verified) necessary since pyplot is stateful
figure  = plotter.figure(figsize=(6.4,6), dpi=80, facecolor='white', edgecolor='black') # figsize = (width,height) in (fractional) inches

left, bottom, width, height = 0.1, 0.07, 0.9, 0.84 # values in 0-1 relative figure coordinates -- allow space on left for colorbar
this  = figure.add_axes((left, bottom, width, height))
this.pcolormesh(XI, YI, image, cmap=colormap)
this.imshow(image, interpolation='bilinear', cmap=colormap, origin='upper', extent=[0,rows_cnt,0,columns_cnt], vmin=0, vmax=255)

plotter.savefig('%s.png' %colormap.name, facecolor='black', edgecolor='black')
plotter.close()
