// Script to correct shear in atomic resolution images caused by drift
// D. R. G. Mitchell, adminnospam@dmscripting.com.au (remove the nospam to make this work)
// version:20200110, v1.1, August 2019, www.dmscripting.com

// This assumes that rows of atoms are horiontal and that what would otherwise be vertical columns
// of atoms are inclined, due to the effects of drift. Run the script to create an arrow (or add one manually
// with the line tool with ALT held down). Align the arrow so that it points downwards and is parallel to the
// rows of atoms you want to make vertical. 

// Run the script again and the image will be sheared to correct the inclination and make the atom columns
// vertical. 

// This is a very crude script since it only shears in the horizontal direction. Also, it applies a uniform
// shear gradient - and distortion due to drift many not be constant.


// Ensure at least one image is displayed

number nodocs=countdocumentwindowsoftype(5)
if(nodocs<1)
	{
		showalert("Please ensure an image for deshearing is displayed>",2)
		exit(0)
	}


// Source the front image and some info

image front:=getfrontimage()
number xsize, ysize
getsize(front, xsize, ysize)
imagedisplay imgdisp=front.imagegetimagedisplay(0)

if(ysize<2)
	{
		showalert("This script does not work on spectra and profiles.", 2)
		exit(0)
	}


// Check to see if an arrow is present. 

number noarrows=imgdisp.componentcountchildrenoftype(3)


// If no arrows are present one is added with a suitable prompt

if(noarrows<1)
	{
		// Create some position coordinates for the arrow
		
		number top=ysize/10
		number left=xsize/10
		number bottom=ysize-top
		number right=left
		
		
		// Add the arrow
		
		component newarrow=newarrowannotation(top, left, bottom, right)
		newarrow.componentsetforegroundcolor(1,0,0)
		imgdisp.componentaddchildatend(newarrow)
		noarrows=imgdisp.componentcountchildrenoftype(3)


		// Add a text prompt
		
		number fontsize=round(ysize/20)
		component label=newtextannotation(left+10, top+10, "Align arrow with\natomic columns\n - rerun this script.", fontsize)
		label.componentsetfontfaceName("Microsoft Sans serif")
		label.componentsetforegroundcolor(1,0,0)
		label.componentsetbackgroundcolor(1,1,1)
		
		label.componentsetdrawingmode(1)
		label.componentsetfillmode(1)
		imgdisp.componentaddchildatend(label)
		
		exit(0)
	}

// if an arrow is present, apply the shear


// Source the arrow and its coordinates

component thisarrow=imgdisp.componentgetnthchildoftype(3,0)
number arrowstartx, arrowstarty, arrowendx, arrowendy
thisarrow.componentgetcontrolpoint(4, arrowstartx, arrowstarty)
thisarrow.componentgetcontrolpoint(7, arrowendx, arrowendy)


// Calculate the angle the arrow makes. Vertical is zero, pointing to lower right is +ve
// pointing to lower left is -ve

number deltax=arrowendx-arrowstartx
number deltay=arrowendy-arrowstarty
number angle=atan(deltax/deltay)


// create a padded image. This is padded with one empty image on either side. This ensure that the 
// shearing can occur without the content wrapping aroud ie image details pushed off one side appearing 
// on the other side of the image. The image content is in the central third, with empty left and right thirds

image paddedfront=realimage("", 4, xsize*3, ysize)
paddedfront[0, xsize, ysize, (2*xsize)]=front


// Calculate the shear to apply

number ypixelshift=1
number xpixelshift=atan(angle)*ypixelshift


// Create a clone of the padded image to store the sheared results

image outputimage=imageclone(paddedfront)*0


// Main processing loop - shear each line by the appropriate amount

number i
for(i=0; i<ysize; i++)
	{
		number shift=i*xpixelshift
		outputimage[i, 0, i+1, (3*xsize)]=offset(paddedfront[i, 0, i+1, (3*xsize)], shift,0)
	}

// Remember the maximum shear applied

number finalshift=(ysize-1)*xpixelshift


// Prompt to include or exclude black space resulting from shearing

number cropimage

if(twobuttondialog("Exclude black space from shearing?", "Exclude", "Include")) cropimage=1
else cropimage=0


// Calculate the shifted x coordinates of the four corners of the sheared image

number midpointx=(3*xsize)/2
number topleft=midpointx-(xsize/2)
number topright=midpointx+(xsize/2)
number bottomleft=topleft-finalshift
number bottomright=topright-finalshift


image finalimage 

if(cropimage==0)
	{
		// Include the black space from shearing. The image is extracted out to it widest 
		// bounds - but black space from shearing is left in.
		
		number minimumleft=minimum(topleft, bottomleft)
		number maximumright=maximum(topright, bottomright)
		finalimage=outputimage[0,minimumleft, ysize, maximumright]
	}
else
	{
		// The image is extracted out to it narrowest 
		// bounds - so that black space from shearing is excluded.
		
		number maximumleft=maximum(topleft, bottomleft)
		number minimumright=minimum(topright, bottomright)
		finalimage=outputimage[0,maximumleft, ysize, minimumright]
	}


// Copy tags and calibration

imagecopycalibrationfrom(finalimage, front)
taggroup fronttags=front.imagegettaggroup()
taggroup finaltags=finalimage.imagegettaggroup()
taggroupcopytagsfrom(finaltags, fronttags)


// Display the front image

showimage(front)
documentwindow frontwin=getdocumentwindow(0)
number top, left
frontwin.windowgetframeposition(left, top)


// Display the final image

string imgname=getname(front)
setname(finalimage, imgname+" (Sheared)")
showimage(finalimage)
documentwindow finalwin=getdocumentwindow(0)
finalwin.windowsetframeposition(left+30, top+30)

// Remove the arrows and label from the front image

try
	{
		component thisarrow=imgdisp.componentgetnthchildoftype(3,0)
		thisarrow.componentremovefromparent()
		component thislabel=imgdisp.componentgetnthchildoftype(13,0)
		thislabel.componentremovefromparent()
	}
catch break