#' Network Plot
#'
#' Draw an undirected graph from a square matrix
#' where each variable is a node and the pairwise
#' entries are the strength of the relationship
#' (ex: correlation, partial correlation).
#' Based on https://r-graph-gallery.com/310-custom-hierarchical-edge-bundling.html
#'
#' @param R A square numeric matrix with names columns and rows
#' @param title Optional plot title
#' @param groups Optional data.frame mapping nodes to groups
#'
#' @return A ggraph object
#'
#' @examples
#' X <- matrix(data = rnorm(100), nrow = 10)
#' R <- cor(X)
#' networkplot(R, title = 'Example Correlation Plot')
#' colnames(R) <- as.character(seq_len(ncol(R)))
#' g <- data.frame(
#'	node = colnames(R),
#'	group = rep(c('A', 'B'), length.out = ncol(R))
#' )
#' networkplot(R, groups = g, title = 'Example Correlation Plot')
#'
#' @import ggraph
#' @import igraph
#'
#' @export
networkplot <- function(R, groups = NULL, title = NULL){
	Rm <- tryCatch({
		if(is.matrix(R)){
			if(nrow(R) != ncol(R)) stop('R must be square')
			Rm <- R
		}else if(is.vector(R)){
			m <- length(R)
			n <- (1 + sqrt(1 + 8 * m)) / 2
			if(n != floor(n)) stop('Vector length not compatible with upper triangle')
			n <- as.integer(n)
			Rm <- matrix(0, n, n)
			Rm[upper.tri(Rm)] <- R
			Rm <- Rm + t(Rm)
			diag(Rm) <- 1
		}else{
			Rm <- as.matrix(R)
			if(nrow(Rm) != ncol(Rm)) stop('R must be square')
		}

		if(is.null(rownames(Rm))){
			rownames(Rm) <- as.character(seq_len(nrow(Rm)))
		}

		if(is.null(colnames(Rm))){
			colnames(Rm) <- as.character(seq_len(ncol(Rm)))
		}

		Rm
	}, error = function(e){
		stop(paste('Invalid input for networkplot:', e$message))
	})

	# Rm is p x p matrix
	# upper.tri(Rm) returns a logical p x p matrix with upper right TRUE
	# which() gives the coordinate pairs for unique edges
	idx <- which(upper.tri(Rm), arr.ind = TRUE)
	
	# edges is a data.frame of k rows with all unique
	# combinations of pairwise entries and value r
	edges <- data.frame(
		from = rownames(Rm)[idx[,1]],
		to = colnames(Rm)[idx[,2]],
		r = Rm[idx]
	)

	# groups must be a data.frame with 2 columns 'node' and 'group'
	if(!is.null(groups)){
		if(!all(c('node','group') %in% colnames(groups))){
			stop("groups must have columns 'node' and 'group'")
		}
		if(!all(rownames(Rm) %in% groups$node)){
			stop('All nodes must appear in groups$node')
		}

		# map hierarchy connecting root -> each group
		d1 <- data.frame(from = 'root', to = unique(groups$group))
		# map hierarchy of each group -> each node
		d2 <- data.frame(from = groups$group, to = groups$node)
		# concatenate them (for 2 levels of the hierarhcy)
		hierarchy <- rbind(d1, d2)
	}else{
		# fallback to single group
		# map root -> each node (no grouping)
		hierarchy <- data.frame(from = 'root', to = rownames(Rm))
	}

	# this is a data.frame with one column listing all vertices
	# this includes the hierarchy vertices (root, groups)
	# this is required by igraph as a node ID table
	vertices <- data.frame(name = unique(c(hierarchy$from, hierarchy$to)))

	# If groups were passed, add a column to vertices called group
	if(!is.null(groups)){
		# maps each vertex to its group, returns NA for root and the
		# group vertices themselves
		# we will use this to also color the nodes
		vertices$group <- groups$group[match(vertices$name, groups$node)]
	}else{
		vertices$group <- 'all'
	}

	# hierarchy is a data.frame listing all the edges
	# vertices is a data.frame listing all the nodes with their attributes
	# returns an igraph object representing the tree structure
	g <- igraph::graph_from_data_frame(
		hierarchy,
		vertices = vertices
	)

	# converts the node names to vertex indeces
	# required by ggraph::get_con()
	from <- match(edges$from, vertices$name)
	to <- match(edges$to, vertices$name)

	# takes the integer vectors from and to and returns the paths along tree
	# transforms node-node edges into paths through hierarchy
	con <- ggraph::get_con(from = from, to = to, value = edges$r)

	# g is the igraph tree
	# dendrogram defines tree embedding
	# circlular arranges the nodes in a circle
	# g contains vertices
	p <- ggraph::ggraph(graph = g, layout = 'dendrogram', circular = TRUE) +
	ggraph::geom_conn_bundle(
		data = con,
		ggplot2::aes(
			edge_alpha = abs(.data$value),
			edge_colour = .data$value
		),
		tension = 0.5
	) +
	ggraph::scale_edge_colour_gradient2(
		low = 'red',
		mid = 'white',
		high = 'blue',
		midpoint = 0,
		limits = c(-1, 1),
		guide = 'none'
	) +
	# operates on node data from the layout
	# leaf is a logical vector (provided by the ggraph object g)
	# where leaf is TRUE, those are the terminal nodes (the ones we want to draw)
	# .data$group is the character vector from vertices
	ggraph::geom_node_point(
		ggplot2::aes(
			filter = .data$leaf,
			colour = .data$group,
			fill = .data$group
		),
		size = 3
	) +
	# labels nodes and groups
	ggraph::geom_node_text(
		ggplot2::aes(label = .data$name),
		# prevents text labels from overlapping
		repel = TRUE,
		size = 3
	) +
	ggplot2::theme_void() +
	ggplot2::theme(legend.position = 'none')

	if(!is.null(title)){
		p <- p +
		ggplot2::ggtitle(title) +
		ggplot2::theme(
			plot.title = ggplot2::element_text(hjust = 0.5)
		)
	}

	return(p)
}