import numpy as np
from scipy import ndimage
from scipy import misc
from matplotlib import pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

def display_plt_image(data_array, name = "data_array"):
    
    # Manually control the size of the figure to assure correct display
    fig = plt.figure(dpi=100, figsize=(10, 10))       

    # Fill the image with data
    plt.imshow(data_array, 
               cmap='hot',                  # others are possible
               interpolation='nearest',     
               origin='lower')    

    #save it 
    fig.savefig(name,
                dpi=100,                    # assures the correct display
                bbox_inches='tight',        # No white space around plot
                pad_inches=0)               # idem

    # display it
    plt.show()
                
    # always close an image. It stays 'active' in mem otherwise
    plt.close()

# Display an image from data    
a = np.array([[1, 2, 0, 0],  [5, 3, 0, 4],   [0, 0, 0, 7],  [9, 3, 0, 0]])
k = np.array([[1,1,1],[1,1,0],[1,0,0]])
# Setting cval=1.0 is equivalent to padding 
convolved = ndimage.convolve(a, k, mode='constant', cval=0.0) 
name = "data_array.png"
display_plt_image(convolved, name)


# Now open data from an existing image
data_as_nd_array = misc.imread("einstein.jpg")
plt.imshow(data_as_nd_array)
plt.show()
plt.close()

# add some black bars
data_as_nd_array[100:200,:] = 0
data_as_nd_array[:,100:200] = 0
plt.imshow(data_as_nd_array)
plt.show()
plt.close()

# Inspect the data.
print (data_as_nd_array.shape)
print (data_as_nd_array[0][0])
# We have 4 channels (rgb + z)

for idx in range(500, data_as_nd_array.shape[0]):
    for idy in range(500, data_as_nd_array.shape[1]):
        if data_as_nd_array[idx][idy][0] < 240:
            data_as_nd_array[idx][idy][0] = 0
            
        data_as_nd_array[idx][idy][1] = 200

plt.imshow(data_as_nd_array)
plt.show()
plt.close()

# A 3d plot
X = np.arange(-5, 5, 0.25)
Y = np.arange(-5, 5, 0.25)
X, Y = np.meshgrid(X, Y)
R = np.sqrt(X**2 + Y**2)
Z = np.sin(R)

fig = plt.figure()
ax = Axes3D(fig)
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.viridis)

plt.show()



