from PIL import Image
import math

def normalized(a):
    
    factor = 1.0/math.sqrt( np.sum(a*a)) # normalize
    return a*factor

def my_gauss(im):
	im_smooth = im.astype(float)
	kernel = np.array([0.5,1.0,0.5]).astype(float);
	kernel=normalized(kernel)

	im_smooth = scipy.ndimage.convolve(im_smooth, kernel[np.newaxis])
	im_smooth = scipy.ndimage.convolve(im_smooth, kernel[np.newaxis].T)

	return im_smooth

#calculate a local internal shadowing/AO aproximation using 'difference of gaussians'
def shadow(im):
    im1 = im.astype(float)
    im0 = im1.copy()
    im00 = im1.copy()
    im000 = im1.copy()
    for i in range(0,2):
        im00 = my_gauss(im00)

    for i in range(0,64):
        im0 = my_gauss(im0)

    for i in range(0,128):
        im1 = my_gauss(im1)
    im000=normalized(im000)
    im00=normalized(im00)
    im0=normalized(im0)
    im1=normalized(im1)
    im00=normalized(im00)

    shadow=im00*2.0+im000-im1*2.0-im0 # np.clip(im1, 0.0, 1000000.0)
    shadow=normalized(shadow)
    mean = np.mean(shadow)
    rmse = rms_error(shadow)
    shadow = np.clip(shadow, mean-rmse*2.0,mean+rmse*0.5)

    return shadow

# ....
# added to the end of main()

   # save the pure normal map and AO/shadow images as seperatejpgs
    scipy.misc.imsave("normal.jpg", normal_map)


    im_shadow = shadow(im)



    scipy.misc.imsave("shadow.jpg", im_shadow)

    # read these images back in, and combine them to make an image with RGB=normal, Alpha=AO
    # really I did this because i couldn't get it working from numpy arrays directly.
    imgsh=Image.open("shadow.jpg")
    #imgsh.show()
    imgnorm=Image.open("normal.jpg")
    #imgnorm.show()
    imgnormao = np.zeros( (img.size[0],img.size[1],4), np.float)
    imgnormao[ :, :, 0:3] = np.asarray(imgnorm)[:,:,0:3]   #RGB = normal
    imgnormao[ :, :, 3] = np.asarray(imgsh)[:,:]                #Alpha = AO

    scipy.misc.imsave(output_file, imgnormao)
    imgnormao=Image.open(output_file)           # save the RGBA image containing Normal,AO