from matplotlib import pylab
import sys
from presto.ppgplot import *
import numpy as Num





class Palette:
    # Set the color palette
    def setpalette(self, palette):
        """
        setpalette(self, palette):
            Set the color palette for imag-style routines
        """
        if (palette == 'rainbow'):
            self.l = Num.array([0.0, 0.015, 0.225, 0.4, 0.59,
                                0.6, 0.775, 0.955, 0.965, 1.0])
            self.r = Num.array([1.0, 1.0, 1.0, 0.0, 0.0,
                                0.0, 0.0, 0.947, 1.0, 1.0])
            self.g = Num.array([0.0, 0.0, 1.0, 1.0, 1.0,
                                0.946, 0.0, 0.8, 0.844, 1.0])
            self.b = Num.array([0.0, 0.0, 0.0, 0.0, 0.95,
                                1.0, 1.0, 1.0, 1.0, 1.0])
        elif (palette == 'antirainbow'):
            self.l = Num.array([0.0, 0.035, 0.045, 0.225, 0.4,
                                0.41, 0.6, 0.775, 0.985, 1.0])
            self.r = Num.array([1.0, 1.0, 0.947, 0.0, 0.0,
                                0.0, 0.0, 1.0, 1.0, 1.0])
            self.g = Num.array([1.0, 0.844, 0.8, 0.0, 0.946,
                                1.0, 1.0, 1.0, 0.0, 0.0])
            self.b = Num.array([1.0, 1.0, 1.0, 1.0, 1.0,
                                0.95, 0.0, 0.0, 0.0, 0.0])
        elif (palette == 'astro'):
            self.l = Num.array([0.0, 0.167, 0.333, 0.5,
                                0.667, 0.833, 1.0])
            self.r = Num.array([0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0])
            self.g = Num.array([0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0])
            self.b = Num.array([0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0])
        elif (palette == 'hue'):
            self.l = Num.array([0.0, 0.167, 0.333, 0.5,
                                0.667, 0.833, 1.0])
            self.r = Num.array([1.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0])
            self.g = Num.array([0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0])
            self.b = Num.array([0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0])
        elif (palette == 'heat'):
            self.l = Num.array([0.0, 0.48, 0.7, 0.75, 1.0])
            self.r = Num.array([0.0, 1.0, 1.0, 1.0, 1.0])
            self.g = Num.array([0.0, 0.0, 0.423, 0.519, 1.0])
            self.b = Num.array([0.0, 0.0, 0.0, 0.0, 1.0])
        elif (palette == 'afmhot'):
            self.l = Num.array([0.  ,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.1 ,0.11,0.12,0.13,
                                0.14,0.15,0.16,0.17,0.18,0.19,0.2 ,0.21,0.22,0.23,0.24,0.25,0.26,0.27,
                                0.28,0.29,0.3 ,0.31,0.32,0.33,0.34,0.35,0.36,0.37,0.38,0.39,0.4 ,0.41,
                                0.42,0.43,0.44,0.45,0.46,0.47,0.48,0.49,0.5 ,0.51,0.52,0.53,0.54,0.55,
                                0.56,0.57,0.58,0.59,0.6 ,0.61,0.62,0.63,0.64,0.65,0.66,0.67,0.68,0.69,
                                0.7 ,0.71,0.72,0.73,0.74,0.75,0.76,0.77,0.78,0.79,0.8 ,0.81,0.82,0.83,
                                0.84,0.85,0.86,0.87,0.88,0.89,0.9 ,0.91,0.92,0.93,0.94,0.95,0.96,0.97,
                                0.98,0.99])
            self.r = Num.array([0.    ,0.0157,0.0392,0.0549,0.0784,0.0941,0.1176,0.1412,0.1569,0.1804,
                                0.1961,0.2196,0.2431,0.2588,0.2824,0.298 ,0.3216,0.3373,0.3608,0.3843,
                                0.4   ,0.4235,0.4392,0.4627,0.4863,0.502 ,0.5255,0.5412,0.5647,0.5804,
                                0.6039,0.6275,0.6431,0.6667,0.6824,0.7059,0.7294,0.7451,0.7686,0.7843,
                                0.8078,0.8314,0.8471,0.8706,0.8863,0.9098,0.9255,0.949 ,0.9725,0.9882,
                                1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,
                                1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,
                                1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,
                                1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,
                                1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ])
            self.g = Num.array([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.002 ,0.0255,0.0412,0.0647,0.0804,
                                0.1039,0.1275,0.1431,0.1667,0.1824,0.2059,0.2294,0.2451,0.2686,0.2843,
                                0.3078,0.3314,0.3471,0.3706,0.3863,0.4098,0.4255,0.449 ,0.4725,0.4882,
                                0.5118,0.5275,0.551 ,0.5745,0.5902,0.6137,0.6294,0.6529,0.6686,0.6922,
                                0.7157,0.7314,0.7549,0.7706,0.7941,0.8176,0.8333,0.8569,0.8725,0.8961,
                                0.9196,0.9353,0.9588,0.9745,0.998 ,1.    ,1.    ,1.    ,1.    ,1.    ,
                                1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,
                                1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ,1.    ])
            self.b = Num.array([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.0118,0.0275,0.051 ,0.0745,0.0902,0.1137,0.1294,0.1529,0.1686,0.1922,
                                0.2157,0.2314,0.2549,0.2706,0.2941,0.3176,0.3333,0.3569,0.3725,0.3961,
                                0.4196,0.4353,0.4588,0.4745,0.498 ,0.5137,0.5373,0.5608,0.5765,0.6   ,
                                0.6157,0.6392,0.6627,0.6784,0.702 ,0.7176,0.7412,0.7569,0.7804,0.8039,
                                0.8196,0.8431,0.8588,0.8824,0.9059,0.9216,0.9451,0.9608,0.9843,1.    ])
            self.r = Num.flip(self.r)
            self.g = Num.flip(self.g)
            self.b = Num.flip(self.b)
        elif (palette == 'hot'):
            self.l = Num.array([0.  ,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.1 ,0.11,0.12,0.13,
 0.14,0.15,0.16,0.17,0.18,0.19,0.2 ,0.21,0.22,0.23,0.24,0.25,0.26,0.27,
 0.28,0.29,0.3 ,0.31,0.32,0.33,0.34,0.35,0.36,0.37,0.38,0.39,0.4 ,0.41,
 0.42,0.43,0.44,0.45,0.46,0.47,0.48,0.49,0.5 ,0.51,0.52,0.53,0.54,0.55,
 0.56,0.57,0.58,0.59,0.6 ,0.61,0.62,0.63,0.64,0.65,0.66,0.67,0.68,0.69,
 0.7 ,0.71,0.72,0.73,0.74,0.75,0.76,0.77,0.78,0.79,0.8 ,0.81,0.82,0.83,
 0.84,0.85,0.86,0.87,0.88,0.89,0.9 ,0.91,0.92,0.93,0.94,0.95,0.96,0.97,
 0.98,0.99])
            self.r = Num.array([0.0416    ,0.06218969,0.09307422,0.11366391,0.14454844,0.16513813,
 0.19602266,0.22690719,0.24749688,0.27838141,0.2989711 ,0.32985563,
 0.36074017,0.38132985,0.41221439,0.43280407,0.46368861,0.48427829,
 0.51516283,0.54604736,0.56663705,0.59752158,0.61811127,0.6489958 ,
 0.67988033,0.70047002,0.73135455,0.75194424,0.78282877,0.80341846,
 0.83430299,0.86518752,0.88577721,0.91666174,0.93725143,0.96813596,
 0.9990205 ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ])
            self.g = Num.array([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.01960877,0.05049107,0.07107927,0.10196158,0.13284388,
 0.15343208,0.18431438,0.20490259,0.23578489,0.25637309,0.28725539,
 0.3181377 ,0.3387259 ,0.3696082 ,0.3901964 ,0.42107871,0.45196101,
 0.47254921,0.50343151,0.52401971,0.55490202,0.57549022,0.60637252,
 0.63725482,0.65784303,0.68872533,0.70931353,0.74019583,0.77107814,
 0.79166634,0.82254864,0.84313684,0.87401915,0.90490145,0.92548965,
 0.95637195,0.97696016,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ,1.        ,1.        ,
 1.        ,1.        ,1.        ,1.        ])
            self.b = Num.array([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.01176372,0.0426461 ,0.08896968,0.13529325,
 0.16617564,0.21249921,0.2433816 ,0.28970517,0.33602875,0.36691113,
 0.41323471,0.44411709,0.49044067,0.52132305,0.56764663,0.6139702 ,
 0.64485259,0.69117616,0.72205855,0.76838212,0.8147057 ,0.84558808,
 0.89191166,0.92279404,0.96911762,1.        ])

            self.r = Num.flip(self.r)
            self.g = Num.flip(self.g)
            self.b = Num.flip(self.b)
        elif (palette == 'gamma'):
            self.l = Num.array([0.0, 0.33, 0.66, 1.0])
            self.r = Num.array([0.3, 1.0, 0.0, 0.0])
            self.g = Num.array([0.0, 0.3, 1.0, 0.0])
            self.b = Num.array([0.0, 0.0, 0.3, 1.0])
        elif (palette == 'antigray' or palette == 'antigrey'):
            self.l = Num.array([0.0, 1.0])
            self.r = Num.array([1.0, 0.0])
            self.g = Num.array([1.0, 0.0])
            self.b = Num.array([1.0, 0.0])
        elif (palette == 'apjgray' or palette == 'apjgrey'):
            self.l = Num.array([0.0, 1.0])
            self.r = Num.array([1.0, 0.25])
            self.g = Num.array([1.0, 0.25])
            self.b = Num.array([1.0, 0.25])
        else:
            self.l = Num.array([0.0, 1.0])
            self.r = Num.array([0.0, 1.0])
            self.g = Num.array([0.0, 1.0])
            self.b = Num.array([0.0, 1.0])

def set_palette(cmap_name, nlevel):

    cm = pylab.get_cmap(cmap_name)

    colors = cm.colors

    step = len(colors)//nlevel

    cmap = colors[::step]

    for ii,f in enumerate(cmap):
        pgscr(ii+16, f[0], f[1], f[2])
