#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @importFrom grid segmentsGrob polylineGrob gpar
#' @importFrom dplyr group_by_ do ungroup
#' @importFrom ggforce interpolateDataFrame
#' @export
GeomEdgePath <- ggproto(
  'GeomEdgePath',
  GeomPath,
  draw_panel = function(
    self,
    data,
    panel_scales,
    coord,
    arrow = NULL,
    lineend = 'butt',
    linejoin = 'round',
    linemitre = 1,
    na.rm = FALSE,
    interpolate = TRUE,
    label_colour = 'black',
    label_alpha = 1,
    label_parse = FALSE,
    check_overlap = FALSE,
    angle_calc = 'none',
    force_flip = TRUE,
    label_dodge = NULL,
    label_push = NULL
  ) {
    if (!anyDuplicated(data$group)) {
      cli::cli_inform(c(
        "{.fn {snake_class(self)}}: Each group consists of only one observation.",
        i = "Do you need to adjust the {.field group} aesthetic?"
      ))
    }
    data <- data[order(data$group), , drop = FALSE]
    data$group <- match(data$group, unique0(data$group))
    if (interpolate) {
      geometries <- vapply(data, is.geometry, logical(1))
      geom_cols <- data[geometries]
      data <- cbind(interpolateDataFrame(data[!geometries]), geom_cols)
    }
    data <- coord$transform(data, panel_scales)
    zero <- coord$transform(data_frame0(x = 0, y = 0), panel_scales)
    if (nrow(data) < 2) {
      return(zeroGrob())
    }
    attr <- pathAttr(
      as.integer(data$group),
      as.numeric(data$edge_alpha),
      as.numeric(data$edge_width),
      as.character(data$edge_linetype),
      as.character(data$edge_colour),
      length(unique0(data$group))
    )

    if (all(is.na(data$start_cap))) {
      start_captype <- 'circle'
      start_cap <- NULL
      start_cap2 <- NULL
    } else {
      scap <- data$start_cap[!duplicated(data$group)]
      start_captype <- geo_type(scap)
      start_cap <- geo_width(scap)
      start_cap2 <- geo_height(scap)
      native_start <- grepl('native', as.character(start_cap))
      native_start2 <- grepl('native', as.character(start_cap2))
      if (any(native_start) || any(native_start2)) {
        recalc <- coord$transform(
          data_frame0(x = as.vector(start_cap), y = as.vector(start_cap2)),
          panel_scales
        )
        start_cap[native_start] <- unit(recalc$x - zero$x, 'npc')[native_start]
        start_cap2[native_start2] <- unit(recalc$y - zero$y, 'npc')[
          native_start2
        ]
      }
    }
    if (all(is.na(data$end_cap))) {
      end_captype <- 'circle'
      end_cap <- NULL
      end_cap2 <- NULL
    } else {
      ecap <- data$end_cap[!duplicated(data$group)]
      end_captype <- geo_type(ecap)
      end_cap <- geo_width(ecap)
      end_cap2 <- geo_height(ecap)
      native_end <- grepl('native', as.character(end_cap))
      native_end2 <- grepl('native', as.character(end_cap2))
      if (any(native_end) || any(native_end2)) {
        recalc <- coord$transform(
          data_frame0(x = as.vector(end_cap), y = as.vector(end_cap2)),
          panel_scales
        )
        end_cap[native_end] <- unit(recalc$x - zero$x, 'native')[native_end]
        end_cap2[native_end2] <- unit(recalc$y - zero$y, 'native')[native_end2]
      }
    }
    if (
      (is.null(start_cap) && !is.null(end_cap)) ||
        (!is.null(start_cap) && is.null(end_cap))
    ) {
      if (is.null(start_cap)) {
        start_cap <- 0
      }
      if (is.null(end_cap)) end_cap <- 0
    }

    solid_lines <- all(attr$solid)
    constant <- all(attr$constant)
    if (!solid_lines && !constant) {
      cli::cli_abort(
        "{.fn {snake_class(self)}} can't have varying {.field edge_colour}, {.field edge_width}, and/or {.field edge_alpha} along the line when {.field edge_linetype} isn't solid"
      )
    }

    gp <- gpar(
      col = alpha(data$edge_colour, data$edge_alpha),
      fill = alpha(data$edge_colour, data$edge_alpha),
      lwd = data$edge_width * .pt,
      lty = data$edge_linetype,
      lineend = lineend,
      linejoin = linejoin,
      linemitre = linemitre
    )
    edge_grob <- cappedPathGrob(
      x = data$x,
      y = data$y,
      id = data$group,
      arrow = arrow,
      start.cap = start_cap,
      start.cap2 = start_cap2,
      start.captype = start_captype,
      end.cap = end_cap,
      end.cap2 = end_cap2,
      end.captype = end_captype,
      default.units = 'native',
      gp = gp,
      constant = constant
    )
    if (any(!is.na(data$label))) {
      if (angle_calc == 'none') {
        angle_calc <- 'rot'
      }
      edge_ind <- split(seq_len(nrow(data)), data$group)
      edge_length <- lengths(edge_ind)
      edge_start <- c(0, cumsum(edge_length)[-length(edge_length)]) + 1
      label_pos <- data$label_pos[vapply(edge_ind, head, integer(1), n = 1)]
      label_pos <- 1 + floor((edge_length - 1) * label_pos)
      label_ind <- edge_start + label_pos
      label_data <- data[label_ind, ]
      label_data$label_colour <- if (is.na(label_colour)) {
        label_data$edge_colour
      } else {
        label_colour
      }
      label_data$label_alpha <- if (is.na(label_alpha)) {
        label_data$edge_alpha
      } else {
        label_alpha
      }
      angle_start <- ifelse(label_pos == 1, 1, label_pos - 1)
      angle_end <- ifelse(label_pos == edge_length, edge_length, label_pos + 1)
      label_x0 <- data$x[angle_start + edge_start]
      label_y0 <- data$y[angle_start + edge_start]
      label_x1 <- data$x[angle_end + edge_start]
      label_y1 <- data$y[angle_end + edge_start]
      lab <- as.character(label_data$label)
      if (label_parse) {
        lab_expressions <- sapply(X = lab, FUN = str2expression)
        lab_len <- sapply(X = lab_expressions, FUN = length)
        if (any(lab_len == 0)) {
          lab[lab_len == 0] <- "` `"
          lab <- sapply(X = lab, FUN = str2expression)
        } else {
          lab <- lab_expressions
        }
      }
      label_grob <- textAlongGrob(
        lab,
        label_data$x,
        label_data$y,
        default.units = 'native',
        hjust = label_data$hjust,
        vjust = label_data$vjust,
        rot = label_data$angle,
        gp = gpar(
          col = alpha(label_data$label_colour, label_data$label_alpha),
          fontsize = label_data$label_size * .pt,
          fontfamily = label_data$family,
          fontface = label_data$fontface,
          lineheight = label_data$lineheight
        ),
        check.overlap = check_overlap,
        rot.type = angle_calc,
        x0 = label_x0,
        y0 = label_y0,
        x1 = label_x1,
        y1 = label_y1,
        force.rot = force_flip,
        dodge = label_dodge,
        push = label_push
      )
      gList(edge_grob, label_grob)
    } else {
      edge_grob
    }
  },
  draw_key = function(data, params, size) {
    segmentsGrob(
      0.1,
      0.5,
      0.9,
      0.5,
      gp = gpar(
        col = alpha(data$edge_colour, data$edge_alpha),
        fill = alpha(
          params$arrow.fill %||%
            data$edge_colour %||%
            data$edge_fill %||%
            "black",
          data$edge_alpha
        ),
        lwd = data$edge_width * .pt,
        lty = data$edge_linetype,
        lineend = 'butt'
      ),
      arrow = params$arrow
    )
  },
  handle_na = function(self, data, params) {
    if (params$interpolate) {
      return(data)
    }
    keep <- function(x) {
      first <- match(FALSE, x, nomatch = 1) - 1
      last <- length(x) - match(FALSE, rev(x), nomatch = 1) + 1
      c(
        rep(FALSE, first),
        rep(TRUE, last - first),
        rep(FALSE, length(x) - last)
      )
    }
    missing <- !stats::complete.cases(
      data[c('x', 'y', 'edge_width', 'edge_colour', 'edge_linetype')]
    )
    kept <- stats::ave(missing, data$group, FUN = keep)
    data <- data[kept, ]
    if (!all(kept) && !params$na.rm) {
      cli::cli_warn(
        "Removed {sum(!kept)} row{?s} containing missing values ({.fn {snake_class(self)}})."
      )
    }
    data
  },
  default_aes = aes(
    edge_colour = 'black',
    edge_width = 0.5,
    edge_linetype = 'solid',
    edge_alpha = NA,
    start_cap = NA,
    end_cap = NA,
    label = NA,
    label_pos = 0.5,
    label_size = 3.88,
    angle = 0,
    hjust = 0.5,
    vjust = 0.5,
    family = '',
    fontface = 1,
    lineheight = 1.2
  ),
  rename_size = FALSE
)
#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @importFrom grid gpar segmentsGrob
#' @export
GeomEdgeParallelPath <- ggproto(
  'GeomEdgeParallelPath',
  GeomEdgePath,
  draw_panel = function(
    data,
    panel_scales,
    coord,
    arrow = NULL,
    sep = unit(2, 'mm'),
    lineend = 'butt',
    linejoin = 'round',
    linemitre = 1,
    na.rm = FALSE,
    interpolate = TRUE,
    label_colour = 'black',
    label_alpha = 1,
    label_parse = FALSE,
    check_overlap = FALSE,
    angle_calc = 'none',
    force_flip = TRUE,
    label_dodge = NULL,
    label_push = NULL
  ) {
    data <- data[order(data$group), , drop = FALSE]
    panel <- GeomEdgePath$draw_panel(
      data,
      panel_scales,
      coord,
      arrow = arrow,
      lineend = lineend,
      linejoin = linejoin,
      linemitre = linemitre,
      na.rm = na.rm,
      interpolate = interpolate,
      label_colour = label_colour,
      label_alpha = label_alpha,
      label_parse = label_parse,
      check_overlap = check_overlap,
      angle_calc = angle_calc,
      force_flip = force_flip,
      label_dodge = label_dodge,
      label_push = label_push
    )
    if (inherits(panel, 'gList')) {
      panel[[1]]$sep <- (data$.position[!duplicated(data$group)] * sep)[
        panel[[1]]$id
      ]
      class(panel[[1]]) <- c('parallelPath', class(panel[[1]]))
    } else {
      panel$sep <- (data$.position[!duplicated(data$group)] * sep)[panel$id]
      class(panel) <- c('parallelPath', class(panel))
    }
    panel
  }
)
#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @importFrom grid gpar segmentsGrob
#' @export
GeomEdgeSegment <- ggproto(
  'GeomEdgeSegment',
  GeomSegment,
  draw_panel = function(
    data,
    panel_scales,
    coord,
    arrow = NULL,
    lineend = 'butt',
    na.rm = FALSE
  ) {
    if (empty_data(data)) {
      return(zeroGrob())
    }
    coord <- coord$transform(data, panel_scales)
    segmentsGrob(
      coord$x,
      coord$y,
      coord$xend,
      coord$yend,
      default.units = 'native',
      gp = gpar(
        col = alpha(coord$edge_colour, coord$edge_alpha),
        fill = alpha(coord$edge_colour, coord$edge_alpha),
        lwd = coord$edge_width * .pt,
        lty = coord$edge_linetype,
        lineend = lineend
      ),
      arrow = arrow
    )
  },
  draw_key = function(data, params, size) {
    segmentsGrob(
      0.1,
      0.5,
      0.9,
      0.5,
      gp = gpar(
        col = alpha(data$edge_colour, data$edge_alpha),
        fill = alpha(
          params$arrow.fill %||%
            data$edge_colour %||%
            data$edge_fill %||%
            "black",
          data$edge_alpha
        ),
        lwd = data$edge_width * .pt,
        lty = data$edge_linetype,
        lineend = 'butt'
      ),
      arrow = params$arrow
    )
  },
  default_aes = aes(
    edge_colour = 'black',
    edge_width = 0.5,
    edge_linetype = 1,
    edge_alpha = NA
  ),
  rename_size = FALSE
)
#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @export
GeomEdgeParallelSegment <- ggproto(
  'GeomEdgeParallelSegment',
  GeomEdgeSegment,
  draw_panel = function(
    data,
    panel_scales,
    coord,
    arrow = NULL,
    lineend = 'butt',
    na.rm = FALSE,
    sep = unit(2, 'mm')
  ) {
    data <- data[order(data$group), , drop = FALSE]
    panel <- GeomEdgeSegment$draw_panel(
      data,
      panel_scales,
      coord,
      arrow = arrow,
      lineend = lineend,
      na.rm = na.rm
    )
    panel$sep <- data$.position * sep
    class(panel) <- c('parallelPath', class(panel))
    panel
  }
)
#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @importFrom grid pointsGrob unit gpar grobTree
#' @export
GeomEdgeSpanPath <- ggproto(
  'GeomEdgeSpanPath',
  GeomEdgePath,
  draw_panel = function(
    data,
    panel_scales,
    coord,
    arrow = NULL,
    end_shape = NA,
    lineend = 'butt',
    linejoin = 'round',
    linemitre = 1,
    na.rm = FALSE,
    interpolate = TRUE,
    label_colour = 'black',
    label_alpha = 1,
    label_parse = FALSE,
    check_overlap = FALSE,
    angle_calc = 'none',
    force_flip = TRUE,
    label_dodge = NULL,
    label_push = NULL
  ) {
    panel <- GeomEdgePath$draw_panel(
      data,
      panel_scales,
      coord,
      arrow = arrow,
      lineend = lineend,
      linejoin = linejoin,
      linemitre = linemitre,
      na.rm = na.rm,
      interpolate = interpolate,
      label_colour = label_colour,
      label_alpha = label_alpha,
      label_parse = label_parse,
      check_overlap = FALSE,
      angle_calc = 'none',
      force_flip = TRUE,
      label_dodge = label_dodge,
      label_push = label_push
    )
    if (is.na(end_shape)) {
      return(panel)
    }

    data <- data[data$index == 0 | data$index == 1, ]
    end_shape <- translate_shape_string(end_shape)
    coords <- coord$transform(data, panel_scales)
    ends <- pointsGrob(
      coords$x,
      coords$y,
      pch = end_shape,
      size = unit(
        1 / (panel_scales$x.range[2] - (panel_scales$x.range[1])),
        'npc'
      ),
      gp = gpar(
        col = alpha(coords$edge_colour, coords$edge_alpha),
        fill = alpha(coords$edge_colour, coords$edge_alpha),
        lwd = 0
      )
    )
    grobTree(panel, ends)
  }
)
#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @importFrom grid pointsGrob unit gpar grobTree
#' @export
GeomEdgeSpanSegment <- ggproto(
  'GeomEdgeSpanSegment',
  GeomEdgeSegment,
  draw_panel = function(
    data,
    panel_scales,
    coord,
    arrow = NULL,
    lineend = 'butt',
    na.rm = FALSE,
    end_shape = NA
  ) {
    panel <- GeomEdgeSegment$draw_panel(
      data,
      panel_scales,
      coord,
      arrow = arrow,
      lineend = lineend,
      na.rm = na.rm
    )
    if (is.na(end_shape)) {
      return(panel)
    }

    data2 <- data
    data2$x <- data2$xend
    data2$y <- data2$yend
    data <- vec_rbind(data, data2)
    data$xend <- NULL
    data$yend <- NULL

    end_shape <- translate_shape_string(end_shape)
    coords <- coord$transform(data, panel_scales)
    ends <- pointsGrob(
      coords$x,
      coords$y,
      pch = end_shape,
      size = unit(
        1 / (panel_scales$x.range[2] - (panel_scales$x.range[1])),
        'npc'
      ),
      gp = gpar(
        col = alpha(coords$edge_colour, coords$edge_alpha),
        fill = alpha(coords$edge_colour, coords$edge_alpha),
        lwd = 0
      )
    )
    grobTree(panel, ends)
  }
)
#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @importFrom grid gpar pointsGrob
#' @export
GeomEdgePoint <- ggproto(
  'GeomEdgePoint',
  GeomPoint,
  setup_data = function(data, params) {
    if (params$mirror) {
      data2 <- data
      data2[, c('x', 'y')] <- data2[, c('y', 'x'), drop = FALSE]
      data2$x <- abs(data2$x) * sign(data$x)
      data2$y <- abs(data2$y) * sign(data$y)
      data <- vec_rbind(data, data2)
    }
    data
  },
  draw_panel = function(data, panel_scales, coord, na.rm = FALSE) {
    if (empty_data(data)) {
      return(zeroGrob())
    }
    if (is.character(data$edge_shape)) {
      data$edge_shape <- translate_shape_string(data$edge_shape)
    }
    coords <- coord$transform(data, panel_scales)
    coords <- coords[order(coords$edge_size, decreasing = TRUE), , drop = FALSE]
    pointsGrob(
      coords$x,
      coords$y,
      pch = coords$edge_shape,
      gp = gpar(
        col = alpha(coords$edge_colour, coords$edge_alpha),
        fill = alpha(coords$edge_fill, coords$edge_alpha),
        fontsize = coords$edge_size *
          .pt +
          coords$stroke *
            .stroke /
            2,
        lwd = coords$stroke * .stroke / 2
      )
    )
  },
  draw_key = function(data, params, size) {
    pointsGrob(
      0.5,
      0.5,
      pch = data$edge_shape,
      gp = gpar(
        col = alpha(data$edge_colour, data$edge_alpha),
        fill = alpha(data$edge_fill, data$edge_alpha),
        fontsize = data$edge_size * .pt + data$stroke * .stroke / 2,
        lwd = data$stroke * .stroke / 2
      )
    )
  },
  default_aes = aes(
    edge_shape = 19,
    edge_colour = 'black',
    edge_size = 1.5,
    edge_fill = NA,
    edge_alpha = NA,
    stroke = 0.5
  ),
  extra_params = c('na.rm', 'mirror')
)
#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @importFrom grid gpar rectGrob
#' @export
GeomEdgeTile <- ggproto(
  'GeomEdgeTile',
  GeomTile,
  setup_data = function(data, params) {
    if (params$mirror) {
      data2 <- data
      data2[, c('x', 'y')] <- data2[, c('y', 'x'), drop = FALSE]
      data2$x <- abs(data2$x) * sign(data$x)
      data2$y <- abs(data2$y) * sign(data$y)
      data <- vec_rbind(data, data2)
    }
    data
  },
  draw_panel = function(
    data,
    panel_params,
    coord,
    lineend = 'butt',
    linejoin = 'mitre',
    na.rm = FALSE
  ) {
    if (empty_data(data)) {
      return(zeroGrob())
    }
    data$xmin <- data$x - 0.5
    data$xmax <- data$x + 0.5
    data$ymin <- data$y - 0.5
    data$ymax <- data$y + 0.5
    coords <- coord$transform(data, panel_params)
    rectGrob(
      coords$xmin,
      coords$ymax,
      width = coords$xmax - coords$xmin,
      height = coords$ymax - coords$ymin,
      default.units = "native",
      just = 'centre',
      gp = gpar(
        col = coords$edge_colour,
        fill = alpha(coords$edge_fill, coords$edge_alpha),
        lwd = coords$edge_width * .pt,
        lty = coords$edge_linetype,
        linejoin = linejoin,
        lineend = lineend
      )
    )
  },
  draw_key = function(data, params, size) {
    if (is.null(data$edge_width)) {
      data$edge_width <- 0.5
    }

    lwd <- min(data$edge_width, min(size) / 4)

    rectGrob(
      width = unit(1, "npc") - unit(lwd, "mm"),
      height = unit(1, "npc") - unit(lwd, "mm"),
      gp = gpar(
        col = data$edge_colour %||% NA,
        fill = alpha(data$edge_fill %||% "grey20", data$edge_alpha),
        lty = data$edge_linetype %||% 1,
        lwd = lwd * .pt,
        linejoin = "mitre",
        # `lineend` is a workaround for Windows and intentionally kept unexposed
        # as an argument. (c.f. https://github.com/tidyverse/ggplot2/issues/3037#issuecomment-457504667)
        lineend = "square"
      )
    )
  },
  default_aes = aes(
    edge_fill = 'grey20',
    edge_colour = NA,
    edge_width = 0.1,
    edge_linetype = 1,
    edge_alpha = NA
  ),
  rename_size = FALSE,
  extra_params = c('na.rm', 'mirror')
)
#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @importFrom ggforce GeomBezier0
#' @export
GeomEdgeBezier <- ggproto(
  'GeomEdgeBezier',
  GeomBezier0,
  draw_panel = function(
    data,
    panel_scales,
    coord,
    arrow = NULL,
    lineend = 'butt',
    linejoin = 'round',
    linemitre = 1,
    na.rm = FALSE
  ) {
    names(data) <- sub('edge_', '', names(data))
    names(data)[names(data) == 'width'] <- 'linewidth'
    GeomBezier0$draw_panel(
      data,
      panel_scales,
      coord,
      arrow,
      lineend,
      linejoin,
      linemitre,
      na.rm
    )
  },
  draw_key = function(data, params, size) {
    segmentsGrob(
      0.1,
      0.5,
      0.9,
      0.5,
      gp = gpar(
        col = alpha(data$edge_colour, data$edge_alpha),
        fill = alpha(
          params$arrow.fill %||%
            data$edge_colour %||%
            data$edge_fill %||%
            "black",
          data$edge_alpha
        ),
        lwd = data$edge_width * .pt,
        lty = data$edge_linetype,
        lineend = 'butt'
      ),
      arrow = params$arrow
    )
  },
  default_aes = aes(
    edge_colour = 'black',
    edge_width = 0.5,
    edge_linetype = 1,
    edge_alpha = NA
  ),
  rename_size = FALSE,
  handle_na = function(data, ...) {
    data
  }
)
#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @importFrom ggforce GeomBspline0
#' @export
GeomEdgeBspline <- ggproto(
  'GeomEdgeBspline',
  GeomBspline0,
  draw_panel = function(
    data,
    panel_scales,
    coord,
    arrow = NULL,
    lineend = 'butt',
    linejoin = 'round',
    linemitre = 1,
    na.rm = FALSE
  ) {
    names(data) <- sub('edge_', '', names(data))
    names(data)[names(data) == 'width'] <- 'linewidth'
    GeomBspline0$draw_panel(
      data,
      panel_scales,
      coord,
      arrow,
      type = "clamped",
      lineend,
      linejoin,
      linemitre,
      na.rm
    )
  },
  draw_key = function(data, params, size) {
    segmentsGrob(
      0.1,
      0.5,
      0.9,
      0.5,
      gp = gpar(
        col = alpha(data$edge_colour, data$edge_alpha),
        fill = alpha(
          params$arrow.fill %||%
            data$edge_colour %||%
            data$edge_fill %||%
            "black",
          data$edge_alpha
        ),
        lwd = data$edge_width * .pt,
        lty = data$edge_linetype,
        lineend = 'butt'
      ),
      arrow = params$arrow
    )
  },
  default_aes = aes(
    edge_colour = 'black',
    edge_width = 0.5,
    edge_linetype = 1,
    edge_alpha = NA
  ),
  rename_size = FALSE,
  handle_na = function(data, ...) {
    data
  }
)

remove_loop <- function(data) {
  if (nrow(data) == 0) {
    return(data)
  }

  data[!(data$x == data$xend & data$y == data$yend), , drop = FALSE]
}
remove_loop2 <- function(data) {
  if (nrow(data) == 0) {
    return(data)
  }

  data <- data[order(data$group), ]
  loop <- data$x[c(TRUE, FALSE)] == data$x[c(FALSE, TRUE)] &
    data$y[c(TRUE, FALSE)] == data$y[c(FALSE, TRUE)]
  data[rep(!loop, each = 2), , drop = FALSE]
}
