#@ String (label = "Time increment") increment_str
#@ String (label = "Time start") time_start_str
#@ String (label = "Timestamp format", choices={"Simple", "h:mm / m:ss", "h:mm:ss"}) time_format
#@ String (label = "Suffix after timestamp") timestamp_suffix
#@ Integer (label = "Font size") font_size
#@ String (label = "Location", choices={"Upper left", "Upper right", "Lower left", "Lower right", "ROI location"}) location


import os
import re

from ij import IJ, ImageStack, ImagePlus
from ij.gui import Overlay, Roi, TextRoi, WaitForUserDialog, NonBlockingGenericDialog

from java.awt import Font, Color

# To stop the IDE from complaining
increment_str = increment_str
time_start_str = time_start_str
time_format = time_format
timestamp_suffix = timestamp_suffix
font_size = font_size
location = location


def time_string_to_value(string):
	split_string = string.split(':')
	return sum([int(string) * 60 ** i for i, string in enumerate(reversed(split_string))])


def time_value_to_string(value):
	if time_format == "Simple":
		timestamp = str(value)
	elif time_format == "h:mm / m:ss":
		timestamp = str(value // 60) + ':' + str(value % 60).zfill(2)
	elif time_format == "h:mm:ss":
		timestamp = str(value // 3600) + ':' + str(value % 3600 // 60).zfill(2) + ':' + str(value % 60).zfill(2)
	return timestamp + " " + timestamp_suffix


def timestamp_to_width(timestamp):
	roi = TextRoi(0, 0, timestamp, font)
	return roi.getBounds().width

increment = time_string_to_value(increment_str)
time_start = time_string_to_value(time_start_str)

imp = IJ.getImage()
dimX, dimY, dimC, dimZ, dimT = imp.getDimensions()

if not imp.isHyperStack():
	mode = 'Stack'
	maxSlice = imp.getStackSize()
else:
	if dimT > 1:
		mode = 'T'
		maxSlice = dimT
	elif dimZ > 1:
		mode = 'Z'
		maxSlice = dimZ
	else:
		WaitForUserDialog("Error", "The active image is not a stack").show()
		exit(0)

overlay = imp.getOverlay()
if overlay is None:
	overlay = Overlay()
	imp.setOverlay(overlay)

last_time_value = time_start + increment * maxSlice
last_timestamp = time_value_to_string(last_time_value)

padding = 0.5 * font_size

if location == "ROI location":
	imp.deleteRoi()
	IJ.selectWindow(imp.getID())
	IJ.setTool("rectangle")
	WaitForUserDialog("Timestamp location", "Please draw a rectangle at timestamp location").show()
	location_roi = imp.getRoi()
	if location_roi is None:
		error_msg = "This option requires a ROI."
		WaitForUserDialog("Error", error_msg).show()
		raise Exception(error_msg)
	rect = location_roi.getBounds()
	x, y = rect.x, rect.y
	font = Font("Arial", Font.PLAIN, rect.height)

else:
	font = Font("Arial", Font.PLAIN, font_size)
	max_timestamp_width = timestamp_to_width(last_timestamp)
	if location.startswith("Upper"):
		y = padding
	elif location.startswith("Lower"):
		y = dimY - padding - font_size
	if location.endswith("left"):
		x = padding
	elif location.endswith("right"):
		x = dimX - padding - max_timestamp_width

for s in range(1, maxSlice + 1):
	time_value = time_start + increment * (s - 1)
	timestamp = time_value_to_string(time_value).zfill(len(last_timestamp))

	roi = TextRoi(x, y, timestamp, font)

	roi.setStrokeColor(Color.WHITE)
	if mode == 'T':
		roi.setPosition(0, 0, s)
	elif mode == 'Z':
		roi.setPosition(0, s, 0)
	elif mode == 'Stack':
		roi.setPosition(s)
	overlay.add(roi)