pro nonrigid

;Read in two images
template = read_tiff('..\gyj\ct_001.tif')
target = read_tiff('..\gyj\ct_002.tif')

;Get the control points on both images
ret = 'Yes'
while ret eq 'Yes' do begin
	temp_pts = GetPoints(template,0,'Template')
	targ_pts = GetPoints(target,1,'Target')
	;Before going forward, make sure there are the same number of control points in both images
	pts_in_temp = n_elements(temp_pts)/2
	pts_in_targ = n_elements(targ_pts)/2
	if pts_in_temp ne pts_in_targ then begin
		ret = Dialog_Message('The number of Control points should be the same!  ReDo?',/Question,title='Info')
		; if no redo, then cut the longer one
		if ret EQ 'No' then begin
			if pts_in_temp GT pts_in_targ then begin
				temp_pts = temp_pts[*,0:pts_in_targ-1]
				pts_in_temp = pts_in_targ
			endif else begin
				targ_pts = targ_pts[*,0:pts_in_temp-1]
				pts_in_targ = pts_in_temp
			endelse
		endif
	endif else begin
	; if the same number, go out of the loop
		ret = 'No'
	endelse
endwhile

; Now we have the same number of control points in both images
; and they are in the same sequence, for those corresponding points
;Do the TPS calculation
tpsinfo = CalcTPS(temp_pts,targ_pts,pts_in_temp)
;tpsinfoBack = CalcTPS(targ_pts,temp_pts,pts_in_temp)

;check if tps can work properly on the control points
;zero=testTps(tpsinfo)

;Deform Image
defimg = DeformImage(template,tpsinfo)
;Find the difference Image
diff1 = defimg-template
;diff1[where (diff1 lt 0)] = 0
diff1 = datascale(diff1,256)
diff2 = defimg-target
;diff2[where (diff2 lt 0)] = 0
diff2 = datascale(diff2,256)
;diff3 = diff1-diff2

def_size = size(defimg)
;window,2,xsize=def_size(1),ysize=def_size(2),title='deformed image-temp'
;tv,diff3
window,3,xsize=def_size(1),ysize=def_size(2),title='deformed image-targ'
;tv,diff2
tv,defimg

write_tiff,'..\gyj\ct_def_2.tif',defimg
write_tiff,'..\gyj\ct_diff1_2.tif',diff1
write_tiff,'..\gyj\ct_diff2_2.tif',diff2

end

;Function to check if the TPS can work properly on those control points
; only check those control points themself, then see what happened
function testTps,tpsinfo
	test_pts = tpsinfo.points
	print,'the control points in template image are:',test_pts
	print,' the control points in target image are:',tpsinfo.points_b
	def_pts = deformpoints(test_pts,tpsinfo)
	print,' the points after deformation are:',def_pts
	return,0
end

;Function to do the deformation of the image
; Input is the image data to be deformed and the applied tpsinfo
function DeformImage,img,tpsinfo
	;first get the image dimension
	d_size = size(img)
	nRows = d_size(1)
	nCols = d_size(2)

	defImg = intarr(nRows,nCols)
	;it maybe too huge for IDL to do the whole image at one time
	; so we will do one row at one time
	for i=0,nRows-1 do begin
		pts = dblarr(2,nCols)
		yVect = findgen(nCols)
		pts[0,0:nCols-1] = i
		pts[1,0:nCols-1] = yVect

		defPts = DeformPoints(pts,tpsinfo)
		defPts = round(defPts)
		;check the boundary points
		for j=0,nCols-1 do begin
			if(defPts[0,j] GE 0 and defPts[0,j] lt nRows $
				and defPts[1,j] GE 0 and defPts[1,j] lt nCols) then begin
				;ptLocs = nCols*pts[0,j]+pts[1,j]
				;defptLocs = nCols*defPts[0,j]+defPts[1,j]
				defImg[pts[0,j],pts[1,j]] = img[defPts[0,j],defPts[1,j]]
			endif
		endfor
		;locs = where(defPts[0,*] GE 0 AND defPts[0,*] lt nRows $
		;		AND defPts[1,*] GE 0 AND defPts[1,*] lt nCols)

		;ptLocs = nCols*pts[0,locs]+pts[1,locs]
		;defPtLocs = nCols*defPts[0,locs]+defPts[1,locs]
		;if locs gt 0 then	begin
		;	defImg[locs] = img[locs]
		;endif
		;defImg[ptLocs] = img[defPtLocs]
	endfor

	return, defImg
end

;Function to deform the points
; using a thin-plate spline function
; TPSINFO contains the necessary information for using TPS
; f(x,y) = a1+a2*x+a3*y+sum(w_i U(p_i-(x,y)));
function DeformPoints,pts,tpsinfo
	num_pts = n_elements(pts)/2
	;since those coordinates must appear in pairs
	P = dblarr(num_pts)
	P[*] = 1
	P = [transpose(P),pts]
	;P=[1 x y;1 x y;...]

	ct_pts = tpsinfo.points
	num_ct = n_elements(ct_pts)/2

	;define two unit array to be used later
	unitp = dblarr(num_pts)
	unitp[*] = 1
	unitct = dblarr(num_ct)
	unitct[*] = 1

	; Fill a matrix with the x location, and a matrix with the y location
	px = unitp # ct_pts[0,*]
	;Find the distance from every control point to every point
	x = unitct # pts[0,*]
	px = px - transpose(x)

	py = unitp # ct_pts[1,*]
	;Find the distance from every control point to every point
	y =  unitct # pts[1,*]
	py = py - transpose(y)

	;Now get the Eulerian distance. We need the square difference
	r_sq = px^2+py^2

	r_size = size(r_sq)
	; The U matrix is r^2*log(r^2)
	U = dblarr(r_size(1),r_size(2))
	;filter the element with value 0 to avoid the exception of LOG
	;locs = where(r_sq > 0)
	;U[locs] = r_sq[locs]*alog(r_sq[locs])
	for i=0,r_size(1)-1 do begin
	for j=0,r_size(2)-1 do begin
		if r_sq[i,j] gt 0 then U[i,j] = r_sq[i,j]*alog(r_sq[i,j])
	endfor
	endfor

	;defPts should be M*2, the first part is the affine portion
	; The second part is the thin-plate spline portion
	defPts = transpose(transpose(tpsinfo.a) ## transpose(P) + transpose(tpsinfo.w) ## U)

	return,defPts
end

; Function to calculate the TPS parameters
function CalcTPS,temp_pts,targ_pts,pts_count
	;Read in the control points. No need to count again

	; First, Calculate the U matrix of each point to the other point
	x = temp_pts[0,*]
	y = temp_pts[1,*]
	; Fill a matrix with the x location, and a matrix with the y location
	px = dblarr(pts_count)
	px[*] = 1
	px = px # x
	px = px-transpose(px)
	py = dblarr(pts_count)
	py[*] = 1
	py = py # y
	py = py-transpose(py)
	;Now get the Eulerian distance. We need the square difference
	r_sq = px^2+py^2

	; The U matrix is r^2*log(r^2)
	r_size = size(r_sq)
	U = dblarr(r_size(1),r_size(2))
	;filter the element with value 0 to avoid the exception of LOG
	locs = where(r_sq > 0)
	U[locs] = r_sq[locs]*alog(r_sq[locs])

	; P is [1 x1,y1;1 x2 y2;...]
	P = dblarr(pts_count)
	P[*] = 1
	P = [transpose(P),x,y]
	; Now is L's turn
	L = dblarr(pts_count+3,pts_count+3)
	L[0:pts_count-1,0:pts_count-1] = U
	L[0:pts_count-1,pts_count:pts_count+2] = transpose(P)
	L[pts_count:pts_count+2,0:pts_count-1] = P

	;Now define the V vector, which is the orignal control points
	;add three more 0s
	V = dblarr(2,pts_count+3)
	V[*,0:pts_count-1] = targ_pts

	; Now to determine the weights.
	W = invert(L)##V

	tpsinfo = { $
		points:temp_pts, $
		w:W[*,0:pts_count-1], $
		a:W[*,pts_count:pts_count+2] $
		}
	return, tpsinfo
end

;Function to return the control points selected interactively
function GetPoints,data,winid,wintitle

	im_size = size(data)
	im_height = im_size(1)
	im_width = im_size(2)

	window,winid,xsize = im_height,ysize = im_width,title=wintitle
	TV,data;,/order
	; Dialog_Message('Left Mouse: Select the control points; Both: End the selection')
	print,'To select the control points:'
	print,'Left Button: select point'
	print,'Right Button: End Selection'

	Cursor,xi,yi,/UP,/DEVICE
	PLOTS, xi, yi, PSYM=7, /DEVICE

	tmp_pts = [xi,yi]
	print,xi,yi,data[xi,yi]
	;Continue to collect control points until right mouse button.
	while !MOUSE.BUTTON ne 4 do begin
   		CURSOR, xi, yi, /UP, /DEVICE
   		PLOTS, xi,yi, PSYM = 7, /DEVICE
		tmp_pts = [[tmp_pts],[xi,yi]]
		print,xi,yi,data[xi,yi]
	endwhile

	pts_size = size(tmp_pts)
	num_pts = pts_size(2)
	print,'You have ',num_pts,' control points in this image'
	;Link those control points together to show their correspondance
	x0 = tmp_pts[0,0]
	y0 = tmp_pts[1,0]
	for i=1,num_pts-1 do begin
		x1 = tmp_pts[0,i]
		y1 = tmp_pts[1,i]
		plots, [x0,x1], [y0,y1],/device
		x0 = x1
		y0 = y1
	endfor

	annotate

	return,tmp_pts
end