# == title
# The AnnotationFunction class
#
# == details
# The heatmap annotation is basically graphics aligned to the heatmap columns
# if it is column annotation or heatmap rows if it is row annotation, while
# there is no restrictions for the graphic types, e.g. it can be heatmap-like
# annotation or points. Here the AnnotationFunction class is designed for
# creating complex and flexible annotation graphics. As the main part of the class, it uses
# a user-defined function to define the graphics. It also keeps information of
# the size of the plotting regions of the annotation. And most importantly, it
# allows subsetting of the annotation to draw a subset of the graphics, which
# is the base for the splitting of the annotations.
#
# See `AnnotationFunction` constructor for details.
#
AnnotationFunction = setClass("AnnotationFunction",
	slots = list(
		which = "character",
		fun_name = "character",
		width = "ANY",
		height = "ANY",
		n = "ANY",
		var_env = "environment",
		fun = "function",
		subset_rule = "list",
		subsetable = "logical",
		data_scale = "numeric",
		extended = "ANY",
		show_name = "logical"
	),
	prototype = list(
		fun_name = "",
		width = unit(1, "npc"),
		height = unit(1, "npc"),
		subset_rule = list(),
		subsetable = FALSE,
		data_scale = c(0, 1),
		n = NA_integer_,
		extended = unit(c(0, 0, 0, 0), "mm"),
		show_name = TRUE
	)
)


anno_width_and_height = function(which, width = NULL, height = NULL, 
	default = unit(10, "mm")) {

	if(which == "column") {
		if(is.null(height)) {
			height = default
		} else {
			if(!is_abs_unit(height)) {
				stop_wrap("height of the annotation can only be an absolute unit.")
			} else {
				height = convertHeight(height, "mm")
			}
		}
		if(is.null(width)) {
			width = unit(1, "npc")
		}
	}
	if(which == "row") {
		if(is.null(width)) {
			width = default
		} else {
			if(!is_abs_unit(width)) {
				stop_wrap("width of the annotation can only be an absolute unit.")
			} else {
				width = convertWidth(width, "mm")
			}
		}
		if(is.null(height)) {
			height = unit(1, "npc")
		}
	}
	return(list(height = height, width = width))
}


# == title
# Constructor of AnnotationFunction Class
#
# == param
# -fun A function which defines how to draw the annotation. See **Details** section.
# -fun_name The name of the function. It is only used for printing the object.
# -which Whether it is drawn as a column annotation or a row annotation?
# -var_import The names of the variables or the variable themselves that the annotation function depends on. See **Details** section.
# -n Number of observations in the annotation. It is not mandatory, but it is better to provide this information
#   so that the higher order `HeatmapAnnotation` knows it and it can perform check on the consistency of annotations and heatmaps.
# -data_scale The data scale on the data axis (y-axis for column annotation and x-axis for row annotation). It is only used
#             when `decorate_annotation` is used with "native" unit coordinates.
# -subset_rule The rule of subsetting variables in ``var_import``. It should be set when users want the final object to
#              be subsetable. See **Details** section.
# -subsetable Whether the object is subsetable?
# -show_name It is used to turn off the drawing of annotation names in `HeatmapAnnotation`. Annotations always have names
#        associated and normally they will be drawn beside the annotation graphics to tell what the annotation is about.
#        e.g. the annotation names put beside the points annotation graphics. However, for some of the annotations, the names
#        are not necessarily to be drawn, such as text annotations drawn by `anno_text` or an empty annotation drawn by `anno_empty`.
#        In this case, when ``show_names`` is set to ``FALSE``, there will be no annotation names drawn for the annotation.
# -width The width of the plotting region (the viewport) that the annotation is drawn. If it is a row annotation,
#        the width must be an absolute unit. Since the ``AnnotationFunction`` object is always contained by the `SingleAnnotation-class`object,
#        you can only set the width of row annotations or height of column annotations, while e.g. the height of the row annotation is always ``unit(1, "npc")``
#        which means it always fully filled in the parent ``SingleAnnotation`` and only in `SingleAnnotation` or even `HeatmapAnnotation`
#        can adjust the height of the row annotations.
# -height The height of the plotting region (the viewport) that the annotation is drawn. If it is a column annotation,
#        the width must be an absolute unit.
#
# == details
# We use a normal R function defines how to draw the annotation graphics. As
# expected, the main part of the AnnotationFunction class is this function.
# The function defines how to draw at specific positions which correspond to
# rows or columns in the heatmap. The function should have three arguments:
# ``index``, ``k`` and ``n`` (the names of the arguments can be arbitory)
# where ``k`` and ``n`` are optional. ``index`` corresponds to the indices of
# rows or columns of the heatmap. The value of ``index`` is not necessarily to
# be the whole row indices or column indices of the heatmap. It can be a
# subset of the indices if the annotation is split into slices according to
# the split of the heatmap. ``index`` is always reordered according to the
# reordering of heatmap rows or columns (e.g. by clustering). So, ``index``
# actually contains a list of row or column indices for the current slice
# after row or column reordering.
# 
# As mentioned, annotation can be split into slices. ``k`` corresponds to the
# current slice and ``n`` corresponds to the total number of slices. As you can image, 
# when ``n > 1``, the annotation function will be executed for all ``k``s. The
# information of ``k`` and ``n`` sometimes can be useful, for example, we want
# to add axis ot the right side of a column annotation, if this column annotation
# is split into several slices, the axis is only drawn when``k == n``.
#
# Since the function only allows ``index``, ``k`` and ``n``, the function
# sometimes uses several external variables which can not be defined inside
# the function, e.g. the data points for the annotation. These variables
# should be imported into the AnnotationFunction class so that the function
# can correctly find these variables (by ``var_import`` argument).
#
# One important feature for AnnotationFunction class is it can be subsetable.
# To allow subsetting of the object, users need to define the rules for the
# imported variables. The rules are simple function which
# accpets the variable and indices, and returns the subset of the variable.
# The subset rule functions implemented in this package are `subset_gp`,
# `subset_matrix_by_row` and `subset_vector`. These three functions are enough
# for most of the cases.
#
# In following, we defined three AnnotationFunction objects: 
#
# 1. It needs external variable and support subsetting
#
#	x = 1:10
#	anno1 = AnnotationFunction(
#		fun = function(index) {
#			n = length(index)
#			pushViewport(viewport())
#			grid.points(1:n, x[index])
#			popViewport()
#		},
#		var_imported = list(x = x),
#		n = 10,
#		subset_rule = list(x = subset_vector),
#		subsetable = TRUE
#	)
#
# 2. The data variable is defined inside the function and no need to import other variables.
#
#	anno2 = AnnotationFunction(
#		fun = function(index) {
#			x = 1:10
#			n = length(index)
#			pushViewport(viewport())
#			grid.points(1:n, x[index])
#			popViewport()
#		},
#		n = 10,
#		subsetable = TRUE
#	)
#
# 3. Only specify the function to the constructor. ``anno3`` is not subsettable.
#
#	anno3 = AnnotationFunction(
#		fun = function(index) {
#			x = 1:10
#			n = length(index)
#			pushViewport(viewport())
#			grid.points(1:n, x[index])
#			popViewport()
#		}
#	)
#
# As you can see from the examples, you need to push a viewport for graphics and finally pop the viewport.
#
# In the package, we have implemted quite a lot annotation function by `AnnotationFunction` constructor:
# `anno_empty`, `anno_image`, `anno_points`, `anno_lines`, `anno_barplot`, `anno_boxplot`, `anno_histogram`,
# `anno_density`, `anno_joyplot`, `anno_horizon`, `anno_text` and `anno_mark`. These built-in annotation functions
# support as both row annotations and column annotations and they are are all subsettable.
#
# == seealso
# The build-in annotation functions are already enough for most of the analysis, nevertheless, if users
# want to know more about how to construct the AnnotationFunction class manually, they can refer to
# ComplexHeatmap Complete Reference ().
AnnotationFunction = function(fun, fun_name = "", which = c("column", "row"), 
	var_import = list(), n = NA, data_scale = c(0, 1), subset_rule = list(), 
	subsetable = length(subset_rule) > 0, show_name = TRUE, width = NULL, height = NULL) {

	which = match.arg(which)[1]

	verbose = ht_opt$verbose
	
	anno = new("AnnotationFunction")

	anno@which = which
	anno@fun_name = fun_name

	if(verbose) qqcat("construct AnnotationFunction with '@{fun_name}()'\n")

	anno_size = anno_width_and_height(which, width, height, unit(1, "cm"))
	anno@width = anno_size$width
	anno@height = anno_size$height

	anno@show_name = show_name

	anno@n = n
	anno@data_scale = data_scale

	if(length(var_import)) {
		anno@var_env = new.env()
		if(is.character(var_import)) {
			for(nm in var_import) {
				anno@var_env[[nm]] = get(nm, envir = parent.frame())
			}
		} else if(inherits(var_import, "list")) {
			if(is.null(names(var_import))) {
				var_imported_nm = sapply(as.list(substitute(var_import))[-1], as.character)
				names(var_import) = var_imported_nm
			}

			for(nm in names(var_import)) {
				anno@var_env[[nm]] = var_import[[nm]]
			}
		} else {
			stop_wrap("`var_import` needs to be a character vector which contains variable names or a list of variables.")
		}
		environment(fun) = anno@var_env
	} else {
		anno@var_env = environment(fun)
	}

	if(length(as.list(formals(fun))) == 1) {
		formals(fun) = alist(index = , k = 1, n = 1)
	}
	anno@fun = fun
	
	if(is.null(subset_rule)) {
		for(nm in names(anno@var_env)) {
			if(is.matrix(anno@var_env[[nm]])) {
				anno@subset_rule[[nm]] = subset_matrix_by_row
			} else if(inherits(anno@var_env[[nm]], "gpar")) {
				anno@subset_rule[[nm]] = subset_gp
			} else if(is.vector(anno@var_env[[nm]])) {
				if(length(anno@var_env[[nm]]) > 1) {
					anno@subset_rule[[nm]] = subset_vector
				}
			}
		}
	} else {
		for(nm in names(subset_rule)) {
			anno@subset_rule[[nm]] = subset_rule[[nm]]
		}
	}

	if(missing(subsetable)) {
		# is user defined subset rule
		if(length(anno@subset_rule)) {
			anno@subsetable = TRUE
		}
	} else {
		anno@subsetable = subsetable
	}

	return(anno)
}

# == title
# Subset an AnnotationFunction Object
#
# == param
# -x An `AnnotationFunction-class` object.
# -i A vector of indices.
#
# == details
# One good thing for designing the `AnnotationFunction-class` is it can be subsetted,
# and this is the base for the splitting of the annotations.
#
# == example
# anno = anno_simple(1:10)
# anno[1:5]
# draw(anno[1:5], test = "subset of column annotation")
"[.AnnotationFunction" = function(x, i) {
	if(nargs() == 1) {
		return(x)
	} else {
		if(!x@subsetable) {
			stop_wrap("This object is not subsetable.")
		}
		x = copy_all(x)
		for(var in names(x@subset_rule)) {
			oe = try(x@var_env[[var]] <- x@subset_rule[[var]](x@var_env[[var]], i), silent = TRUE)
			if(inherits(oe, "try-error")) {
				message(paste0("An error when subsetting ", var))
				stop_wrap(oe)
			}
		}
		if(is.logical(i)) {
			x@n = sum(i)
		} else {
			x@n = length(i)
		}
		return(x)
	}
}

# == title
# Draw the AnnotationFunction Object
#
# == param
# -object The `AnnotationFunction-class` object.
# -index Index of observations.
# -k Current index of slice.
# -n Total number of slices.
# -test Is it in test mode? The value can be logical or a text which is plotted as the title of plot.
#
# == detail
# Normally it is called internally by the `SingleAnnotation` class.
#
# When ``test`` is set to ``TRUE``, the annotation graphic is directly drawn,
# which is generally for testing purpose.
# 
setMethod(f = "draw",
	signature = "AnnotationFunction",
	definition = function(object, index, k = 1, n = 1, test = FALSE) {
		
	if(is.character(test)) {
		test2 = TRUE
	} else {
		test2 = test
		test = ""
	}
	if(test2) {
        grid.newpage()
        pushViewport(viewport(width = 0.8, height = 0.8))
        if(is.na(object@n)) {
        	object@n = 1
        }
    }

    verbose = ht_opt$verbose
    if(verbose) qqcat("draw annotation generated by @{object@fun_name}\n")

    if(missing(index)) index = seq_len(object@n)

    anno_height = object@height
    anno_width = object@width

    # names should be passed to the data viewport
	pushViewport(viewport(width = anno_width, height = anno_height))
	vp_name1 = current.viewport()$name
	object@fun(index, k, n)
	if(test2) {
		grid.text(test, y = unit(1, "npc") + unit(2, "mm"), just = "bottom")

		if(!identical(unit(0, "mm"), object@extended[1])) {
			grid.rect(y = 1, height = unit(1, "npc") + object@extended[1], just = "top",
				gp = gpar(fill = "transparent", col = "red", lty = 2))
		} else if(!identical(unit(0, "mm"), object@extended[[2]])) {
			grid.rect(x = 1, width = unit(1, "npc") + object@extended[2], just = "right",
				gp = gpar(fill = "transparent", col = "red", lty = 2))
		} else if(!identical(unit(0, "mm"), object@extended[[3]])) {
			grid.rect(y = 0, height = unit(1, "npc") + object@extended[3], just = "bottom",
				gp = gpar(fill = "transparent", col = "red", lty = 2))
		} else if(!identical(unit(0, "mm"), object@extended[[4]])) {
			grid.rect(x = 0, width = unit(1, "npc") + object@extended[4], just = "left",
				gp = gpar(fill = "transparent", col = "red", lty = 2))
		}
		
	}
	vp_name2 = current.viewport()$name
	if(vp_name1 != vp_name2) {
		stop_wrap("Viewports should be the same before and after plotting the annotation graphics.")
	}
	popViewport()

	if(test2) {
		popViewport()
	}
	
})

# == title
# Copy the AnnotationFunction Object
#
# == param
# -object The `AnnotationFunction-class` object.
#
# == detail 
# In `AnnotationFunction-class`, there is an environment which
# stores some external variables for the annotation function (specified by the
# ``var_import`` argument by constructing the `AnnotationFunction-class`
# object. This `copy_all,AnnotationFunction-method` hard copies all the
# variables into a new isolated environment.
#
# The environment is at ``object@var_env``.
#
setMethod(f = "copy_all",
	signature = "AnnotationFunction",
	definition = function(object) {
		object2 = object
		object2@var_env = new.env()
		for(var in names(object@var_env)) {
			object2@var_env[[var]] = object@var_env[[var]]
		}
		environment(object2@fun) = object2@var_env
		return(object2)
})

# == title
# Print the AnnotationFunction Object
#
# == param
# -object The `AnnotationFunction-class` object.
#
setMethod(f = "show",
	signature = "AnnotationFunction",
	definition = function(object) {

	cat("An AnnotationFunction object\n")
	if(object@fun_name == "") {
		cat("  function: user-defined\n")
	} else {
		cat("  function: ", object@fun_name, "()\n", sep = "")
	}
	cat("  position:", object@which, "\n")
	cat("  items:", ifelse(object@n == 0, "unknown", object@n), "\n")
	cat("  width:", as.character(object@width), "\n")
	cat("  height:", as.character(object@height), "\n")
	var_imported = names(object@var_env)
	if(length(var_imported)) {
		cat("  imported variable:", paste(var_imported, collapse = ", "), "\n")
		var_subsetable = names(object@subset_rule)
		if(length(var_subsetable)) {
			cat("  subsetable variable:", paste(var_subsetable, collapse = ", "), "\n")
		}
	}
	cat("  this object is", ifelse(object@subsetable, "\b", "not"), "subsetable\n")
	dirt = c("bottom", "left", "top", "right")
	for(i in 1:4) {
		if(!identical(unit(0, "mm"), object@extended[i])) {
			cat(" ", as.character(object@extended[i]), "extension on the", dirt[i], "\n")
		}
	}
	
})

# == title
# Number of Observations
#
# == param
# -object The `AnnotationFunction-class` object.
# -... Other arguments.
#
# == details It returns the ``n`` slot in the object. If it does not exist, it
# returns ``NA``.
#
# == example
# anno = anno_points(1:10)
# nobs(anno)
nobs.AnnotationFunction = function(object, ...) {
	if(is.na(object@n)) {
		return(NA)
	} else if(object@n > 0) {
		object@n
	} else {
		NA
	}
}