library(testthat)
library(dplyr)
library(stringr)

test_that("HLA_mismatched_alleles correctly identifies mismatched alleles", {
  GL_string_recip <- "HLA-A*01:01+HLA-A*02:01^HLA-B*07:02+HLA-B*08:01^HLA-DRB1*13:01+HLA-DRB1*15:01"
  GL_string_donor <- "HLA-A*01:01+HLA-A*03:01^HLA-B*07:02+HLA-B*44:02^HLA-DRB1*13:01+HLA-DRB1*14:01"
  loci <- c("HLA-A", "HLA-B", "HLA-DRB1")

  result_HvG <- HLA_mismatched_alleles(GL_string_recip, GL_string_donor, loci, direction = "HvG")
  result_GvH <- HLA_mismatched_alleles(GL_string_recip, GL_string_donor, loci, direction = "GvH")
  result_bidirectional <- HLA_mismatched_alleles(GL_string_recip, GL_string_donor, loci, direction = "bidirectional")
  result_SOT <- HLA_mismatched_alleles(GL_string_recip, GL_string_donor, loci, direction = "SOT")

  expect_equal(result_HvG, "HLA-A=HLA-A*03:01, HLA-B=HLA-B*44:02, HLA-DRB1=HLA-DRB1*14:01")
  expect_equal(result_GvH, "HLA-A=HLA-A*02:01, HLA-B=HLA-B*08:01, HLA-DRB1=HLA-DRB1*15:01")
  expect_equal(result_SOT, "HLA-A=HLA-A*03:01, HLA-B=HLA-B*44:02, HLA-DRB1=HLA-DRB1*14:01")
  expect_equal(result_bidirectional, "HvG;HLA-A=HLA-A*03:01, HLA-B=HLA-B*44:02, HLA-DRB1=HLA-DRB1*14:01<>GvH;HLA-A=HLA-A*02:01, HLA-B=HLA-B*08:01, HLA-DRB1=HLA-DRB1*15:01")

  result_single_locus <- HLA_mismatched_alleles(GL_string_recip, GL_string_donor, "HLA-A", direction = "HvG")
  expect_equal(result_single_locus, "HLA-A*03:01")

  result_homozygous_1 <- HLA_mismatched_alleles(GL_string_recip, GL_string_donor, "HLA-B", direction = "HvG", homozygous_count = 1)
  expect_equal(result_homozygous_1, "HLA-B*44:02")

  result_homozygous_2 <- HLA_mismatched_alleles(GL_string_recip, GL_string_donor, "HLA-B", direction = "HvG", homozygous_count = 2)
  expect_equal(result_homozygous_2, "HLA-B*44:02")

  expect_error(HLA_mismatched_alleles(GL_string_recip, GL_string_donor, loci, direction = "invalid"))
})


# --- Input validation tests ---

test_that("HLA_mismatched_alleles rejects NULL inputs", {
  gl <- "HLA-A*01:01+HLA-A*02:01"
  expect_error(HLA_mismatched_alleles(NULL, gl, "HLA-A", "HvG"), "GL_string_recip")
  expect_error(HLA_mismatched_alleles(gl, NULL, "HLA-A", "HvG"), "GL_string_donor")
  expect_error(HLA_mismatched_alleles(gl, gl, NULL, "HvG"), "loci")
})

test_that("HLA_mismatched_alleles rejects non-character GL strings", {
  gl <- "HLA-A*01:01+HLA-A*02:01"
  expect_error(HLA_mismatched_alleles(123, gl, "HLA-A", "HvG"), "must be a character")
})

test_that("HLA_mismatched_alleles rejects invalid homozygous_count", {
  gl <- "HLA-A*01:01+HLA-A*02:01"
  expect_error(HLA_mismatched_alleles(gl, gl, "HLA-A", "HvG", homozygous_count = 0), "must be.*1.*or.*2")
})

test_that("HLA_mismatched_alleles returns NA for perfect match at single locus", {
  gl <- "HLA-A*01:01+HLA-A*02:01^HLA-B*07:02+HLA-B*08:01"
  result <- HLA_mismatched_alleles(gl, gl, "HLA-A", "HvG")
  expect_true(is.na(result))
})

test_that("HLA_mismatched_alleles agrees with mismatch_table_2010 consensus (homozygous_count = 1)", {
  # Define symbolic alleles to substitute into the table.
  A <- "HLA-A*01:01"
  B <- "HLA-A*02:05"
  C <- "HLA-A*24:02"
  D <- "HLA-A*31:03"
  N <- "HLA-A*68:11N"

  # Build GL strings from the 2010 table's symbolic allele pairs.
  mismatch_table_2010_alleles <- mismatch_table_2010 %>%
    rename(Patient = patient, Donor = donor) %>%
    # Assign the example alleles based on symbolic codes.
    mutate(recipient_1 = case_when(
      str_detect(Patient, "^A") ~ A,
      str_detect(Patient, "^N") ~ N
    ), recipient_2 = case_when(
      str_detect(Patient, "A$") ~ A,
      str_detect(Patient, "B$") ~ B,
      str_detect(Patient, "N$") ~ N
    ), donor_1 = case_when(
      str_detect(Donor, "^A") ~ A,
      str_detect(Donor, "^B") ~ B,
      str_detect(Donor, "^C") ~ C,
      str_detect(Donor, "^N") ~ N
    ), donor_2 = case_when(
      str_detect(Donor, "A$") ~ A,
      str_detect(Donor, "B$") ~ B,
      str_detect(Donor, "C$") ~ C,
      str_detect(Donor, "D$") ~ D,
      str_detect(Donor, "N$") ~ N
    )) %>%
    # Turn the two alleles into a GL string.
    mutate(
      GL_string_recip = str_c(recipient_1, recipient_2, sep = "+"),
      GL_string_donor = str_c(donor_1, donor_2, sep = "+")
    ) %>%
    select(Patient, Donor, GL_string_recip, GL_string_donor, "#GvH", "#HvG") %>%
    rename(GvH_count = "#GvH", HvG_count = "#HvG") %>%
    # Calculate mismatched alleles for each direction with homozygous_count = 1.
    mutate(
      GvH_alleles = HLA_mismatched_alleles(GL_string_recip, GL_string_donor, "HLA-A", "GvH", homozygous_count = 1),
      HvG_alleles = HLA_mismatched_alleles(GL_string_recip, GL_string_donor, "HLA-A", "HvG", homozygous_count = 1)
    ) %>%
    # Verify: NA when mismatch count is 0, non-NA when count > 0.
    mutate(
      GvH_na_correct = (is.na(GvH_alleles) == (GvH_count == 0)),
      HvG_na_correct = (is.na(HvG_alleles) == (HvG_count == 0))
    ) %>%
    # Verify: number of alleles returned matches the mismatch count.
    # Count alleles by splitting on "+" delimiter.
    mutate(
      GvH_n_alleles = if_else(is.na(GvH_alleles), 0L, as.integer(str_count(GvH_alleles, "\\+") + 1)),
      HvG_n_alleles = if_else(is.na(HvG_alleles), 0L, as.integer(str_count(HvG_alleles, "\\+") + 1))
    ) %>%
    mutate(
      GvH_count_correct = (GvH_n_alleles == GvH_count),
      HvG_count_correct = (HvG_n_alleles == HvG_count)
    ) %>%
    # Check that all validations pass for every row.
    mutate(total_result = if_all(c(GvH_na_correct, HvG_na_correct, GvH_count_correct, HvG_count_correct))) %>%
    distinct(total_result) %>%
    pull(total_result)

  expect_equal(mismatch_table_2010_alleles, TRUE)
})

test_that("HLA_mismatched_alleles agrees with mismatch_table_2016 consensus (homozygous_count = 2)", {
  # Define symbolic alleles to substitute into the table.
  A <- "HLA-A*01:01"
  B <- "HLA-A*02:05"
  C <- "HLA-A*24:02"
  D <- "HLA-A*31:03"
  N <- "HLA-A*68:11N"

  # Build GL strings from the 2016 table's symbolic allele pairs.
  mismatch_table_2016_alleles <- mismatch_table_2016 %>%
    # Assign the example alleles based on symbolic codes.
    mutate(recipient_1 = case_when(
      str_detect(Patient, "^A") ~ A,
      str_detect(Patient, "^N") ~ N
    ), recipient_2 = case_when(
      str_detect(Patient, "A$") ~ A,
      str_detect(Patient, "B$") ~ B,
      str_detect(Patient, "N$") ~ N
    ), donor_1 = case_when(
      str_detect(Donor, "^A") ~ A,
      str_detect(Donor, "^B") ~ B,
      str_detect(Donor, "^C") ~ C,
      str_detect(Donor, "^N") ~ N
    ), donor_2 = case_when(
      str_detect(Donor, "A$") ~ A,
      str_detect(Donor, "B$") ~ B,
      str_detect(Donor, "C$") ~ C,
      str_detect(Donor, "D$") ~ D,
      str_detect(Donor, "N$") ~ N
    )) %>%
    # Turn the two alleles into a GL string.
    mutate(
      GL_string_recip = str_c(recipient_1, recipient_2, sep = "+"),
      GL_string_donor = str_c(donor_1, donor_2, sep = "+")
    ) %>%
    select(Patient, Donor, GL_string_recip, GL_string_donor, "#GvH", "#HvG") %>%
    rename(GvH_count = "#GvH", HvG_count = "#HvG") %>%
    # Calculate mismatched alleles for each direction with homozygous_count = 2 (default).
    mutate(
      GvH_alleles = HLA_mismatched_alleles(GL_string_recip, GL_string_donor, "HLA-A", "GvH", homozygous_count = 2),
      HvG_alleles = HLA_mismatched_alleles(GL_string_recip, GL_string_donor, "HLA-A", "HvG", homozygous_count = 2)
    ) %>%
    # Verify: NA when mismatch count is 0, non-NA when count > 0.
    mutate(
      GvH_na_correct = (is.na(GvH_alleles) == (GvH_count == 0)),
      HvG_na_correct = (is.na(HvG_alleles) == (HvG_count == 0))
    ) %>%
    # Verify: number of alleles returned matches the mismatch count.
    # Count alleles by splitting on "+" delimiter.
    mutate(
      GvH_n_alleles = if_else(is.na(GvH_alleles), 0L, as.integer(str_count(GvH_alleles, "\\+") + 1)),
      HvG_n_alleles = if_else(is.na(HvG_alleles), 0L, as.integer(str_count(HvG_alleles, "\\+") + 1))
    ) %>%
    mutate(
      GvH_count_correct = (GvH_n_alleles == GvH_count),
      HvG_count_correct = (HvG_n_alleles == HvG_count)
    ) %>%
    # Check that all validations pass for every row.
    mutate(total_result = if_all(c(GvH_na_correct, HvG_na_correct, GvH_count_correct, HvG_count_correct))) %>%
    distinct(total_result) %>%
    pull(total_result)

  expect_equal(mismatch_table_2016_alleles, TRUE)
})

test_that("HLA_mismatched_alleles works with vectorized inputs", {
  recip <- c(
    "HLA-A*01:01+HLA-A*02:01",
    "HLA-A*03:01+HLA-A*24:02"
  )
  donor <- c(
    "HLA-A*01:01+HLA-A*03:01",
    "HLA-A*03:01+HLA-A*24:02"
  )
  result <- HLA_mismatched_alleles(recip, donor, "HLA-A", "HvG")
  expect_length(result, 2)
  # First pair has a mismatch
  expect_equal(result[1], "HLA-A*03:01")
  # Second pair is a perfect match
  expect_true(is.na(result[2]))
})
