## -----------------------------------------------------------------------------
#| label: setup
#| message: false
#| include: false
library(dplyr)
library(sf)
library(tmap)
library(pycnogrid)


## -----------------------------------------------------------------------------
#| label: fig-sample_population_count
#| fig-cap: "Sample population count data for Lower Manhattan"
#| warning: false
#| echo: false
tm_shape(nyc_ct_small) + tm_fill(
  fill = "populationE",
  fill.scale = tm_scale_intervals(
    values = "viridis",
    style = "jenks",
    n = 6,
    label.format = tm_label_format(big.num.abbr = c("K" = 3))
  ),
  fill_alpha = .75,
  fill.legend = tm_legend(title = "Population Count", frame = FALSE)
) +
  #tm_basemap("CartoDB.PositronNoLabels", alpha = .5) +
  tm_shape(nyc_ct) + tm_fill(fill = "grey97", zindex = 3) +
  tm_layout(
    frame = FALSE,
    bg.color = "grey90"
  )


## -----------------------------------------------------------------------------
#| label: tbl-grids
#| tbl-cap: "Supported target grid systems and their principal geometric and hierarchical properties"
#| echo: false

tibble::tribble(
  ~grid_system, ~cell_geometry, ~area_property, ~extent_and_hierarchy,
  ~subdivision_or_aperture, ~neighbour_topology, ~analytical_implication,

  "H3",
  "Mostly hexagons, with twelve pentagons",
  "Approximately equal-area",
  "Global discrete grid with a nested hierarchy",
  "Aperture-7 hierarchy with alternating cell orientation across resolutions",
  "Usually six edge-neighbours; pentagons have five",
  "Strong indexing and neighbourhood structure, particularly useful for mobility, accessibility, and spatial aggregation",

  "ISEA3H, ISEA4H", # ISEA7H",
  "Mostly hexagons, with twelve pentagons on an icosahedral projection",
  "Equal-area",
  "Global discrete grid with a nested hierarchy",
  "Aperture-3 and Aperture-4 hierarchies, #, and Aperture-7 hierarchies",
  "Usually six edge-neighbours; pentagons have five",
  "Equal-area hexagonal option with relatively fine-grained hierarchical scaling between resolutions",

  #"ISEA4H",
  #"Mostly hexagons, with twelve pentagons on an icosahedral projection",
  #"Equal-area",
  #"Global discrete grid with a nested hierarchy",
  #"Aperture-4 hierarchy",
  #"Usually six edge-neighbours; pentagons have five",
  #"Equal-area hexagonal option with a different resolution progression and nesting structure from aperture-3 systems",

  #"ISEA7H",
  #"Mostly hexagons, with twelve pentagons on an icosahedral projection",
  #"Equal-area",
  #"Global discrete grid with a nested hierarchy",
  #"Aperture-7 hierarchy",
  #"Usually six edge-neighbours; pentagons have five",
  #"Equal-area hexagonal option with resolution scaling broadly comparable to H3 but without H3's specific indexing system",

  #"ISEA4/3H",
  #"Mostly hexagons, with twelve pentagons on an icosahedral projection",
  #"Equal-area",
  #"Global discrete grid with a nested hierarchy",
  #"Mixed aperture-4 and aperture-3 hierarchy",
  #"Usually six edge-neighbours; pentagons have five",
  #"Allows a mixed resolution progression where the scale change between successive levels follows the selected aperture sequence",
  
  "A5",
  "Equal-area pentagonal cells",
  "Equal-area",
  "Global discrete grid with a hierarchical structure",
  "Five-way initial refinement, followed by four-way refinement",
  "Predominantly five edge-neighbours",
  "Provides a global equal-area alternative with strong indexing, although its pentagonal geometry may produce more directional variation than hexagonal grids",

  "S2",
  "Quadrilateral cells projected from the faces of a cube",
  "Not equal-area",
  "Global discrete grid with a nested hierarchy",
  "Quadtree subdivision: each cell has four children",
  "Variable topology across cube-face boundaries",
  "Strong global indexing and web-mapping infrastructure, but cell areas vary substantially across locations",
  
  "Local Raster",
  "Rectangular cells in a projected CRS",
  "Equal-area only when constructed in an appropriate projected CRS",
  "Usually regional or local; hierarchy is not intrinsic",
  "Resolution specified directly in map units",
  "Four- or eight-neighbour structure, depending on rook or queen contiguity",
  "Simple and familiar benchmark, but geometric properties depend on the selected projection and resolution",
  
  "Local Hex",
  "Regular hexagonal cells in a projected CRS",
  "Equal-area only when constructed in an appropriate projected CRS",
  "Usually regional or local; hierarchy is not intrinsic",
  "Resolution specified directly in map units",
  "Usually six edge-neighbours in a complete tessellation",
  "Local hexagonal support with six symmetric first-order neighbours, but geometric properties depend on the selected projection and resolution"
) |>
  gt::gt() |>
  gt::cols_label(
    grid_system = "Grid System",
    cell_geometry = "Cell Geometry",
    area_property = "Area Property",
    extent_and_hierarchy = "Extent and Hierarchy",
    subdivision_or_aperture = "Subdivision / Aperture",
    neighbour_topology = "Neighbour Topology",
    analytical_implication = "Analytical Implication"
  ) |>
  gt::tab_options(
    table.width = gt::pct(100),
    table.font.size = 12
  )


## -----------------------------------------------------------------------------
#| label: fig-sample_dggs_grids
#| fig-cap: "Sample DGGS hierarchies at different resolutions"
#| fig-height: 6
#| fig-width: 9
#| warning: false
#| echo: false
sample_dggs_grids <- bind_rows(
  
  h3_grids <- purrr::map_dfr(
    list(8, 9, 10),
    ~ pycnogrid:::prepare_target_cells(
      source_polys = nyc_ct_small,
      grid_type = "h3",
      resolution = .x,
      cell_inclusion = "intersect"
    ) |> mutate(grid_type = "h3", cell_resolution = .x, zorder = cell_resolution)
  ),
  
  isea3h_grids <- purrr::map_dfr(
    list(17, 18, 19),
    ~ pycnogrid:::prepare_target_cells(
      source_polys = nyc_ct_small,
      grid_type = "isea3h",
      resolution = .x,
      cell_inclusion = "intersect"
    ) |> mutate(grid_type = "isea3h", cell_resolution = .x, zorder = cell_resolution)
  ),
  
  isea4h_grids <- purrr::map_dfr(
    list(13, 14, 15),
    ~ pycnogrid:::prepare_target_cells(
      source_polys = nyc_ct_small,
      grid_type = "isea4h",
      resolution = .x,
      cell_inclusion = "intersect"
    ) |> mutate(grid_type = "isea4h", cell_resolution = .x, zorder = cell_resolution)
  ),
  
  a5_grids <- purrr::map_dfr(
    list(13, 14, 15),
    ~ pycnogrid:::prepare_target_cells(
      source_polys = nyc_ct_small,
      grid_type = "a5",
      resolution = .x,
      cell_inclusion = "intersect"
    ) |> mutate(grid_type = "a5", cell_resolution = .x, zorder = cell_resolution)
  ),
  
  s2_grids <- purrr::map_dfr(
    list(14, 15, 16),
    ~ pycnogrid:::prepare_target_cells(
      source_polys = nyc_ct_small,
      grid_type = "s2",
      resolution = .x,
      cell_inclusion = "intersect"
    ) |> mutate(grid_type = "s2", cell_resolution = .x, zorder = cell_resolution)
  )
) |>
  mutate(grid_type = forcats::fct_relevel(grid_type, "h3", "isea3h", "isea4h", "a5", "s2")) |> 
  arrange(desc(zorder))

tm_shape(sample_dggs_grids) +
  tm_lines(
    col = "cell_resolution",
    col.scale = tm_scale_categorical(),
    col.free = TRUE,
    col.legend = tm_legend(title = "resolution", frame = FALSE)
  ) +
  #tm_basemap("CartoDB.PositronNoLabels", alpha = .5) +
  tm_facets(by = "grid_type",
            ncol = 3,
            free.coords = FALSE) +
  tm_shape(nyc_ct) + tm_fill(fill = "grey97", zindex = 3) +
  tm_layout(
    panel.label.frame = FALSE,
    panel.label.bg = FALSE,
    frame = FALSE,
    bg.color = "grey90"
  )


## -----------------------------------------------------------------------------
#| label: fig-sample_local_grids
#| fig-cap: "Sample local grid hierarchies at different resolutions"
#| fig-height: 3
#| fig-width: 9
#| message: false
#| echo: false
sample_local_grids <- dplyr::bind_rows(
  raster_grids <- purrr::map_dfr(
    list(1000, 500, 250),
    ~ pycnogrid:::prepare_target_cells(
      source_polys = nyc_ct_small,
      grid_type = "raster",
      resolution = .x,
      cell_inclusion = "intersect"
    ) |>
      mutate(
        grid_type = "raster",
        cell_resolution = .x,
        zorder = 1 / cell_resolution
      )
  ),
  hex_grids <- purrr::map_dfr(
    list(1000, 500, 250),
    ~ pycnogrid:::prepare_target_cells(
      source_polys = nyc_ct_small,
      grid_type = "hex",
      resolution = .x,
      cell_inclusion = "intersect"
    ) |>
      mutate(
        grid_type = "hex",
        cell_resolution = .x,
        zorder = 1 / cell_resolution
      )
  )
) |> mutate(grid_type = forcats::fct_relevel(grid_type, "raster", "hex")) |>
  arrange(desc(zorder))

tm_shape(sample_local_grids) +
  tm_lines(
    col = "cell_resolution",
    col.scale = tm_scale_categorical(),
    col.free = TRUE,
    col.legend = tm_legend(title = "resolution", frame = FALSE)
  ) +
  #tm_basemap("CartoDB.PositronNoLabels", alpha = .5) +
  tm_facets_wrap(by = "grid_type",
            #ncol = 2,
            free.coords = FALSE) +
  tm_shape(nyc_ct) + tm_fill(fill = "grey97", zindex = 3) +
  tm_layout(
    panel.label.frame = FALSE,
    panel.label.bg = FALSE,
    frame = FALSE,
    bg.color = "grey90"
  )


## -----------------------------------------------------------------------------
#| label: run to_grid
pycno_nyc_ct_small <- nyc_ct_small |>
  pycnogrid::to_grid(
    value_col = populationE,
    grid_type = "h3",
    resolution = 10
  )


## -----------------------------------------------------------------------------
#| label: fig-pycno_nyc_ct_small
#| fig-cap: "Census tract population counts interpolated to an H3 grid"
#| echo: false
tm_shape(pycno_nyc_ct_small) + tm_fill(
  fill = "pycno_populationE",
  fill.scale = tm_scale_intervals(
    values = "viridis",
    style = "jenks",
    n = 6,
    label.format = tm_label_format(big.num.abbr = c("K" = 3))
  ),
  fill_alpha = .75,
  fill.legend = tm_legend(title = "Interpolated \nPopulation", frame = FALSE)
) +
  tm_shape(nyc_ct_small) + tm_lines(col = "grey45") +
  #tm_basemap("CartoDB.PositronNoLabels", alpha = .5) +
  tm_shape(nyc_ct) + tm_fill(fill = "grey97", zindex = 3) +
  tm_layout(
    frame = FALSE,
    bg.color = "grey90"
  )


## -----------------------------------------------------------------------------
#| label: fig-pycno_nyc_ct_small_resolutions
#| fig-cap: "Census tract population counts interpolated to all four grid types"
#| fig-height: 8
#| fig-width: 10
#| echo: false
grid_list <- list("h3", "isea3h", "isea4h", "a5", "s2", "raster", "hex")
resolution_list <- list(10, 20, 16, 16, 16, 100, 100)

pycno_nyc_ct_small_resolutions <- purrr::map2_dfr(
  grid_list, resolution_list, ~ to_grid(
    source = nyc_ct_small,
    value_col = "populationE",
    grid_type = .x,
    resolution = .y
  ) |> 
    mutate(grid_type = .x, resolution = .y, facet_label = paste0(grid_type, " at resolution ", resolution))
)

tm_shape(pycno_nyc_ct_small_resolutions |> filter(grid_type != "h3" & grid_type != "isea4h" & grid_type != "hex")) + tm_fill(
  fill = "pycno_populationE",
  fill.scale = tm_scale_intervals(
    values = "viridis",
    style = "jenks",
    n = 6,
    label.format = tm_label_format(big.num.abbr = c("K" = 3))
  ),
  fill_alpha = .75,
  fill.legend = tm_legend(title = "Interpolated \nPopulation", frame = FALSE),
  fill.free = TRUE
) +
  tm_facets(by = "facet_label", ncol = 2,
                 free.coords = TRUE) +
  tm_shape(nyc_ct_small) + tm_lines(col = "grey45") +
  #tm_basemap("CartoDB.PositronNoLabels", alpha = .5) +
  tm_shape(nyc_ct) + tm_fill(fill = "grey97", zindex = 3) +
  tm_layout(panel.label.frame = FALSE, panel.label.bg = FALSE, frame = FALSE, bg.color = "grey90")


## -----------------------------------------------------------------------------
#| label: tbl-grid_descriptives
#| tbl-cap: "Population descriptive statistics for different grid types"
#| echo: false
pycno_nyc_ct_small_resolutions |> 
  st_drop_geometry() |> 
  group_by(grid_type) |> 
  summarize(
    resolution = min(resolution),
    cell_count = n(), 
    mean_pop = mean(pycno_populationE), 
    sum_pop = sum(pycno_populationE)
    ) |> 
  ungroup() |> 
  mutate(grid_type = forcats::fct_relevel(grid_type, "h3", "isea3h", "isea4h", "a5", "s2", "raster", "hex")) |> 
  arrange(grid_type) |> 
  gt::gt() |> 
  gt::cols_label(
    grid_type = "grid type",
    resolution = "grid resolution",
    cell_count = "grid cell count",
    mean_pop = "mean cell population",
    sum_pop = "total population"
  )


## -----------------------------------------------------------------------------
#| label: fig-pycno_nyc_ct_small_combinations
#| fig-cap: "Interpolated population counts with varying inclusion and allocation parameters"
#| fig-height: 8
#| fig-width: 8
#| echo: false
combinations <- tidyr::expand_grid(
  inclusion_list = list("intersect", "centroid"),
  allocation_list = list("area", "centroid")
)

pycno_nyc_ct_small_combinations <- purrr::map2_dfr(
  combinations |> purrr::pluck("inclusion_list"),
  combinations |> purrr::pluck("allocation_list"),
  ~ to_grid(
    source = nyc_ct_small,
    value_col = "populationE",
    id_col = id,
    grid_type = "h3",
    resolution = 10,
    cell_inclusion = .x,
    cell_allocation = .y,
    nb_order = 1
  ) |>
    mutate(
      inclusion = paste0("inclusion: ", .x),
      allocation = paste0("allocation: ", .y)
    )
) |>
  mutate(inclusion = forcats::fct_relevel(inclusion, "inclusion: intersect", "inclusion: centroid"))

tm_shape(pycno_nyc_ct_small_combinations) + tm_fill(
  fill = "pycno_populationE",
  fill.scale = tm_scale_intervals(
    values = "viridis",
    style = "jenks",
    n = 6,
    label.format = tm_label_format(big.num.abbr = c("K" = 3))
  ),
  fill_alpha = .75,
  fill.legend = tm_legend(title = "Interpolated \nPopulation", frame = FALSE)
) +
  tm_facets_grid(rows = "inclusion",
                 columns = "allocation",
                 free.coords = FALSE) +
  tm_shape(nyc_ct_small) + tm_lines(col = "grey45") +
  #tm_basemap("CartoDB.PositronNoLabels", alpha = .5) +
  tm_shape(nyc_ct) + tm_fill(fill = "grey97", zindex = 3) +
  tm_layout(panel.label.frame = FALSE, panel.label.bg = FALSE, frame = FALSE, bg.color = "grey90")


## -----------------------------------------------------------------------------
#| label: fig-pycno_nyc_ct_small_nb_order
#| fig-cap: "Interpolated population counts with increasing neighbourhood order"
#| fig-height: 8
#| fig-width: 8
#| echo: false
neighbours_list <- seq(1, 4, by = 1)

pycno_nyc_ct_small_neighbourss <- purrr::map_dfr(
  neighbours_list, ~ to_grid(
    source = nyc_ct_small,
    value_col = "populationE",
    grid_type = "h3",
    resolution = 10,
    nb_order = .x
  ) |> 
    mutate(nb_order = .x, facet_label = paste0("nb_order = ", nb_order))
)

tm_shape(pycno_nyc_ct_small_neighbourss) + tm_fill(
  fill = "pycno_populationE",
  fill.scale = tm_scale_intervals(
    values = "viridis",
    style = "jenks",
    n = 6,
    label.format = tm_label_format(big.num.abbr = c("K" = 3))
  ),
  fill_alpha = .75,
  fill.legend = tm_legend(title = "Interpolated \nPopulation", frame = FALSE),
  fill.free = FALSE
) +
  tm_facets(by = "facet_label",
            ncol = 2,
            free.coords = TRUE) +
  tm_shape(nyc_ct_small) + tm_lines(col = "grey45") +
  #tm_basemap("CartoDB.PositronNoLabels", alpha = .5) +
  tm_shape(nyc_ct) + tm_fill(fill = "grey97", zindex = 3) +
  tm_layout(panel.label.frame = FALSE, panel.label.bg = FALSE, frame = FALSE, bg.color = "grey90")


## -----------------------------------------------------------------------------
#| label: fig-nyc_ct_subset
#| fig-cap: "Source zone subset and target cells with centroid-based inclusion criteria"
#| echo: false
subset_ids <- c(
  "36061005800", 
  "36061008700",
  "36061009100"
  )

nyc_ct_subset <- nyc_ct |> 
  filter(id %in% subset_ids) |> 
  mutate(source_id = paste0("S", row_number()))

prepped_source <- pycnogrid:::prep_source(source = nyc_ct_subset,
                                          value_col = "populationE",
                                          id_col = "id")

source_polys <- prepped_source |> 
  purrr::pluck("source") |> 
  mutate(source_label_id = paste0("S", .sid))
  
source_values <- prepped_source |> 
  purrr::pluck("source_values")

input_total_original <- prepped_source |> 
  purrr::pluck("input_total_original")

input_total_represented <- prepped_source |> 
  purrr::pluck("input_total_represented")

target_cells <- pycnogrid:::prepare_target_cells(
  source_polys,
  grid_type = "h3",
  resolution = 9,
  cell_inclusion = "centroid"
  #cell_inclusion = "intersect"
) |> 
  mutate(target_label_id = paste0("T", .tid))


tm_shape(source_polys) +
  tm_fill(col = "grey35",
          col_alpha = .8,
          fill = "grey",
          fill_alpha = .1) + tm_text("source_label_id", ymod = -1, angle = 330) +
  tm_shape(target_cells) +
  tm_fill(col = "blue",
          fill = "blue",
          fill_alpha = .1) +
  tm_text("target_label_id", col = "blue") +
  #tm_basemap("CartoDB.PositronNoLabels", alpha = .5) +
  tm_shape(nyc_ct) + tm_fill(fill = "grey97", zindex = 3) +
  tm_layout(panel.label.frame = FALSE, panel.label.bg = FALSE, frame = FALSE, bg.color = "grey90")


## -----------------------------------------------------------------------------
#| echo: false
#| warning: false

# manually:
inter_raw <- sf::st_intersection(source_polys, target_cells) |>
  dplyr::mutate(.weight = as.numeric(sf::st_area(geometry))) |>
  sf::st_drop_geometry() |>
  dplyr::filter(.weight > 0) |>
  dplyr::select(source_id, .tid, .weight)

represented_sources <- inter_raw |>
  dplyr::distinct(source_id)

source_index <- source_values |>
  dplyr::mutate(.sid = dplyr::row_number()) |>
  dplyr::select(source_id, .sid)

inter <- inter_raw |>
  # semi join to drop any target cells that don't have any sources
  dplyr::semi_join(source_values, by = "source_id") |>
  # attach the numeric source index .sid.
  dplyr::left_join(source_index, by = "source_id") |>
  # Sum weights in case the same source-cell pair appears more than once,
  # which can happen after geometric operations involving multipart polygons
  # or multiple intersection fragments.
  dplyr::group_by(.sid, .tid, source_id) |>
  dplyr::summarise(.weight = sum(.weight, na.rm = TRUE), .groups = "drop")

n_target <- nrow(target_cells)
n_source <- nrow(source_values)

A_matrix <- Matrix::sparseMatrix(
  i = inter$.sid,
  j = inter$.tid,
  x = inter$.weight,
  dims = c(n_source, n_target)
)

# row and column normalize the A_matrix overlap matrix
# need some sparse matrix trickery here for row and column sums
row_sums <- Matrix::rowSums(A_matrix)
W_matrix <- Matrix::Diagonal(x = 1 / pmax(row_sums, .Machine$double.eps),
                             n = n_source) %*% A_matrix

col_sums <- as.numeric(Matrix::colSums(A_matrix))
V_matrix <- A_matrix %*% Matrix::Diagonal(x = 1 / pmax(col_sums, .Machine$double.eps),
                                          n = n_target)

# cross-reference with package function:
source_cell_weights <- pycnogrid:::prepare_source_cell_weights(
  source_polys,
  source_values,
  target_cells,
  grid_type = "h3",
  resolution = 9,
  cell_allocation = "area",
  missing_policy = "warn",
  input_total_original
)

#A_matrix == source_cell_weights$A_matrix
#W_matrix == source_cell_weights$W_matrix
#V_matrix == source_cell_weights$V_matrix


## -----------------------------------------------------------------------------
#| echo: false
y <- source_index |>
  dplyr::left_join(source_values, by = "source_id") |>
  dplyr::arrange(.sid) |>
  dplyr::pull(source_value)

# Calculate initial areal allocation
x <- as.numeric(Matrix::t(W_matrix) %*% y)


## -----------------------------------------------------------------------------
#| echo: false
density <- x / col_sums


## -----------------------------------------------------------------------------
#| echo: false
S_matrix <- pycnogrid:::build_smoothing_matrix(
  target_cells = target_cells,
  nb_order = 1,
  include_self = TRUE
  ) 


## -----------------------------------------------------------------------------
#| echo: false
d_smooth <- as.numeric(S_matrix %*% density) 


## -----------------------------------------------------------------------------
#| echo: false
source_est <- as.numeric(A_matrix %*% density)


## -----------------------------------------------------------------------------
#| echo: false
r_ratio <- y / source_est


## -----------------------------------------------------------------------------
#| echo: false
q_correction <- as.numeric(Matrix::t(V_matrix) %*% r_ratio)


## -----------------------------------------------------------------------------
#| echo: false
density_t_plus_one <- density * q_correction


## -----------------------------------------------------------------------------
#| echo: false
abs_change <- abs(density_t_plus_one - density)
denom <- pmax(abs(density), 1e-12)
last_error <- mean(abs_change / denom, na.rm = TRUE)


## -----------------------------------------------------------------------------
#| label: fig-pycno_nyc_ct_subset
#| fig-cap: "Source zone and interpolated population values"
#| fig-width: 8
#| fig-height: 4
#| echo: false
subset_results <- dplyr::bind_rows(nyc_ct_subset |> 
  to_grid(
  value_col = "populationE",
  id_col = id,
  grid_type = "h3",
  resolution = 9,
  #grid_type = "raster",
  #resolution = 275,
  cell_inclusion = "centroid",
  cell_allocation = "area",
  nb_order = 1
) |> mutate(
  label_id = paste0("T", .tid), 
  population = pycno_populationE,
  type = "target cells"
  ),
source_polys |> mutate(
  population = .source_value, 
  label_id = paste0("S", .sid),
  type = "source zones")
)

subset_overlay <- dplyr::bind_rows(
  source_polys |> mutate(
  type = "target cells"),
  target_cells |> mutate(type = "source zones")
)

tm_shape(subset_results) + tm_fill(
  fill = "population",
  fill.scale = tm_scale_ordinal(
    values = "viridis",
    label.format = tm_label_format(big.num.abbr = c("K" = 3))
  ),
  fill_alpha = .75,
  fill.legend = tm_legend(title = "Population", frame = FALSE, position = tm_pos_in()),
  fill.free = TRUE
) +
  #tm_text("label_id") +
  tm_facets(by = "type", free.coords = FALSE) +
  tm_shape(subset_overlay) +
  tm_lines(col = "grey35",
          col_alpha = .8) +
  tm_facets(by = "type", free.coords = FALSE) +
  #tm_basemap("CartoDB.PositronNoLabels", alpha = .5) +
  tm_shape(nyc_ct) + tm_fill(fill = "grey97", zindex = 3) +
  tm_layout(panel.label.frame = FALSE, panel.label.bg = FALSE, frame = FALSE, bg.color = "grey90")