# Piper diagram in modern R
#
# A self-contained replacement for the 2019 version of this plot, which relied
# on hydrogeo::toPercent and a sourced ggplot_Piper.R. Both are gone: hydrogeo
# was removed from CRAN, as was rgdal. This file needs only ggplot2.
#
# Geometry (the standard Piper layout, all units in plot space):
#   cation ternary   (0,0) (100,0) (50, 86.603)
#   anion ternary    (120,0) (220,0) (170, 86.603)
#   diamond          bottom (110, 17.3206)  left  (60, 103.9236)
#                    right  (160, 103.9236) top   (110, 190.5266)
# The diamond vertices are computed as the intersections of the two ternaries'
# projection lines at +/- tan(60 deg).
#
# Usage:
#   source("piper.R")
#   df <- piper_percent(raw, cations = c("Ca","Mg","Na","K"),
#                            anions  = c("Cl","SO4","CO3","HCO3"))
#   piper_plot(piper_transform(df), colour = df$Formation)

library(ggplot2)

GRAD <- tan(60 * pi / 180)   # 1.732051
APEX <- 50 * GRAD            # 86.60254
OFFSET <- 120                # x shift of the anion ternary

# ---- data preparation ------------------------------------------------------

#' Convert milliequivalent concentrations to percentages of cation / anion sums.
#' Replaces hydrogeo::toPercent.
piper_percent <- function(df,
                          cations = c("Ca", "Mg", "Na", "K"),
                          anions  = c("Cl", "SO4", "CO3", "HCO3")) {
  stopifnot(all(c(cations, anions) %in% names(df)))
  cat_sum <- rowSums(df[, cations, drop = FALSE])
  an_sum  <- rowSums(df[, anions,  drop = FALSE])
  if (any(cat_sum == 0 | an_sum == 0)) {
    stop("a sample has a zero cation or anion sum; cannot convert to percent")
  }
  df[, cations] <- 100 * df[, cations] / cat_sum
  df[, anions]  <- 100 * df[, anions]  / an_sum
  df
}

#' Project percentage data into Piper plot coordinates.
#' Returns one row per sample with the three point positions.
piper_transform <- function(df) {
  ca  <- df$Ca
  mg  <- df$Mg
  cl  <- df$Cl
  so4 <- df$SO4

  # cation ternary
  cx <- 100 * (1 - (ca / 100) - (mg / 200))
  cy <- mg * GRAD / 2

  # anion ternary
  ax <- OFFSET + cl + 0.5 * so4
  ay <- so4 * GRAD / 2

  # diamond: where the up-slope from the cation point meets the down-slope
  # from the anion point
  dx <- (GRAD * cx - (-GRAD) * ax + ay - cy) / (2 * GRAD)
  dy <- GRAD * (dx - cx) + cy

  data.frame(
    cation_x = cx, cation_y = cy,
    anion_x  = ax, anion_y  = ay,
    diamond_x = dx, diamond_y = dy
  )
}

# ---- plot furniture --------------------------------------------------------

# Internal grid lines at 20/40/60/80 for one ternary, offset in x.
ternary_grid <- function(dx = 0) {
  t <- c(20, 40, 60, 80)
  h <- t * GRAD / 2
  rbind(
    # left edge -> base
    data.frame(x = dx + t / 2,        y = h,        xend = dx + 100 - t / 2, yend = h),
    # base -> right edge
    data.frame(x = dx + t,            y = 0,        xend = dx + t + (100 - t) / 2, yend = (100 - t) * GRAD / 2),
    # base -> left edge
    data.frame(x = dx + t,            y = 0,        xend = dx + t / 2,       yend = t * GRAD / 2)
  )
}

diamond_grid <- function() {
  t <- c(20, 40, 60, 80)
  # lines parallel to the two diamond axes
  rbind(
    data.frame(x = 110 - t / 2,        y = 17.3206 + t * GRAD / 2,
               xend = 160 - t / 2,     yend = 103.9236 + t * GRAD / 2),
    data.frame(x = 110 + t / 2,        y = 17.3206 + t * GRAD / 2,
               xend = 60 + t / 2,      yend = 103.9236 + t * GRAD / 2)
  )
}

piper_base <- function() {
  tri_l <- data.frame(x = c(0, 100, 50, 0), y = c(0, 0, APEX, 0))
  tri_r <- data.frame(x = c(0, 100, 50, 0) + OFFSET, y = c(0, 0, APEX, 0))
  diam  <- data.frame(x = c(110, 60, 110, 160, 110),
                      y = c(17.3206, 103.9236, 190.5266, 103.9236, 17.3206))

  ggplot() +
    geom_segment(data = rbind(ternary_grid(0), ternary_grid(OFFSET), diamond_grid()),
                 aes(x = x, y = y, xend = xend, yend = yend),
                 linetype = "dashed", colour = "grey60", linewidth = 0.3) +
    geom_path(data = tri_l, aes(x, y), linewidth = 0.6, colour = "black") +
    geom_path(data = tri_r, aes(x, y), linewidth = 0.6, colour = "black") +
    geom_path(data = diam,  aes(x, y), linewidth = 0.6, colour = "black") +
    coord_equal(clip = "off") +
    theme_void()
}

# Axis tick labels. Rotation follows the original: the two outer near-vertical
# edges (Mg, SO4) read horizontally, every other edge is rotated to sit along
# its own grid family.
piper_labels <- function(p) {
  t <- c(20, 40, 60, 80)
  pad <- 4.6
  dx <- pad * 0.87
  dy <- pad * 0.5

  lab <- rbind(
    # --- cation ternary ---
    # Mg, left edge, horizontal
    data.frame(x = t / 2 - dx - 1.6, y = t * GRAD / 2 + dy, l = t, a = 0),
    # Ca, base
    data.frame(x = t, y = -pad, l = rev(t), a = 60),
    # Na + K, right edge
    data.frame(x = 100 - t / 2 + dx, y = t * GRAD / 2 + dy, l = rev(t), a = -60),

    # --- anion ternary ---
    # alkalinity, left edge
    data.frame(x = OFFSET + t / 2 - dx, y = t * GRAD / 2 + dy, l = rev(t), a = 60),
    # Cl, base
    data.frame(x = OFFSET + t, y = -pad, l = t, a = 60),
    # SO4, right edge, horizontal
    data.frame(x = OFFSET + 100 - t / 2 + dx + 1.6, y = t * GRAD / 2 + dy, l = t, a = 0),

    # --- diamond ---
    # SO4 + Cl, upper-left edge, increasing toward the top vertex
    data.frame(x = 60 + t / 2 - dx, y = 103.9236 + t * GRAD / 2 + dy, l = t, a = 60),
    # Ca + Mg, upper-right edge
    data.frame(x = 160 - t / 2 + dx, y = 103.9236 + t * GRAD / 2 + dy, l = t, a = -60)
  )
  p + geom_text(data = lab, aes(x, y, label = l, angle = a), size = 2.6, colour = "grey20")
}

piper_titles <- function(p) {
  # Positions are computed on each edge's outward normal. Identical to piper.py.
  ann <- list(
    list(76.2, 158.1,      'SO[4]^{"2-"}~+~Cl^{"-"}',    60),
    list(143.8, 158.1,     'Ca^{"2+"}~+~Mg^{"2+"}',     -60),
    list(21.5, 59.2,       'Mg^{"2+"}',                  60),
    list(50, -13,          'Ca^{"2+"}',                   0),
    list(78.5, 59.2,       'Na^{"+"}~+~K^{"+"}',        -60),
    list(OFFSET + 21.5, 59.2, 'Alkalinity~as~HCO[3]^{"-"}', 60),
    list(OFFSET + 50, -13, 'Cl^{"-"}',                    0),
    list(OFFSET + 78.5, 59.2, 'SO[4]^{"2-"}',           -60)
  )
  for (a in ann) {
    p <- p + annotate("text", x = a[[1]], y = a[[2]], label = a[[3]],
                      parse = TRUE, angle = a[[4]], size = 3.4)
  }
  p
}

# ---- location map ----------------------------------------------------------

M_PER_DEG_LAT <- 111132

#' Read a single-polygon GeoJSON into a data frame of lon/lat vertices.
#' Avoids a GDAL/sf dependency for the one polygon this needs.
basin_polygon <- function(path) {
  g <- jsonlite::fromJSON(path, simplifyVector = FALSE)
  geom <- g$features[[1]]$geometry
  ring <- if (geom$type == "Polygon") geom$coordinates[[1]] else geom$coordinates[[1]][[1]]
  data.frame(
    lon = vapply(ring, function(p) p[[1]], numeric(1)),
    lat = vapply(ring, function(p) p[[2]], numeric(1))
  )
}

deg_label <- function(v, suffix) {
  paste0(formatC(abs(v), format = "f", digits = 2), "°", suffix)
}

#' Location map with graticule labels, a north arrow and a scale bar.
#'
#' @param basin  data frame of lon/lat from basin_polygon()
#' @param wells  data frame with Longitude, Latitude
#' @param colour optional grouping vector, one value per well
#' @param km     scale bar length in kilometres
basin_map <- function(basin, wells, colour = NULL, km = 5, legend_title = "Formation") {
  lat_mid <- mean(range(basin$lat))
  # One degree of longitude shortens with latitude; fixing the ratio this way
  # keeps distances true without reprojecting.
  aspect <- 1 / cos(lat_mid * pi / 180)

  xr <- range(basin$lon); yr <- range(basin$lat)
  padx <- diff(xr) * 0.06; pady <- diff(yr) * 0.06

  # Scale bar, anchored bottom-left, sized in degrees of longitude.
  deg_per_km <- 1000 / (M_PER_DEG_LAT * cos(lat_mid * pi / 180))
  bx0 <- xr[1] + padx * 0.3
  bx1 <- bx0 + km * deg_per_km
  by <- yr[1] - pady * 0.35

  # North arrow, anchored top-right.
  nx <- xr[2] + padx * 0.35
  ny0 <- yr[2] - diff(yr) * 0.13
  ny1 <- yr[2] - diff(yr) * 0.02

  if (!is.null(colour)) wells$grp <- colour

  p <- ggplot() +
    geom_polygon(data = basin, aes(lon, lat),
                 fill = "grey72", colour = "grey35", linewidth = 0.4)

  if (is.null(colour)) {
    p <- p + geom_point(data = wells, aes(Longitude, Latitude), size = 1.8)
  } else {
    p <- p + geom_point(data = wells, aes(Longitude, Latitude, colour = grp), size = 1.8) +
      labs(colour = legend_title)
  }

  p +
    # scale bar
    annotate("segment", x = bx0, xend = bx1, y = by, yend = by, linewidth = 0.8) +
    annotate("segment", x = bx0, xend = bx0, y = by - pady * 0.12, yend = by + pady * 0.12,
             linewidth = 0.8) +
    annotate("segment", x = bx1, xend = bx1, y = by - pady * 0.12, yend = by + pady * 0.12,
             linewidth = 0.8) +
    annotate("text", x = (bx0 + bx1) / 2, y = by + pady * 0.45,
             label = paste0(km, " km"), size = 2.8) +
    # north arrow
    annotate("segment", x = nx, xend = nx, y = ny0, yend = ny1,
             arrow = arrow(length = unit(0.16, "cm"), type = "closed"), linewidth = 0.6) +
    annotate("text", x = nx, y = ny1 + diff(yr) * 0.035, label = "N",
             size = 3.2, fontface = "bold") +
    scale_x_continuous(labels = function(v) deg_label(v, "W"),
                       expand = expansion(mult = 0.10)) +
    scale_y_continuous(labels = function(v) deg_label(v, "N"),
                       expand = expansion(mult = 0.10)) +
    coord_fixed(ratio = aspect, clip = "off") +
    labs(x = NULL, y = NULL) +
    theme_bw(base_size = 9) +
    theme(panel.grid = element_line(colour = "grey88", linewidth = 0.3),
          axis.text = element_text(size = 7, colour = "grey25"),
          legend.position = "none")
}

#' Draw a full Piper diagram.
#' @param pts  output of piper_transform()
#' @param colour optional grouping vector, one value per sample
piper_plot <- function(pts, colour = NULL, legend_title = "Formation") {
  long <- rbind(
    data.frame(x = pts$cation_x,  y = pts$cation_y),
    data.frame(x = pts$anion_x,   y = pts$anion_y),
    data.frame(x = pts$diamond_x, y = pts$diamond_y)
  )
  if (!is.null(colour)) long$grp <- rep(colour, 3)

  p <- piper_titles(piper_labels(piper_base()))
  if (is.null(colour)) {
    p <- p + geom_point(data = long, aes(x, y), size = 2, alpha = 0.9)
  } else {
    p <- p + geom_point(data = long, aes(x, y, colour = grp), size = 2, alpha = 0.9) +
      labs(colour = legend_title) +
      theme(legend.position = "inside",
            legend.position.inside = c(0.14, 0.88),
            legend.title = element_text(face = "bold", size = 8),
            legend.text = element_text(size = 7),
            legend.key.size = unit(0.8, "lines"))
  }
  p
}