import sys,time
from xfab import tools,detector
from fabio import edfimage,tifimage
import variables,check_input
from scipy import ndimage
from scipy.stats import norm
import numpy as n
from PIL import Image

A_id = variables.refarray().A_id

class make_image:
	def __init__(self,graindata,killfile=None):
		self.graindata = graindata
		self.killfile = killfile

	def make_image(self,frame_number=None):
		"""
		makeimage script produces edf diffraction images using the reflection information
		
                Henning Osholm Sorensen, June 23, 2006.
		python translation Jette Oddershede, March 31, 2008
		"""
		peakshape = self.graindata.param['peakshape']
	
		if peakshape[0] == 0: # spike peak, 2x2 pixels
			peak_add = 1
			frame_add = 1
			peakwsig = 0
		elif peakshape[0] == 1: # 3d Gaussian peak
			peak_add = max(1,int(round(peakshape[1])))
			frame_add = max(1,int(round(peakshape[1])))
			peakwsig = peakshape[2]
		
		framedimy = self.graindata.param['dety_size']+2*frame_add
		framedimz = self.graindata.param['detz_size']+2*frame_add
		
		totalrefl = 0
		if frame_number == None:
			no_frames = range(len(self.graindata.frameinfo))
			print 'Generating diffraction images'
		else:
			no_frames = [frame_number]
	
		for i in no_frames:
			check_input.interrupt(self.killfile)
			t1 = time.clock()
			nrefl = 0
			frame = n.zeros((framedimy,framedimz))
			omega = self.graindata.frameinfo[i].omega
			omega_step = self.graindata.param['omega_step']
			# Jettes hack to add relative movement of sample and detector, modelled to be Gaussian in y and z direction with a spread of 1 micron
                # movement of 1 micron along x judged to be irrelevant, at least for farfield data 
			y_move = n.random.normal(0,1./self.graindata.param['dety_size'])
			z_move = n.random.normal(0,1./self.graindata.param['detz_size'])
                # loop over grains
			for j in range(self.graindata.param['no_grains']):
				# loop over reflections for each grain
				for k in range(len(self.graindata.grain[j].refs)):
					# exploit that the reflection list is sorted according to omega
					if self.graindata.grain[j].refs[k,A_id['omega']]*180/n.pi > \
						    omega+omega_step+2*peakwsig:
						break
					elif self.graindata.grain[j].refs[k,A_id['omega']]*180/n.pi < \
						    omega-2*peakwsig:
						continue
					dety = self.graindata.grain[j].refs[k,A_id['dety']]
					detz = self.graindata.grain[j].refs[k,A_id['detz']]
                #apply hack
#					dety = self.graindata.grain[j].refs[k,A_id['dety']] + y_move
#					detz = self.graindata.grain[j].refs[k,A_id['detz']] + z_move
					ndety = int(round(dety))
					ndetz = int(round(detz))
					yrange = range(ndety+frame_add-peak_add,ndety+frame_add+peak_add+1)
					zrange = range(ndetz+frame_add-peak_add,ndetz+frame_add+peak_add+1)
					intensity = int(round(self.graindata.grain[j].refs[k,A_id['Int']]))
					nrefl = nrefl + 1
					totalrefl = totalrefl + 1
					# Gaussian along omega
					if peakshape[0] == 1:
						fraction = norm.cdf((omega-self.graindata.grain[j].refs[k,A_id['omega']]*180/n.pi+omega_step)/(0.5*peakwsig))\
								  -norm.cdf((omega-self.graindata.grain[j].refs[k,A_id['omega']]*180/n.pi)/(0.5*peakwsig))
					else:
						fraction = 1.                       
					# Generate spikes, 2x2 pixels
					for y in yrange:
						for z in zrange:
							if y > 0 and y < framedimy and z > 0 and z < framedimz and abs(dety+frame_add-y) < 1 and abs(detz+frame_add-z) < 1:
#								frame[y-1,z] = frame[y-1,z] + fraction*intensity*(1-abs(dety+frame_add-y))*(1-abs(detz+frame_add-z))
								frame[y,z] = frame[y,z] + fraction*intensity*(1-abs(dety+frame_add-y))*(1-abs(detz+frame_add-z))

			# 2D Gaussian on detector					
			if peakshape[0] == 1:
				frame = ndimage.gaussian_filter(frame,peakshape[1]*0.5)
			# add background
			if self.graindata.param['bg'] > 0:
				frame = frame + self.graindata.param['bg']*n.ones((framedimy,framedimz))
			# add noise
			if self.graindata.param['noise'] != 0:
				frame = n.random.poisson(frame)
			# apply psf
			if self.graindata.param['psf'] != 0:
				frame = ndimage.gaussian_filter(frame,self.graindata.param['psf']*0.5)
			# resize, convert to integers and flip to same orientation as experimental frames
			frame = frame[frame_add:framedimy-frame_add,frame_add:framedimz-frame_add]

	                # limit values above 16 bit to be 16bit
			frame = n.clip(frame,0,2**16-1)
                        # convert to integers
			frame = n.uint16(frame)

		        #flip detector orientation according to input: o11, o12, o21, o22
			frame = detector.trans_orientation(frame,
							   self.graindata.param['o11'],
							   self.graindata.param['o12'],
							   self.graindata.param['o21'],
							   self.graindata.param['o22'],
							   'inverse')
			# Output frames 
			if '.edf' in self.graindata.param['output']:
				self.write_edf(i,frame)
			if '.tif' in self.graindata.param['output']:
				self.write_tif(i,frame)
			if '.tif16bit' in self.graindata.param['output']:
				self.write_tif16bit(i,frame)
			print '\rDone frame %i took %8f s' %(i+1,time.clock()-t1),
			sys.stdout.flush()
				
	def write_edf(self,framenumber,frame):
		e=edfimage.edfimage()
		e.data=frame
		e.dim2,e.dim1=frame.shape
		e.header = {}
		e.header['Dim_1']=e.dim1
		e.header['Dim_2']=e.dim2
		e.header['col_end']=e.dim1-1
		e.header['row_end']=e.dim2-1
		e.header['DataType']='UnsignedShort'
		e.header['Image']=1
		e.header['ByteOrder']='Low'
		e.header['time']=time.asctime()
		e.header['Omega']= self.graindata.frameinfo[framenumber].omega +\
		    self.graindata.param['omega_step']/2.0
		e.header['OmegaStep']=self.graindata.param['omega_step']
		e.header['grainfile']='%s/%s_%0.4dgrains.txt' \
			%(self.graindata.param['direc'],self.graindata.param['stem'],self.graindata.param['no_grains'])
		e.write('%s%s' %(self.graindata.frameinfo[framenumber].name,'.edf'))
				
	def write_tif(self,framenumber,frame):
		e=tifimage.tifimage()
		e.data=frame
		e.write('%s%s' %(self.graindata.frameinfo[framenumber].name,'.tif'))
	def write_tif16bit(self,framenumber,frame):
		size = frame.shape[:2][::-1]
		pilimage = Image.frombuffer('I',size,frame.tostring(),"raw",'I;16',0,1)
		pilimage.save('%s%s' %(self.graindata.frameinfo[framenumber].name,'.tif'))



