; NAME :	corralgn
; PURPOSE :
;	Align two images by maximizing the cross correlation.
;	Cross correlations can be taken on sub-areas, and 
;	the motions of the sub-areas can be used to warp
;	the sample image with respect to the reference image.
; CATEGORY :
; 	image processing
; CALLING SEQUENCE :
;	pro corralgn, refimage,image, refpos, pos, $
;	ndiv=ndiv, $
;	monitor=monitor, refinfo=refinfo, info=info, $
;	norder=norder, warpedimage=warpedimage, samewarp=samewarp, $
;	px=px, py=py, xshift=xshift, yshift=yshift, $
;	xresidual=xresidual, yresidual=yresidual
;
; INPUTS :
;	refimage = 2-d array (reference image)
;	image    = 2-d array (image to be aligned with respect to refimage)
; OPTIONAL INPUT PARAMETERS : 
;	refpos   = 1-d array of two elements (x,y), giving the reference
;	           position (e.g. sunspot location) on the reference image
;	pos      = 1-d array of two elements (x,y), giving the reference
;	           position on the sample image
;	if not supplied, refpos=[0,0] and pos=[0,0] are assumed
; KEYWORD PARAMETERS :
;	ndiv = number of subarea division, either a scalar or
;	       a two element vector ndiv=[nxdiv,nydiv]
;	monitor=0,1,2,3 : monitor level (0 is terse, 3 is verbose)
;	refinfo=title string for refimage
;	info=title string for sample image
;	norder=order of polynomical warping
;	       0: uniform lateral shift
;            >0: polynomial warp
;            <0: piecewise bilinear warp
;	samewarp=2-d or 3-d array for which the warping procedure applied to
;	         the sample image is also carried out
; OPTIONAL OUTPUTS :
;	warpedimage= warped sample image
;	samewarp= warped image of array 'samewarp' (overwritten)
;	px, py= polynomical coefficients of warping
;	xshift, yshift=shift vector
;	xresidual, yresidual=residual shift vector
;	(polynomical warping subtracted)
;
; COMMON BLOCKS : 	none
; SIDE EFFECTS :	a window number 'mywindow' is created
; RESTRICTIONS :	X-window environment is assumed
; PROCEDURE :
;
; MODIFICATION HISTORY :
;	T.Sakurai, 1993
;
;------------------------------------------------------------------------
;
pro peak_2d, array, xloc, yloc
;
; find a local maximum of data given in array
;
; assume a quadratic function z=array(x,y)
; z = a*x**2 + b*y**2 + c*x*y + d*x + e*y +f
;
; peak of z is at
;  2*a*x + c*y + d =0
;  c*x + 2*b*y + e =0
; that is,
;  x = (c*e -2*b*d)/(4*a*b -c*c)
;  y = (c*d -2*a*e)/(4*a*b -c*c)
;
	nsize = size(array)
	if (n_elements(nsize) eq 4 ) then nz = fix(nsize(3)) else nz=1

	fit = quadfit(array, coeff=coeff)

	a=coeff(0)
	b=coeff(1)
	c=coeff(2)
	d=coeff(3)
	e=coeff(4)
	f=coeff(5)
	det=4.0*a*b - c*c
	xloc = (c*e -2.0*b*d)/det
	yloc = (c*d -2.0*a*e)/det

	return
	end

pro plotarrows, xvec,yvec, xrange,yrange, dx,dy, mag=magnify
	if (not keyword_set(magnify) ) then magnify=1.0
	ndim = size(xvec)
	nx = fix(ndim(1))
	if (ndim(0) gt 1) then ny = fix(ndim(2)) else ny=1

	head=5.0
	fa=0.3
	xb=fltarr(2)      &     yb=xb
	xa1=xb      &     xa2=xa1
	ya1=yb      &     ya2=ya1
	position=[xrange(0),yrange(0),xrange(1),yrange(1)]
	plot, xb, yb, /nodata, /noerase, $
	xrange=xrange, yrange=yrange, $
	position=position, /dev , $
	xstyle=5, ystyle=5, thick=1
	dxh = dx/2
	dyh = dy/2

	for y=0, ny-1 do for x=0, nx-1 do begin
		x1h = xrange(0) + x*dx +dxh
		y1h = yrange(0) + y*dy +dyh
		xb(0) = x1h
		yb(0) = y1h
		xb(1) = x1h + magnify*xvec(x,y)
		yb(1) = y1h + magnify*yvec(x,y)
		length= sqrt(xvec(x,y)^2 + yvec(x,y)^2)
		if (magnify*length lt 1.0 ) then goto, loopend
		oplot, xb, yb
		xa1(0)=xb(1) & xa2(0)=xb(1)
		ya1(0)=yb(1) & ya2(0)=yb(1)
		xa1(1)=xa1(0)-head/length*(xvec(x,y)+fa*yvec(x,y))
		xa2(1)=xa2(0)-head/length*(xvec(x,y)-fa*yvec(x,y))
		ya1(1)=ya1(0)-head/length*(yvec(x,y)-fa*xvec(x,y))
		ya2(1)=ya2(0)-head/length*(yvec(x,y)+fa*xvec(x,y))
		oplot, xa1, ya1
		oplot, xa2, ya2
	loopend:
	endfor
	return
end

;
pro corralgn, refimage,image, refpos, pos, $
	ndiv=ndiv, $
	monitor=monitor, refinfo=refinfo, info=info, $
	norder=norder, warpedimage=warpedimage, samewarp=samewarp, $
	px=px, py=py, xshift=xshift, yshift=yshift, $
	xresidual=xresidual, yresidual=yresidual

	if ( n_elements(monitor) eq 0 ) then monitor=0
	if ( n_elements(norder) eq 0 ) then norder=1
	if ( n_elements(ndiv) eq 1 ) then begin
		nxdiv=ndiv
		nydiv=ndiv
	endif else begin
		nxdiv=ndiv(0)
		nydiv=ndiv(1)
	endelse

	nsize = size(image)
	nx = fix(nsize(1))
	ny = fix(nsize(2))
	warpedimage = fltarr(nx,ny)
	ssize = size(samewarp)
	if(ssize(0) eq 3) then nz=fix(ssize(3)) $
	else if(ssize(0) eq 2) then nz=1 $
	else nz=0

	nx1 = nx / nxdiv
	ny1 = ny / nydiv
	nx1=(nx1/2)*2
	ny1=(ny1/2)*2
      nxh = nx1/2 &     nyh = ny1/2
	x1 = nx1*indgen(nxdiv)
	y1 = ny1*indgen(nydiv)
	xco = intarr(nxdiv, nydiv)
	yco = intarr(nxdiv, nydiv)
	for yy=0, nydiv-1 do xco(*,yy) = nx1*(indgen(nxdiv) + 0.5)
	for xx=0, nxdiv-1 do yco(xx,*) = ny1*(indgen(nydiv) + 0.5)
; (xco,yco) is the center of sub-area
	xci = fltarr(nxdiv,nydiv)
	yci = fltarr(nxdiv,nydiv)
	xh = x1+nxh	&	xh(0)=0	&	xh(nxdiv-1)=nx-1
	yh = y1+nyh	&	yh(0)=0	&	yh(nydiv-1)=ny-1
	xshift = fltarr(nxdiv,nydiv)
	yshift = fltarr(nxdiv,nydiv)
	cosbell = hanning(nx1,ny1)
	buffer = complexarr(nx1,ny1)
	maxref = max(abs(refimage))
	diflevel = 0.05

      ft1 = complexarr(nx1,ny1)
      corr = fltarr(nx1,ny1)
      map = lonarr(nx1,ny1)

      for j=0,ny1-1 do for i=0,nx1-1 do begin
            ii=i+nxh    &     jj=j+nyh
            if (i ge nxh ) then ii = i -nxh
            if (j ge nyh ) then jj = j -nyh
            map(i,j)=long(ii)+long(jj)*long(nx1)
      endfor

	xo = intarr(nx,ny)	&	xo=fltarr(nx,ny)
	yo = intarr(nx,ny)	&	yo=fltarr(nx,ny)
	for j=0, ny-1 do xo(*,j)=indgen(nx)
	for i=0, nx-1 do yo(i,*)=indgen(ny)

	ft0 = complexarr(nx1,ny1,nxdiv,nydiv)
	power0 = fltarr(nxdiv,nydiv)
	for yy=0,nydiv-1 do for xx=0,nxdiv-1 do begin
		xs = x1(xx)	&	ys = y1(yy)
		xe = xs + nx1 -1
		ye = ys + ny1 -1
		buffer = cosbell * refimage(xs:xe, ys:ye)
		ft0(*,*,xx,yy)=  conj( fft( buffer, -1) )
		power0(xx,yy) = total (buffer*buffer)
	endfor
	psize = size(pos)
	rsize = size(refpos)

	w = 1
	corrmin=0.5
	interp=1
	yn = ' '
	d=20
	x0=0
	yd=3*d
	yu=yd+ny+2*d
	magnify=10
	mywindow=2
	charsize = 1.5
	if (monitor ge 1) then $
		window, mywindow, xsize=x0 + 3*nx + 2*d, ysize=yd + 2*ny + 2*d

	if (psize(0) ge 1 and rsize(0) ge 1 ) then begin
		delx = pos(0)-refpos(0)
		dely = pos(1)-refpos(1)
		if (delx gt 0.0) then	idelx = fix(delx + 0.5) $
		else				idelx = fix(delx - 0.5)
		if (dely gt 0.0) then	idely = fix(dely + 0.5) $
		else				idely = fix(dely - 0.5)
; delx,y = reference_pos(image) - reference_pos(refimage)
; shift by integers, in order not to degrade image resolution
		image1 = shift(image,-idelx,-idely)
	endif else begin
		delx = 0	&	dely = 0
		idelx = 0	&	idely = 0
		image1 = image
	endelse

	back:
	for yy=0, nydiv-1 do for xx=0,nxdiv-1 do begin
		xs = x1(xx)	&	ys = y1(yy)
		xe = xs + nx1 -1
		ye = ys + ny1 -1
		buffer = cosbell * image1(xs:xe, ys:ye)
		ft1=fft( buffer, -1)
		power1=total(buffer*buffer)
		ft1 = ft1 * ft0(*,*,xx,yy) *float(nx1)*float(ny1)
		corr1 = float(fft(ft1,1))
		corr = corr1(map)/sqrt(power0(xx,yy)*power1)

; correlation with zero shift is given at corr(nxh,nyh)

		maxcorr = max( corr, maxloc)
		xmax = fix (maxloc mod long(nx1))
		ymax = fix (maxloc / long(nx1) )
		xshiftr = 0
		yshiftr = 0
		if (xmax-w lt 0 or xmax+w ge nx1 or $
			ymax-w lt 0 or ymax+w ge ny1 or $
			abs(xmax-nxh) gt nxh or $
			abs(ymax-nyh) gt nyh ) then begin
			if (monitor ge 3) then $
				print,'too large shift, xx,yy,xshift,yshift=', $
				xx,yy,xmax-nxh,ymax-nyh
		endif else if (maxcorr lt corrmin) then begin
			if (monitor ge 3) then $
				print,'low correlation, xx,yy,corr=', $
				xx,yy,maxcorr
		endif else begin
			peak_2d, corr((xmax-w):(xmax+w), (ymax-w):(ymax+w)), $
				xpeak,ypeak

; peak_2d returns the location of the peak of correlation (xpeak,ypeak)
; as pixel number counted from (xmax-w,ymax-w)

			xshiftr = xpeak + ( xmax - w ) -nxh
			yshiftr = ypeak + ( ymax - w ) -nyh

; (xshiftr,yshiftr) is the shift of image1 with respect to refimage
;  if xshiftr>0, image1 is shifted right with respect to refimage
;  shift(image1,-xshiftr,-yshiftr) will be close to refimage

		endelse

		if (monitor ge 3) then begin
			contour, corr, indgen(nx1)-nxh, indgen(ny1)-nyh
			oplot, [-nxh,nxh-1],[0,0]
			oplot, [0,0],[-nyh,nyh-1]
			oplot, [xshiftr],[yshiftr],psym=1
			tvscl, image1(xs:xe, ys:ye), order=0
			print, 'xx,yy,delx,dely=',xx,yy,delx,dely
			print, 'correlation, xshiftr, yshiftr=', $
				maxcorr,xshiftr,yshiftr
	print,' type "r" for non-stop mode, "s" to suspend, CR for next sub-image'
			read, yn
			if (yn eq 'r') then monitor =2 $
			else if (yn eq 's') then stop
		endif

		xci(xx,yy) = xco(xx,yy) + xshiftr
		yci(xx,yy) = yco(xx,yy) + yshiftr
		xshift(xx,yy)=idelx + xshiftr 
		yshift(xx,yy)=idely + yshiftr

	endfor

	if (monitor ge 1) then begin
		erase
		tvscl, refimage, x0,yu, /device, order=0
		xyouts,x0,yu-d,'reference image', $
			/device, size=charsize
		if (keyword_set(refinfo)) then $
			xyouts, x0, yu, refinfo, /device, size=charsize

		tvscl, image, x0+nx+d,yu, /device, order=0
		xyouts,x0+nx+d,yu-d,'sample image', $
			/device, size=charsize
		if (keyword_set(info)) then $
			xyouts, x0+nx+d, yu, info, /device, size=charsize

		tv, bytscl(rshift(image,-delx,-dely) - refimage, $
			min=-diflevel*maxref, max=diflevel*maxref), $
			x0+nx+d+nx+d, yu, /device, order=0
		xyouts,x0+nx+d+nx+d,yu-d,'shifted sample - ref', $
			/device, size=charsize

		tvscl, image1, x0,yd, /device, order=0
		xyouts,x0,yd-d,'shifted sample image', $
			/device, size=charsize
		xyouts, x0, d/2, 'shift vectors magnified by ' + $
			strtrim(string(magnify),1) + $
			'     order of polynomical warping=' + $
			strtrim(string(norder),1) + $
			'     diff. normalized to ' + $
			strtrim(string(diflevel,format='(f6.3)'),1) + $
			' of max.', $
			/device, size=charsize
; plot shift vectors which would bring refimage to sample image
; namely (xshift,yshift) represents the image motion
		xrange=[x0,x0+nx-1]
		yrange=[yd,yd+ny-1]
		plotarrows, xshift,yshift, xrange,yrange, nx1,ny1, mag=magnify
	endif

	if (norder eq 0) then begin
; uniform lateral shift
		px0=total(xshift)/(float(nxdiv)*float(nydiv))
		py0=total(yshift)/(float(nxdiv)*float(nydiv))
		xshiftr0=px0-idelx
		yshiftr0=py0-idely
		warpedimage=rshift(image1,-xshiftr0,-yshiftr0)
		for z=0, nz-1 do $
			samewarp(*,*,z)=rshift(samewarp(*,*,z), -px0,-py0)
		print,'shift of sample with respect to reference=', px0,py0
		px = [[px0, 0.0], [1.0, 0.0]]
		py = [[py0, 1.0], [0.0, 0.0]]	; these px,py can be used by poly_2d

	endif else if (norder gt 0) then begin
; polywarp over the whole area
		polywarp,xci,yci,xco,yco,norder,px,py
		warpedimage = poly_2d( image1,px,py,interp,missing=0)
		px(0,0)=px(0,0)+idelx
		py(0,0)=py(0,0)+idely
		print,'px=',px(0,0),px(0,1),px(1,0)
		print,'py=',py(0,0),py(0,1),py(1,0)
		for z=0, nz-1 do $
			samewarp(*,*,z)= $
				poly_2d(samewarp(*,*,z),px,py,interp,missing=0)
		xvecw = -xco
		yvecw = -yco
		for j=0, norder do for i=0, norder do begin
			xvecw=xvecw + px(i,j)*float(xco)^j*float(yco)^i
			yvecw=yvecw + py(i,j)*float(xco)^j*float(yco)^i
		endfor
		xresidual = xshift - xvecw
		yresidual = yshift - yvecw

	endif else begin
; bilinear piecewise warp
		xi=xo	&	yi=yo
		xi(*,*)=-1	&	yi(*,*)=-1
		if (nxdiv lt 2 or nydiv lt 2) then begin
			print,'(proc.corralgn) error, ndiv lt 2'
			return
		endif
		for yy = 0,nydiv-2 do for xx =0, nxdiv-2 do begin
			xs = xh(xx)	&	ys = yh(yy)
			xe = xh(xx+1)-1
			ye = yh(yy+1)-1
			dx=(xo(xs:xe,ys:ye)-xco(xx,yy))/ $
				(xco(xx+1,yy)-xco(xx,yy))
			dy=(yo(xs:xe,ys:ye)-yco(xx,yy))/ $
				(yco(xx,yy+1)-yco(xx,yy))
			xi(xs:xe,ys:ye) = xci(xx,yy) $
			+ dx * (xci(xx+1,yy)-xci(xx,yy)) $
			+ dy * (xci(xx,yy+1)-xci(xx,yy)) $
			+ dx*dy*( xci(xx+1,yy+1)-xci(xx+1,yy) $
				-xci(xx,yy+1)+xci(xx,yy))
			yi(xs:xe,ys:ye) = yci(xx,yy) $
			+ dx * (yci(xx+1,yy)-yci(xx,yy)) $
			+ dy * (yci(xx,yy+1)-yci(xx,yy)) $
			+ dx*dy*( yci(xx+1,yy+1)-yci(xx+1,yy) $
				-yci(xx,yy+1)+yci(xx,yy))
		endfor
		loc = where( xi ge 0 and xi lt nx-1 and $
			 yi ge 0 and yi lt ny-1)
		dxi=xi-fix(xi)
		dyi=yi-fix(yi)
		warpedimage( xo(loc),yo(loc))= $
			image1( xi(loc),yi(loc)) + $
		(image1( xi(loc)+1,yi(loc)) -image1(xi(loc),yi(loc)))*dxi(loc) + $
		(image1( xi(loc),yi(loc)+1) -image1(xi(loc),yi(loc)))*dyi(loc) + $
		(image1( xi(loc)+1,yi(loc)+1) -image1(xi(loc)+1,yi(loc)) $
		-image1( xi(loc),yi(loc)+1) + image1(xi(loc),yi(loc)))*dxi(loc)*dyi(loc)

		for z=0, nz-1 do begin
			buf=shift(samewarp(*,*,z),-idelx,-idely)
			samewarp( xo(loc),yo(loc),z) = $
			buf( xi(loc),yi(loc)) + $
			(buf( xi(loc)+1,yi(loc)) -buf(xi(loc),yi(loc)))*dxi(loc) + $
			(buf( xi(loc),yi(loc)+1) -buf(xi(loc),yi(loc)))*dyi(loc) + $
			(buf( xi(loc)+1,yi(loc)+1) -buf(xi(loc)+1,yi(loc)) $
			-buf( xi(loc),yi(loc)+1) + buf(xi(loc),yi(loc)))*dxi(loc)*dyi(loc)
		endfor
	endelse
				
	if (monitor ge 1) then begin

		tvscl, warpedimage, x0+nx+d,yd, /device, order=0
		xyouts,x0+nx+d,yd-d,'warped sample image', $
			/device, size=charsize

		tv, bytscl(warpedimage-refimage, $
			min=-diflevel*maxref, max=diflevel*maxref), $
			x0+nx+d+nx+d,yd, /device, order=0
		xyouts,x0+nx+d+nx+d,yd-d,'warped sample - ref', $
			/device, size=charsize

		if (norder gt 0) then begin
                 	xrange=[x0+nx+d,x0+nx+d+nx-1]
                 	yrange=[yd,yd+ny-1]
			plotarrows, xresidual,yresidual, xrange,yrange, $
				nx1,ny1, mag=magnify
                 	xrange=[x0+nx+d+nx+d,x0+nx+d+nx+d+nx-1]
                 	yrange=[yd,yd+ny-1]
			plotarrows, xvecw,yvecw, xrange,yrange, $
				nx1,ny1, mag=magnify
		endif
	endif
	if (monitor ge 2) then begin
		print,'type "r" for non-stop mode,', $
		' "b" for display correlation, "s" to suspend'
		read, yn
		if (yn eq 'r') then monitor=1 $
		else if (yn eq 's') then stop $
		else if (yn eq 'b') then begin
			monitor=3
			goto, back
		endif
	endif

	return
	end