
import numpy as np   
import matplotlib.pyplot as plt  
import matplotlib

def imread_int( filename ):
    img = plt.imread( filename )
    return np.array( img*255, dtype = 'int')

def colorize_red(grayvalue, cmap):
    return cmap[0][grayvalue][0]

def colorize_green(grayvalue, cmap):
    return cmap[0][grayvalue][1]

def colorize_blue(grayvalue, cmap):
    return cmap[0][grayvalue][2]

def imcolorize(highfield, colormap):
    img = np.zeros((highfield.shape[0],highfield.shape[1],3))
    for i in xrange(highfield.shape[0]):
        for j in xrange(highfield.shape[1]):
            g = highfield[i][j]
            img[i][j][0] = colorize_red(g, colormap)
            img[i][j][1] = colorize_green(g, colormap)
            img[i][j][2] = colorize_blue(g, colormap)
    return img;
   


def exc2a():    
    mappath = 'u2Images/colormap.png'
    colormap = imread_int(mappath)

    plt.subplot(221)
    plt.imshow(plt.imread(mappath))    
    plt.subplot(222)
    rng = np.arange(255)
    plt.plot(rng, np.array([colorize_red(g,colormap) for g in rng]), color='red')
    plt.plot(rng, np.array([colorize_green(g,colormap) for g in rng]), color='green')
    plt.plot(rng, np.array([colorize_blue(g,colormap) for g in rng]), color='blue')
    plt.show()
    
def exc2b():
    colormap = imread_int('u2Images/colormap.png')
    highfield = imread_int('u2Images/usa_small.png')
    plt.imshow(imcolorize(highfield, colormap)/255)
    plt.show()
    

def exc2c():
    colormap = imread_int('u2Images/colormapr.png')
    highfield = imread_int('u2Images/mars-nordpol.png')
    plt.imshow(imcolorize(highfield, colormap)/255)
    plt.show()

exc2a()
#exc2b()
#exc2c()