import snappy
from PIL import Image #Python Imaging Library
from math import *
from time import *   #only used to check running times

def rgb2int(r,g,b):
    return r + 256*(g + 256*b)

def frac(x):
    return x - floor(x)

def positive(vol):   #all positive tetrahedra, coloured green, shaded in bands by volume
    ratio = vol/max_vol
    ratio = frac(num_bands * ratio)
    return rgb2int(0,int(127 + ratio * 128),0)

def negative(vol):   #some negative tetrahedra, coloured blue, shaded in bands by volume
    ratio = vol/max_vol
    ratio = frac(num_bands * ratio)
    return rgb2int(0,0,int(127 + ratio * 128))

def no_sol(vol):
    return rgb2int(127,127,127)  #grey

def flat(vol):
    return rgb2int(255,0,0)  #red

def na(vol):
    return rgb2int(0,127,127)  #cyan

def degen(vol):
    return rgb2int(127,0,127)  #purple

def unrec(vol):
    return rgb2int(255,255,255)   #white

def is_integral(z, tol):
    return z - floor(z) < tol * 0.9

def volextend(M):
    if M.solution_type() == 'no solution found':
        return 0
    else: 
        return M.volume()

choose_col = { 'not attempted' : na, 
	 'all tetrahedra positively oriented' : positive,
	 'contains negatively oriented tetrahedra' : negative,
	 'contains flat tetrahedra' : flat,
	 'contains degenerate tetrahedra' : degen,
	 'unrecognized solution type' : unrec,
	 'no solution found' : no_sol }
###########################################

num_bands = 10  #how many bands of colour from full to zero volume

mfld = 'm004'  

xmin = -5.0
xmax = 5.0
xsteps = 400


ymin = -5.0
ymax = 5.0
ysteps = 400

###########################################
xstep = (xmax - xmin) / xsteps
ystep = (ymax - ymin) / ysteps

M_master = snappy.Manifold(mfld)
max_vol = M_master.volume()

im = Image.new("RGB", (xsteps,ysteps))

t = time()

for i in range(xsteps):
    for j in range(ysteps):
        if(is_integral(i*xstep + xmin, xstep) and is_integral(j*ystep + ymin, ystep)):
            im.putpixel((i,ysteps-1-j), rgb2int(0,0,0)) #black at integral points, reflect across horizontal axis due to pixel coord conventions
        else:
            M = M_master.copy()  #cloning is faster than calling get_manifold again
            M.dehn_fill((i*xstep + xmin,j*ystep + ymin),0)
            col = choose_col[M.solution_type()](volextend(M))
            im.putpixel((i,ysteps-1-j), col)  #reflect across horizontal axis due to pixel coord conventions
            
print 'done ' + mfld + ' in time: ' + str(time() - t)
im.save(mfld + '.png')