# ============================================================
# Cluster analysis of 12 k-root series in Slavic languages
# ============================================================
#
# Input:
#   semantic_maps_2_k_root.xlsx
#
# Place this script and the Excel file in the same directory
# and run the script from that directory.
#
# Required R packages:
#   readxl, dplyr, tidyr, ggplot2, stringr, cluster
#
# The script:
#   1. builds semantic and distributional profiles,
#   2. performs hierarchical clustering (Euclidean distance,
#      Ward.D2 method),
#   3. evaluates k using mean silhouette width,
#   4. reports the best-supported cluster solution,
#   5. compares the k = 2 solution with the preliminary
#      two-subgroup classification using the Adjusted Rand Index.
#
# Output is written to the folder "cluster_analysis_results".
# ============================================================

library(readxl)
library(dplyr)
library(tidyr)
library(ggplot2)
library(stringr)
library(cluster)


# ------------------------------------------------------------
# Dirs
# ------------------------------------------------------------

data_file <- "semantic_maps_2_k_root.xlsx"
output_dir <- "cluster_analysis_results"


dir.create(
  output_dir,
  recursive = TRUE,
  showWarnings = FALSE
)


# ------------------------------------------------------------
# Clean text imported from Excel
# ------------------------------------------------------------

clean_text <- function(x) {
  x %>%
    as.character() %>%
    str_replace_all("\u00A0", " ") %>%
    str_squish()
}


# ------------------------------------------------------------
# Read data
# ------------------------------------------------------------

series <- read_excel(
  data_file,
  sheet = "series"
) %>%
  mutate(
    language = clean_text(language),
    series_id = clean_text(series_id),
    display_name = clean_text(display_name),
    sample_n = as.numeric(sample_n),
    
    preliminary_group = if_else(
      series_id %in% c("PL_06", "PL_10", "CZ_10"),
      "Group 1",
      "Group 2"
    ),
    
    series_label = paste0(
      language,
      " ",
      display_name,
      " [",
      series_id,
      "]"
    )
  )


semantics <- read_excel(
  data_file,
  sheet = "semantic_counts"
) %>%
  mutate(
    language = clean_text(language),
    series_id = clean_text(series_id),
    display_name = clean_text(display_name),
    semantic_function = clean_text(semantic_function),
    count = as.numeric(count)
  ) %>%
  left_join(
    series %>%
      select(
        language,
        series_id,
        sample_n,
        series_label
      ),
    by = c(
      "language",
      "series_id"
    )
  ) %>%
  mutate(
    rate = count / sample_n
  )


contexts <- read_excel(
  data_file,
  sheet = "context_counts"
) %>%
  mutate(
    language = clean_text(language),
    series_id = clean_text(series_id),
    display_name = clean_text(display_name),
    context = clean_text(context),
    count = as.numeric(count)
  ) %>%
  left_join(
    series %>%
      select(
        language,
        series_id,
        sample_n,
        series_label
      ),
    by = c(
      "language",
      "series_id"
    )
  ) %>%
  mutate(
    rate = count / sample_n
  )


stopifnot(!anyNA(semantics$sample_n))
stopifnot(!anyNA(contexts$sample_n))


# ------------------------------------------------------------
# matrices
# ------------------------------------------------------------

make_matrix <- function(data, variable) {
  
  x <- data %>%
    select(
      series_label,
      {{ variable }},
      rate
    ) %>%
    group_by(
      series_label,
      {{ variable }}
    ) %>%
    summarise(
      rate = sum(
        rate,
        na.rm = TRUE
      ),
      .groups = "drop"
    ) %>%
    pivot_wider(
      names_from = {{ variable }},
      values_from = rate,
      values_fill = 0
    )
  
  mat <- as.matrix(
    x[, -1]
  )
  
  rownames(mat) <- x$series_label
  
  mat
}


semantic_matrix <- make_matrix(
  semantics,
  semantic_function
)

context_matrix <- make_matrix(
  contexts,
  context
)


stopifnot(nrow(semantic_matrix) == 12)
stopifnot(nrow(context_matrix) == 12)


cat(
  "\nSemantic matrix:",
  nrow(semantic_matrix),
  "series x",
  ncol(semantic_matrix),
  "functions\n"
)

cat(
  "Distribution matrix:",
  nrow(context_matrix),
  "series x",
  ncol(context_matrix),
  "contexts\n"
)


# ------------------------------------------------------------
# Euclidean dist/Ward clustering
# ------------------------------------------------------------

semantic_dist <- dist(
  semantic_matrix,
  method = "euclidean"
)

context_dist <- dist(
  context_matrix,
  method = "euclidean"
)


semantic_hc <- hclust(
  semantic_dist,
  method = "ward.D2"
)

context_hc <- hclust(
  context_dist,
  method = "ward.D2"
)


# ------------------------------------------------------------
# Silhouette analysis
# ------------------------------------------------------------

get_silhouettes <- function(
    hc,
    distances,
    analysis
) {
  
  bind_rows(
    lapply(
      2:5,
      function(k) {
        
        clusters <- cutree(
          hc,
          k = k
        )
        
        sil <- silhouette(
          clusters,
          distances
        )
        
        tibble(
          analysis = analysis,
          k = k,
          mean_silhouette = mean(
            sil[, "sil_width"]
          )
        )
      }
    )
  )
}


silhouette_results <- bind_rows(
  
  get_silhouettes(
    semantic_hc,
    semantic_dist,
    "Semantics"
  ),
  
  get_silhouettes(
    context_hc,
    context_dist,
    "Distribution"
  )
)


cat(
  "\nSilhouette results:\n"
)

print(
  silhouette_results
)


# ------------------------------------------------------------
# Best number of clusters
# ------------------------------------------------------------

best_k <- silhouette_results %>%
  group_by(
    analysis
  ) %>%
  slice_max(
    mean_silhouette,
    n = 1,
    with_ties = FALSE
  ) %>%
  ungroup()


cat(
  "\nBest number of clusters:\n"
)

print(
  best_k
)


semantic_best_k <- best_k %>%
  filter(
    analysis == "Semantics"
  ) %>%
  pull(
    k
  )

distribution_best_k <- best_k %>%
  filter(
    analysis == "Distribution"
  ) %>%
  pull(
    k
  )


# ------------------------------------------------------------
# Cluster membership
# ------------------------------------------------------------

semantic_best <- cutree(
  semantic_hc,
  k = semantic_best_k
)

distribution_best <- cutree(
  context_hc,
  k = distribution_best_k
)


# k = 2 is needed for comparison with the preliminary two subgroup classification

semantic_k2 <- cutree(
  semantic_hc,
  k = 2
)

distribution_k2 <- cutree(
  context_hc,
  k = 2
)


cluster_membership <- series %>%
  select(
    series_id,
    series_label,
    preliminary_group
  ) %>%
  mutate(
    
    semantic_cluster = semantic_best[
      match(
        series_label,
        names(semantic_best)
      )
    ],
    
    distribution_cluster = distribution_best[
      match(
        series_label,
        names(distribution_best)
      )
    ],
    
    semantic_cluster_k2 = semantic_k2[
      match(
        series_label,
        names(semantic_k2)
      )
    ],
    
    distribution_cluster_k2 = distribution_k2[
      match(
        series_label,
        names(distribution_k2)
      )
    ]
  )


cat(
  "\nCluster membership:\n"
)

print(
  cluster_membership,
  n = Inf
)


# ------------------------------------------------------------
# Adjusted Rand Index
# ------------------------------------------------------------

adjusted_rand_index <- function(x, y) {
  
  tab <- table(
    x,
    y
  )
  
  choose2 <- function(n) {
    n * (n - 1) / 2
  }
  
  observed <- sum(
    choose2(tab)
  )
  
  rows <- sum(
    choose2(
      rowSums(tab)
    )
  )
  
  cols <- sum(
    choose2(
      colSums(tab)
    )
  )
  
  total <- choose2(
    sum(tab)
  )
  
  expected <- rows * cols / total
  
  maximum <- (
    rows + cols
  ) / 2
  
  denominator <- maximum - expected
  
  if (denominator == 0) {
    return(NA_real_)
  }
  
  (
    observed - expected
  ) / denominator
}


semantic_groups <- series$preliminary_group[
  match(
    names(semantic_k2),
    series$series_label
  )
]

distribution_groups <- series$preliminary_group[
  match(
    names(distribution_k2),
    series$series_label
  )
]


semantic_ari <- adjusted_rand_index(
  semantic_k2,
  semantic_groups
)

distribution_ari <- adjusted_rand_index(
  distribution_k2,
  distribution_groups
)


cat(
  "\nAdjusted Rand Index for k = 2:\n",
  "Semantics:     ",
  round(semantic_ari, 3),
  "\n",
  "Distribution:  ",
  round(distribution_ari, 3),
  "\n"
)


# ------------------------------------------------------------
# Silhouette plots
# ------------------------------------------------------------

semantic_plot <- silhouette_results %>%
  filter(
    analysis == "Semantics"
  ) %>%
  ggplot(
    aes(
      x = k,
      y = mean_silhouette
    )
  ) +
  geom_line() +
  geom_point(
    size = 2.5
  ) +
  scale_x_continuous(
    breaks = 2:5
  ) +
  labs(
    title = "Semantics",
    x = "Number of clusters",
    y = "Mean silhouette width"
  ) +
  theme_classic()


distribution_plot <- silhouette_results %>%
  filter(
    analysis == "Distribution"
  ) %>%
  ggplot(
    aes(
      x = k,
      y = mean_silhouette
    )
  ) +
  geom_line() +
  geom_point(
    size = 2.5
  ) +
  scale_x_continuous(
    breaks = 2:5
  ) +
  labs(
    title = "Distribution",
    x = "Number of clusters",
    y = "Mean silhouette width"
  ) +
  theme_classic()


print(
  semantic_plot
)

print(
  distribution_plot
)


# ------------------------------------------------------------
# Save results
# ------------------------------------------------------------

write.csv(
  cluster_membership,
  file.path(
    output_dir,
    "cluster_membership.csv"
  ),
  row.names = FALSE
)


write.csv(
  silhouette_results,
  file.path(
    output_dir,
    "silhouette_results.csv"
  ),
  row.names = FALSE
)


write.csv(
  best_k,
  file.path(
    output_dir,
    "best_k.csv"
  ),
  row.names = FALSE
)


ari_results <- tibble(
  analysis = c(
    "Semantics",
    "Distribution"
  ),
  ARI_k2 = c(
    semantic_ari,
    distribution_ari
  )
)


write.csv(
  ari_results,
  file.path(
    output_dir,
    "ari_results.csv"
  ),
  row.names = FALSE
)


ggsave(
  filename = file.path(
    output_dir,
    "silhouette_semantics.png"
  ),
  plot = semantic_plot,
  width = 5,
  height = 4,
  dpi = 300,
  bg = "white"
)


ggsave(
  filename = file.path(
    output_dir,
    "silhouette_distribution.png"
  ),
  plot = distribution_plot,
  width = 5,
  height = 4,
  dpi = 300,
  bg = "white"
)


cat(
  "\nAnalysis finished.\n",
  "Results saved in:\n",
  output_dir,
  "\n"
)