# script for comparing astro images (find moving objects)
# version 1.42
# Ricardo Serpell / Sep 2021
#
################################## Parameters ##################################
seq_name = "r_F22_00"		# sequence root filename
seq_ext ="fits"				# sequence file extension
border = 100 				# border in pixels to be discarded from the input
num_diff = 200 				# this many largest differences will be checked
max_dist = 24 				# longest distance (pixels) to consider possible find

########################### Additional Parameters ##############################
bkg_res = 16 				# resolution of background extraction
b_rad = 10 					# match-avoid radius around already identified diffs

################################################################################

import numpy as np
from astropy.io import fits
from skimage.transform import rescale, resize
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import math

candidates = 0 
max_dist = int(max_dist / 2)
b_rad = int(b_rad / 2)

# performs power stretching of image
def stretch (im, mn = 0, mx = 1, exp = 1):
	rng = mx-mn
	masklow = im < mn
	maskhigh = im > mx
	im[masklow] = mn
	im[maskhigh] = mx
	output = ((im-mn)/rng)**exp
	return output


# opens fit image, crops border (defined at the begining)
# returns a float32 array with image data in range 0->1 
def openimage (filename, verbose = False):
	hdul = fits.open(filename)
	im1 = hdul[0].data
	date_obs = hdul[0].header['DATE-OBS'] # YYYY-MM-DDThh:mm:ss observation start
	exp_time = hdul[0].header['EXPTIME']  # in seconds
	if verbose:
		print(f"image : {filename}")
		hdul.info()
		print()
		print(repr(hdul[0].header))
		print()
		print(f"data type: {im1.dtype.name}")
	hdul.close()
	div = 255
	if im1.dtype.name=="uint16":
		div = 65536
	if im1.dtype.name=="uint32":
		div = 4294967296
	im10 = np.zeros([np.shape(im1)[0]-border*2,np.shape(im1)[1]-border*2], dtype = np.float32)
	im10 = im1[border:-border,border:-border]/div
	return im10, date_obs, exp_time


# subtracts background from image (gradients)
# returns subtracted image stretched from 0 to 1
def calibrate (image, res = 16, sigmas = 0):
	ii = image.shape[0]
	jj = image.shape[1]
	med = np.median(image)
	dev = np.std(image)
	maxi = np.amax(image)
	i_step = int(ii/res)
	j_step = int(jj/res)
	bkg = np.zeros([res+2, res+2], dtype = np.float32)
	for i in range(res):
		for j in range(res):
			chunk = np.copy(image[i*i_step:(i+1)*i_step, j*j_step:(j+1)*j_step])
			chunk[chunk>maxi/4] = med
			val = np.median(chunk)
			bkg[i+1,j+1] = val
			if i == 0:
				bkg[0,j+1] = val
				if j == 0:
					bkg[0,0] = val
				elif j == res-1:
					bkg[0,j+2] = val
			elif i == res-1:
				bkg[i+2,j+1] = val
				if j == 0:
					bkg[i+2,0] = val
				elif j == res-1:
					bkg[i+2,j+2] = val
			if j == 0:
				bkg[i+1,0] = val
			elif j == res-1:
				bkg[i+1,j+2] = val
	background = resize(bkg, (ii+i_step*2, jj+j_step*2))[i_step:-i_step,j_step:-j_step] + dev*sigmas
	output = image - background
	output[output<0] = 0
	return output/np.amax(image)


# finds the largest differences recorded in each subtraction
# a square is masked around each "find", size defined by 2*mask
def finddiff (subimage, n, mask = 20):
	ii, jj = subimage.shape
	top_n = np.zeros([n,2], dtype = np.int16)
	for x in range(n):
		ind = np.unravel_index(np.argmax(subimage, axis=None), subimage.shape)
		top_n[x, :] = ind[:]
		i, j = ind
		if i < mask:
			i = mask
		elif i > (ii-mask):
			i = ii-mask
		if j < mask:
			j = mask
		elif j > (jj-mask):
			j = jj-mask	
		subimage[i-mask:i+mask,j-mask:j+mask] = -1
	return top_n


# finds differences ocurring in the right order and roughly linear path
# i, j : positions of point in first image (sub1)
# maxdist is the maximum distance in pixels for ocurrence of object in succesive images
def findline (i, j, ts2, ts3, maxdist = 30):
	global candidates
	finds = []
	mdist2 = maxdist**2
	for p2 in range(ts2.shape[0]):
		dist = (i-ts2[p2,0])**2+(j-ts2[p2,1])**2
		if dist == 0:
			continue
		if dist <= mdist2:
			ndist = math.sqrt(dist)
			print()
			print(f"Found pair {round(ndist*2,2)} pixels appart. Position in image 1: i = {i*2}   j = {j*2}")
			di = (ts2[p2,0]-i)
			dj = (ts2[p2,1]-j)
			# print(f"di = {di}   dj = {dj}")
			dec = int(ndist/2)
			if dec == 0:
				dec = 1
			min_i = int(i+di*1.4 - dec*np.sign(di))
			max_i = i+di*2 + int(ndist)*np.sign(di)
			min_j = int(j+dj*1.4 - dec*np.sign(dj))
			max_j = j+dj*2 + int(ndist)*np.sign(dj)
			if min_i > max_i:
				val = min_i
				min_i = max_i
				max_i = val
			if min_j > max_j:
				val = min_j
				min_j = max_j
				max_j = val	
			for p3 in range(ts3.shape[0]):
				if min_i <= ts3[p3,0] <= max_i:
					if min_j <= ts3[p3,1] <= max_j:
						print()
						print("Found one candidate:")
						print(f"In image 1: i = {i*2}   j = {j*2}")
						print(f"In image 2: i = {ts2[p2,0]*2}   j = {ts2[p2,1]*2}")
						print(f"In image 3: i = {ts3[p3,0]*2}   j = {ts3[p3,1]*2}")
						print()
						finds.append([ts2[p2,0], ts2[p2,1]])
	return finds


# animates the results
def animate (c_image):
	fig = plt.figure()
	fig.tight_layout()
	# ax = fig.add_subplot(111)
	ims = []
	for i in range(3):
		ttl = plt.text(0.5, 2, "UTC D: "+utcdate[i]+"  T: "+utctime[i], color = "white")
		ims.append([plt.imshow(stretch(c_image[:,:,i], 0, 1, 0.33), cmap = "gray"), ttl])
	im_ani = animation.ArtistAnimation(fig, ims, interval = 700, repeat = True, blit = True)
	plt.draw()
	return im_ani


############################### main code ##################################

# open images and calibrate a downscaled (0.5x) version of them
print("Opening image files...")
datesstr = [0, 0, 0]
exposures = [0, 0, 0]

im01, datesstr[0], exposures[0] = openimage(seq_name+"1."+seq_ext, verbose = False)
im02, datesstr[1], exposures[1] = openimage(seq_name+"2."+seq_ext, verbose = False)
im03, datesstr[2], exposures[2] = openimage(seq_name+"3."+seq_ext, verbose = False)

utcdate = [0, 0, 0]
utctime = [0, 0, 0]
for i in range(3):
	h = int(datesstr[i][11:13])
	m = int(datesstr[i][14:16])
	s = float(datesstr[i][17:])+exposures[i]/2
	m = m + math.floor(s/60)
	if m > 59:
		m = m-60
		h = h+1
		if h > 23:
			h = h-24
	s = (s/60 - math.floor(s/60))*60
	utcdate[i] = datesstr[i][0:10]
	utctime[i] = "{0:02.0f}".format(h)+":{0:02.0f}".format(m)+":{0:06.3f}".format(s)

print("Calibrating and downsampling (0.5x) images...")
im11 = calibrate(rescale(im01, 0.5, anti_aliasing=False), bkg_res, 0.25)
im12 = calibrate(rescale(im02, 0.5, anti_aliasing=False), bkg_res, 0.25)
im13 = calibrate(rescale(im03, 0.5, anti_aliasing=False), bkg_res, 0.25)

ii, jj = im11.shape

# do the subtractions
print("Subtracting images...") 
sub1 = im11-im12-im13
sub1[sub1<0] = 0
sub2 = im12-im11-im13
sub2[sub2<0] = 0
sub3 = im13-im11-im12
sub3[sub3<0] = 0

# fig3, axis = plt.subplots(3)
# axis[0].imshow(stretch(sub1, 0, np.amax(sub1), 0.2))
# axis[1].imshow(stretch(sub2, 0, np.amax(sub1), 0.2))
# axis[2].imshow(stretch(sub3, 0, np.amax(sub1), 0.2))


print("Maximum differences between images (1-2-3, 2-1-3, 3-1-2):")
print(np.amax(sub1))
print(np.amax(sub2))
print(np.amax(sub3))

print("Compiling list of differences...") 
top_sub1 = finddiff(sub1, num_diff, b_rad)
top_sub2 = finddiff(sub2, num_diff, b_rad)
top_sub3 = finddiff(sub3, num_diff, b_rad)

# print(top_sub1)
# print(top_sub2)
# print(top_sub3)

# prepare image to show results (this consumes a lot of memory)
imc = np.zeros([ii*2,jj*2,3])
med = np.median(im01)
imc[:,:,0] = stretch(np.copy(im01), med, 0.5, 0.2)
med = np.median(im02)
imc[:,:,1] = stretch(np.copy(im02), med, 0.5, 0.25)
med = np.median(im03)
imc[:,:,2] = stretch(np.copy(im03), med, 0.5, 0.2)

# find aligned differences... (i.e. points fairly aligned across images)
print("Searching aligned differences...") 
for p1 in range(top_sub1.shape[0]):
	finds = findline (top_sub1[p1,0], top_sub1[p1,1], top_sub2, top_sub3, max_dist)
	if len(finds)>0:
		for n in range(len(finds)):
			candidates += 1
			i1 = finds[n][0]-max_dist
			j1 = finds[n][1]-max_dist
			i2 = finds[n][0]+max_dist
			j2 = finds[n][1]+max_dist
			if i1 < 0:
				i1 = 0
				i2 = max_dist*2
			elif i2 > ii:
				i1 = ii-(max_dist*2)
				i2 = ii
			if j1 < 0:
				j1 = 0
				j2 = max_dist*2
			elif j2 > jj:
				j1 = jj-(max_dist*2)
				j2 = jj	
			result = np.zeros([max_dist*4, max_dist*4, 3], dtype = np.float32)
			result[:,:,0] = np.copy(im01[i1*2:i2*2,j1*2:j2*2])
			result[:,:,1] = np.copy(im02[i1*2:i2*2,j1*2:j2*2])
			result[:,:,2] = np.copy(im03[i1*2:i2*2,j1*2:j2*2])
			maxi = max([np.amax(result[:,:,0]),np.amax(result[:,:,1]),np.amax(result[:,:,2])])
			result[:,:,0] = result[:,:,0]/maxi
			result[:,:,1] = result[:,:,1]/maxi
			result[:,:,2] = result[:,:,2]/maxi

			imc[i1*2:i2*2,j1*2:j1*2+3,0:2] = 1
			imc[i1*2:i2*2,j2*2-3:j2*2,0:2] = 1
			imc[i1*2:i1*2+3,j1*2:j2*2,0:2] = 1
			imc[i2*2-3:i2*2,j1*2:j2*2,0:2] = 1

			ani = animate(result)
			ani.save(f"{seq_name}_result_0{candidates}.gif", writer="Imagemagick")
			plt.ion()
			plt.plot()
			plt.pause(0.1)
			plt.ioff()

print()
if candidates > 0:
	print(f"Found {candidates} candidate(s) in total")
else:
	print("No moving objects found")

# display location of differences (image 1 = red, image 2 = green, image 3 = blue)
fig2, (ax2) = plt.subplots(1)
fig2.tight_layout()
imm2 = ax2.imshow(imc)

plt.show()