#####################################
# class and methods to map values to colors
#

# == title
# Class for Color Mapping
#
# == details
# The `ColorMapping-class` handles color mapping for discrete values and continuous values.
# Discrete values are mapped by setting a vector of colors and continuous values are mapped by setting
# a color mapping function.
#
# == methods
# The `ColorMapping-class` provides following methods:
#
# - `ColorMapping`: contructor methods.
# - `map_to_colors,ColorMapping-method`: mapping values to colors.
# - `color_mapping_legend,ColorMapping-method`: draw legend or get legend as an object.
#
# == author
# Zuguang Gu <z.gu@dkfz.de>
#
ColorMapping = setClass("ColorMapping",
	slots = list(
		colors  = "character", # a list of colors
		levels  = "ANY", # levels which colors correspond to
		col_fun = "function", # function to map values to colors
		type    = "character",  # continuous or discrete
		name    = "character",  # used to map to the dataset and taken as the title of the legend
		na_col  = "character",
		full_col  = "character"
	)
)

# == title
# Constructor Method for ColorMapping Class
#
# == param
# -name Name for this color mapping. The name is automatically generated if it is not specified.
# -colors Discrete colors.
# -levels Levels that correspond to ``colors``. If ``colors`` is name indexed,
#         ``levels`` can be ignored.
# -col_fun Color mapping function that maps continuous values to colors.
# -breaks Breaks for the continuous color mapping. If ``col_fun`` is
#         generated by `circlize::colorRamp2`, ``breaks`` is automatically inferred from the color mapping function.
# -na_col Colors for ``NA`` values.
# -full_col A super set of ``colors``, used internally.
#
# == detail
# ``colors`` and ``levels`` are used for discrete color mapping, ``col_fun`` and
# ``breaks`` are used for continuous color mapping.
#
# == value
# A `ColorMapping-class` object.
#
# == author
# Zuguang Gu <z.gu@dkfz.de>
#
# == examples
# cm = ColorMapping(colors = c("A" = "red", "B" = "black"))
# cm
# require(circlize)
# col_fun = colorRamp2(c(0, 1), c("white", "red"))
# cm = ColorMapping(col_fun = col_fun)
ColorMapping = function(name, colors = NULL, levels = NULL,
	col_fun = NULL, breaks = NULL, na_col = "#FFFFFF", full_col = NULL) {

	.Object = new("ColorMapping")

	if(missing(name)) {
		increase_color_mapping_index()
		name = paste0("color_mapping_", get_color_mapping_index())
	}
	if(!is.null(colors)) {
		if(is.null(levels)) {
			if(is.null(names(colors))) {
				stop_wrap("either provide `levels` or provide named `colors`.\n")
			}
			levels = names(colors)
		}
		if(length(colors) != length(levels)) {
			stop_wrap("length of colors and length of levels should be the same.\n")
		}
		colors = t(col2rgb(colors, alpha = TRUE))
		colors = rgb(colors[, 1:3, drop = FALSE], alpha = colors[, 4], maxColorValue = 255)
		.Object@colors = colors
		if(is.numeric(levels)) {
			.Object@levels = as.character(levels)
			#attr(.Object@levels, "numeric") = TRUE
		} else {
			.Object@levels = levels
		}
		names(.Object@colors) = levels

		if(!is.null(full_col)) {
			full_col_nm = names(full_col)
			full_col = t(col2rgb(full_col, alpha = TRUE))
			full_col = rgb(full_col[, 1:3, drop = FALSE], alpha = full_col[, 4], maxColorValue = 255)
			.Object@full_col = structure(full_col, names = full_col_nm)
		}
		.Object@type = "discrete"
	} else if(!is.null(col_fun)) {
		if(is.null(breaks)) {
			breaks = attr(col_fun, "breaks")
			if(is.null(breaks)) {
				stop_wrap("You should provide breaks.\n")
			}

			rg = range(breaks)
			diff = rg[2] - rg[1]
			rg[1] = rg[1] + diff*0.05
			rg[2] = rg[2] - diff*0.05

			le = pretty(rg, n = 3)

		} else {
			le = breaks
		}

		.Object@colors = col_fun(le)
		.Object@levels = le
		.Object@col_fun = col_fun
		.Object@type = "continuous"
	} else {
		stop_wrap("initialization failed. Either specify `colors` + `levels` or `col_fun` + `breaks`\n")
	}

	.Object@name = name
	na_col = t(col2rgb(na_col, alpha = TRUE))
	na_col = rgb(na_col[, 1:3, drop = FALSE], alpha = na_col[, 4], maxColorValue = 255)
	.Object@na_col = na_col[1]

	return(.Object)
}

# == title
# Print the ColorMapping Object
#
# == param
# -object A `ColorMapping-class` object.
#
# == value
# This function returns no value.
#
# == author
# Zuguang Gu <z.gu@dkfz.de>
#
setMethod(f = "show",
	signature = "ColorMapping",
	definition = function(object) {
	if(object@type == "discrete") {
		cat("Discrete color mapping:\n")
		cat("name:", object@name, "\n")
		cat("levels:\n")
		print(object@levels)
		cat("\n")
		cat("colors:\n")
		col = object@colors; names(col) = NULL
		print(col)
		cat("\n")
	} else if(object@type == "continuous") {
		cat("Continuous color mapping:\n")
		cat("name:", object@name, "\n")
		cat("default breaks:\n")
		print(object@levels)
		cat("\n")
		cat("colors:\n")
		col = object@colors; names(col) = NULL
		print(col)
		cat("\n")
	}
})

# == title
# Map Values to Colors
#
# == param
# -object A `ColorMapping-class` object.
# -x Input values.
#
# == details
# It maps a vector of values to a vector of colors.
#
# This function provides a uniform way for discrete and continuous color mapping.
#
# == value
# A vector of colors.
#
# == author
# Zuguang Gu <z.gu@dkfz.de>
#
# == example
# cm = ColorMapping(colors = c("A" = "red", "B" = "black"))
# map_to_colors(cm, sample(c("A", "B"), 10, replace = TRUE))
# require(circlize)
# col_fun = colorRamp2(c(0, 1), c("white", "red"))
# cm = ColorMapping(col_fun = col_fun)
# map_to_colors(cm, runif(10))
setMethod(f = "map_to_colors",
	signature = "ColorMapping",
	definition = function(object, x) {

	if(is.factor(x)) x = as.vector(x)
	original_attr = attributes(x)
	x2 = vector(length = length(x))

	if(object@type == "discrete") {
		x[grepl("^\\s*$", x)] = NA
		lna = is.na(x)

		if(length(object@full_col) > 0) {
			if(is.numeric(x)) x = as.character(x)
			if(any(! x[!lna] %in% names(object@full_col))) {
				msg = paste0(object@name, ": cannot map colors to some of the levels:\n", paste(setdiff(x[!lna], names(object@full_col)), sep = ", ", collapse = ", "))
				stop_wrap(msg)
			}

			x2[lna] = object@na_col
			x2[!lna] = object@full_col[ x[!lna] ]
		} else {
			if(is.numeric(x)) x = as.character(x)
			if(any(! x[!lna] %in% object@levels)) {
				msg = paste0(object@name, ": cannot map colors to some of the levels:\n", paste(setdiff(x[!lna], object@levels), sep = ", ", collapse = ", "))
				stop_wrap(msg)
			}

			x2[lna] = object@na_col
			x2[!lna] = object@colors[ x[!lna] ]
		}
	} else {
		lna = is.na(x)
		x2[lna] = object@na_col
		x2[!lna] = object@col_fun(x[!lna])
	}

	# keep original attributes, such as dimension
	attributes(x2) = original_attr
	return(x2)
})


# == title
# Draw Legend Based on Color Mapping
#
# == param
# -object A `ColorMapping-class` object.
# -plot Whether to plot or just return the legend object?
# -... Pass to `draw,Legends-method`.
# -color_bar "continous" or "discrete". It controls whether to show the discrete legend for the continuous color mapping.
# -title Title of the legend, by default it is the name of the legend.
# -title_gp Graphical parameters for legend title.
# -title_position Position of the title. See `Legend` for all possible values.
# -grid_height Height of each legend grid. Pass to `Legend`.
# -grid_width Width of each legend grid. Pass to `Legend`.
# -border Color for legend grid borders. Pass to `Legend`.
# -at Break values of the legend. By default it is the levels in the `ColorMapping-class` object.
# -labels Labels corresponding to break values.
# -labels_gp Graphcial parameters for legend labels.
# -labels_rot Rotation of labels.
# -nrow Pass to `Legend`. It controls the layout of legend grids if they are arranged in multiple rows or columns.
# -ncol Pass to `Legend`. It controls the layout of legend grids if they are arranged in multiple rows or columns.
# -by_row Pass to `Legend`. It controls the order of legend grids if they are arranged in multiple rows or columns.
# -legend_height Height of the legend body. It only works when ``color_bar`` is ``continuous`` and ``direction`` is ``vertical``. Pass to `Legend`.
# -legend_width Width of the legend body. It only works when ``color_bar`` is ``continuous`` and ``direction`` is ``horizontal``. Pass to `Legend`.
# -legend_direction When ``color_bar`` is ``continuous``, whether the legend is vertical or horizontal? Pass to `Legend`.
# -break_dist A zooming factor to control relative distance of two neighbouring break values.The length
#     of it should be ``length(at) - 1`` or a scalar. 
# -graphics Internally used.
# -param All the legend-related parameters can be specified as a single list.
#
# == details
# The legend is constructed by `Legend`. 
#
# == value
# A `Legends-class` object.
#
# == author
# Zuguang Gu <z.gu@dkfz.de>
#
setMethod(f = "color_mapping_legend",
	signature = "ColorMapping",
	definition = function(object,
	plot = TRUE, ...,
	
	color_bar = object@type,
	
	title = object@name,
	title_gp = gpar(fontsize = 10, fontface = "bold"),
	title_position = "topleft",
	grid_height = unit(4, "mm"),
	grid_width = unit(4, "mm"),
	border = NULL,
	at = object@levels,
	labels = at,
	labels_gp = gpar(fontsize = 10),
	labels_rot = 0,
	nrow = NULL,
	ncol = 1,
	by_row = FALSE,
	legend_height = NULL, 
	legend_width = NULL,
	legend_direction = c("vertical", "horizontal"),
	break_dist = NULL,

	graphics = NULL,
	param = NULL) {

	e = environment()
	if(!is.null(param)) {
		for(nm in names(param)) {
			if(nm == "direction") {
				assign("legend_direction", param[[nm]], envir = e)
			} else {
				assign(nm, param[[nm]], envir = e)
			}
		}
	}

	title_gp = check_gp(title_gp)
	labels_gp = check_gp(labels_gp)

	if(object@type == "discrete" && color_bar == "continuous") {
		stop_wrap("'color_bar' can only be set to 'discrete' only if the color mapping is discrete")
	}

	# get labels
	if(length(at) != length(labels)) {
		stop_wrap("Length of 'at' should be same as length of 'labels'.")
	}

	# if it is character color mapping, remove items in `at` which are not in the available optinos
	if(color_bar == "discrete" && is.character(at)) {
		if(length(object@full_col) > 0) {
			l = at %in% names(object@full_col)
			if(any(!l)) {
				message_wrap(paste0("Following `at` are removed: ", paste(at[!l], collapse = ", "), ", because no color was defined for them."))
			}
			at = at[l]
			labels = labels[l]
		} else {
			l = at %in% object@levels
			if(any(!l)) {
				message_wrap(paste0("Following `at` are removed: ", paste(at[!l], collapse = ", "), ", because no color was defined for them."))
			}
			at = at[l]
			labels = labels[l]
		}
	}

	if(color_bar == "discrete") {
		if(object@type == "continuous") {
			at = rev(at)
			labels = rev(labels)
		}
		if(length(at) == 0) {
			gf = Legend(at = "NA", labels = "NA", name = object@name, title = title, title_gp = title_gp, grid_height = grid_height,
				grid_width = grid_width, border = border, labels_gp = labels_gp, direction = legend_direction, nrow = nrow, ncol = ncol,
				legend_gp = gpar(fill = object@na_col), title_position = title_position, by_row = by_row, graphics = graphics, break_dist = break_dist)
		} else {
			gf = Legend(at = at, labels = labels, name = object@name, title = title, title_gp = title_gp, grid_height = grid_height,
				grid_width = grid_width, border = border, labels_gp = labels_gp, direction = legend_direction, nrow = nrow, ncol = ncol,
				legend_gp = gpar(fill = map_to_colors(object, at)), title_position = title_position, by_row = by_row, graphics = graphics, break_dist = break_dist)
		}

	} else {
		gf = Legend(at = at, labels = labels, name = object@name, col_fun = object@col_fun, title = title, title_gp = title_gp, grid_height = grid_height,
				grid_width = grid_width, border = border, labels_gp = labels_gp, labels_rot = labels_rot, direction = legend_direction,
				legend_width = legend_width, legend_height = legend_height, title_position = title_position, by_row = by_row, break_dist = break_dist)

	}

	if(plot) {
		draw(gf, ...)
	}

	return(invisible(gf))
})

# == title
# Concatenate A List of ColorMapping objects
#
# == param
# -... A list of `ColorMapping-class` objects.
# -name Name of the new merged color mapping.
#
# == details
# Only discrete color mappings can be concatenated.
#
# == example
# cm1 = ColorMapping(colors = c("A" = "red", "B" = "black"))
# cm2 = ColorMapping(colors = c("B" = "blue", "C" = "green"))
# c(cm1, cm2)
c.ColorMapping = function(..., name = NULL) {
	cm_list = list(...)
	if(!all(sapply(cm_list, function(x) x@type) == "discrete")) {
		stop_wrap("Only discrete color mappings can be concatenated.")
	}

	all_levels = unlist(lapply(cm_list, function(x) x@levels))
	all_colors = unlist(lapply(cm_list, function(x) x@colors))

	increase_color_mapping_index()
	if(is.null(name)) name = paste0("color_mapping_", get_color_mapping_index())

	l_dup = duplicated(all_levels)

	all_levels = all_levels[!l_dup]
	all_colors = all_colors[!l_dup]
	names(all_colors) = all_levels

	cm = new("ColorMapping")
	cm@colors = all_colors
	cm@levels = all_levels
	cm@type = "discrete"
	cm@name = name
	cm@na_col = cm_list[[1]]@na_col
	return(cm)
}