#!/usr/bin/python
# -*- coding: utf-8 -*-

import math
import PIL, PIL.Image, PIL.ImageDraw


class ImageBuffer:
	
	def __init__(self):
		self.buffer = None
		self.dims = None
	
	def saveImage(self, filename):
		surface = PIL.Image.new('L', self.dims, 0)
		dc = PIL.ImageDraw.Draw(surface)
		maxval = max(self.buffer)
		if maxval == 0:
			raise ZeroDivisionError('Cannot normalize an image whose max value is zero')
		width = self.dims[0]
		for i in range(len(self.buffer)):
			x = i % width
			y = i / width
			#print " ",i,"=>",x,",",y
			val = self.buffer[i] / maxval * 255
			dc.point((x,y), fill=val)
		surface.save(filename)
	
	def sample(self, point):
		pt_x = int(round(point[0]))
		pt_y = int(round(point[1]))
		if pt_x < 0 or pt_x >= self.dims[0] or pt_y < 0 or pt_y >= self.dims[1]:
			return 0.0
		return self.buffer[pt_y*self.dims[0]+pt_x]

	
# Tests (see the test_all function below) show that using window_fact=4 is a safe - all of the ellipse
# fit in the quad, whatever the angle.

class BivNorm(ImageBuffer):
	def __init__(self, sigma1, sigma2, rot=0.0, window_fact=4.0, rot_is_rad=False):
		self.sigma1 = sigma1
		self.sigma2 = sigma2
		self.rot = rot if rot_is_rad else rot/180.0*math.pi
		self.window_fact = window_fact
		self.center = None
		
	def bivgauss(self, x, y):
		return (1. / (2.*math.pi*self.sigma1*self.sigma2)) * math.exp(-0.5 * (x*x/self.s1 + y*y/self.s2))
		
	def render(self):
		
		self.s1, self.s2 = self.sigma1*self.sigma1, self.sigma2*self.sigma2 
		
		sm = max(self.sigma1,self.sigma2)
		H = W = self.window_fact * math.sqrt(2) * sm
		center = [W/2.0, H/2.0]
		
		#print "W=%f, H=%f" % (W,H)
		
	
		width = int(math.ceil(W))
		if width%2==0: width += 1
		height = int(math.ceil(H))
		if height%2==0: height += 1
		
		#print "width=%f, height=%f" % (width, height)
		
		# Generate the "mesh"
		sum = 0
		xmin, ymin = -width/2+1, -height/2+1
		buffer = [0.0]*width*height
		cost = math.cos(self.rot)
		sint = math.sin(self.rot)
		msint = -sint
		for y in range(ymin,height/2+1):
			baseidx = height*(y-ymin)
			for x in range(xmin, width/2+1):
				idx = baseidx + (x-xmin)
				#print " ",x,",",y,"=>",idx
				x2 = x * cost - y * sint
				y2 = x * sint + y * cost
				v = self.bivgauss(x2, y2)
				#print "%d,%d = %f"  % (x,y,v)
				sum += v
				buffer[idx] = v
		#print "sum=",sum # sum seems to converge to Z, where Z=1.15580194518 for s1,s2=3,1 { and Z=<another value> for s1,s2=<other values> }
		
		# Normalize (needed, see comment above)
		buffer = map(lambda x: x/sum, buffer)
		
		# Check norm
		#sum=0
		#for x in buffer:
		#	sum += x
		#print " final sum=",sum 
		
		self.dims = (width, height)
		self.buffer = buffer
		self.center = (-xmin, -ymin)
		#print self.buffer
		
		

class Surface(ImageBuffer):

	def __init__(self, dims):
		self.dims = dims
		self.buffer = [0.0]*self.dims[0]*self.dims[1]
		
	def draw(self, bn, at):
		bnw = bn.dims[0]
		bnh = bn.dims[1]
		width = self.dims[0]
		height = self.dims[1]
		for y in range(bnw):
			for x in range(bnh):
				dst_x = int(round(x - bn.center[0] + at[0]))
				dst_y = int(round(y - bn.center[1] + at[1]))
				if dst_x < 0 or dst_x >= width or dst_y < 0 or dst_y >= height:
					continue
				dst_idx = width*dst_y + dst_x
				src_idx = bnh*y + x
				self.buffer[dst_idx] += bn.buffer[src_idx]	

def test_all():
	ranges = [0.5, 1.0, 1.5, 2.0, 3.0, 5.0, 10.0]
	ratios = ((3,1), (10,5), (2, 4), (30,15))
	rots = [0, 30, 45, 60, 90]
	# ranges = [2.0]
	# ratios = ((2,1),)

	for ratio in ratios:
		for rng in ranges:
			for rot in rots:
				print "Creating bivnorm(%d,%d,%f,%f)" % (ratio[0], ratio[1], rot, rng)
				bv1 = BivNorm(ratio[0], ratio[1], rot=rot, window_fact=rng)
				bv1.render()
				filename = 'ellipse_%d-%d_%f_%f.png' % (ratio[0], ratio[1], rot, rng)
				bv1.saveImage(filename)

def test_draw():
	points = [(100,100,0), (100,200,30), (200, 100, 45), (200, 200, 90)]
	ratio = (20,10)
	dims = (300,300)
	print "Creating %dx%d surface with %d no overlapping points" % (dims[0],dims[1],len(points))
	surface = Surface(dims)
	for point in points:
		bn = BivNorm(*ratio, rot=point[2])
		bn.render()
		surface.draw(bn, point[:2])
	filename = 'ellipse_surface_4_no_overlap.png'
	surface.saveImage(filename)

	import random
	dims = (1000,1000)
	numpoints = [20, 50, 100, 200, 500, 1000]
	functions = ['uniform','gauss']
	for np in numpoints:
		for function in functions:
			points = []
			print "Creating %dx%d surface with %d random points (%s distribution)" % (dims[0],dims[1],np,function)
			if function == 'gauss':
				for i in range(np):
					x, y = int(random.uniform(0,dims[0])), int(random.uniform(0,dims[1]))
					angle = random.uniform(-180,180)
					points.append((x,y,angle))
			else:
				for i in range(np):
					x, y = int(random.gauss(dims[0]/2, dims[0]/2)), int(random.gauss(dims[1]/2, dims[1]/2))
					angle = random.uniform(-180,180)
					points.append((x,y,angle))
			
			surface = Surface(dims)
			for point in points:
				bn = BivNorm(*ratio, rot=point[2])
				bn.render()
				surface.draw(bn, point[:2])
			filename = 'ellipse_surface_%d_random_%s.png' % (np,function)
			surface.saveImage(filename)
	
	
def draw_ref():
	print "Creating reference ellipses used by the scripts"
	# Note: this must be manually kept in sync with the actual values used in GeneralMetrics.Coherency
	mmap_setups = {
		'20_10': {'bifurcation':(20,10), 'ridge_ending':(20,10), 'unknown':(40,20), 'unknown2':(50,40)}
	}
	for setup in mmap_setups:
		for mtype, ratio in mmap_setups[setup].iteritems():
			surf = Surface((300,300))
			bn = BivNorm(*ratio, rot=30)
			bn.render()
			surf.draw(bn, (150,150))
			filename = 'ellipse_ref_%s_%s.png' % (setup, mtype)
			surf.saveImage(filename)
	
	
if __name__ == '__main__':
	import sys
	if 'all' in sys.argv:
		test_all()
	elif 'draw' in sys.argv:
		test_draw()
	elif 'ref' in sys.argv:
		draw_ref()
	else:
		print "Need an argument in %s" % ', '.join("'"+x+"'" for x in ['all','draw','ref'])
		
	