# ============================================================
#  PRDM9 Red Queen Dynamics
#  Motif depletion leads to frequency-dependent selection
#  Two PRDM9 alleles competing for motif availability
# ============================================================

# ------------------------------------------------------------
# 1.  BIOLOGICAL SETUP
# ------------------------------------------------------------

# PRDM9 alleles:
#   A1: binds motif type X (depletes X over time)
#   A2: binds motif type Y (depletes Y over time)

# Motif availability:
#   X(t): frequency of intact motif X sites  [0,1]
#   Y(t): frequency of intact motif Y sites  [0,1]

# Fitness logic:
#   High motif availability → successful crossover initiation → normal meiosis → fitness = 1
#   Low motif availability → failed crossover → meiotic failure → fitness ≈ 0
#   Fitness depends on YOUR motif type availability, not opponent frequency

x0    <- 0.7     # initial frequency of PRDM9 allele A1
X0    <- 0.8     # initial availability of motif X
Y0    <- 0.8     # initial availability of motif Y
t_max <- 150     # generations

# Motif depletion rates
depl_rate_X <- 0.02   # rate at which A1 depletes motif X
depl_rate_Y <- 0.02   # rate at which A2 depletes motif Y
recovery_rate <- 0.005 # slow motif recovery through mutation

# Fitness functions — sigmoid collapse when motifs become rare
# k controls steepness, M_half = motif level where fitness = 0.5
k_steep <- 15
M_half  <- 0.4

fitness_A1 <- function(X) 1 / (1 + exp(-k_steep * (X - M_half)))
fitness_A2 <- function(Y) 1 / (1 + exp(-k_steep * (Y - M_half)))

# Selection strength
delta <- 0.8

# ------------------------------------------------------------
# 2.  UPDATE FUNCTIONS
# ------------------------------------------------------------

# PRDM9 frequency update (standard replicator)
replicator_PRDM9 <- function(x, X, Y, delta) {
  f_A1  <- fitness_A1(X)
  f_A2  <- fitness_A2(Y)
  f_bar <- x * f_A1 + (1 - x) * f_A2
  
  x_new <- x * (1 - delta + delta * f_A1) /
               (1 - delta + delta * f_bar)
  max(1e-6, min(1 - 1e-6, x_new))
}

# Motif availability update (depletion by PRDM9 usage + slow recovery)
motif_step <- function(x, X, Y) {
  # A1 depletes X motifs proportional to its frequency
  X_new <- X - depl_rate_X * x + recovery_rate * (1 - X)
  
  # A2 depletes Y motifs proportional to its frequency  
  Y_new <- Y - depl_rate_Y * (1 - x) + recovery_rate * (1 - Y)
  
  X_new <- max(0, min(1, X_new))
  Y_new <- max(0, min(1, Y_new))
  
  c(X_new, Y_new)
}

# ------------------------------------------------------------
# 3.  SIMULATE
# ------------------------------------------------------------

x <- numeric(t_max + 1); x[1] <- x0
X <- numeric(t_max + 1); X[1] <- X0
Y <- numeric(t_max + 1); Y[1] <- Y0

for (t in seq_len(t_max)) {
  x[t + 1] <- replicator_PRDM9(x[t], X[t], Y[t], delta)
  motifs   <- motif_step(x[t], X[t], Y[t])
  X[t + 1] <- motifs[1]
  Y[t + 1] <- motifs[2]
}

# Build results dataframe with fitness calculations
sim <- data.frame(
  time     = 0:t_max,
  freq_A1  = x,
  freq_A2  = 1 - x,
  motif_X  = X,
  motif_Y  = Y,
  fit_A1   = fitness_A1(X),
  fit_A2   = fitness_A2(Y)
)

# Population average fitness (weighted by frequencies)
sim$pop_fitness <- sim$freq_A1 * sim$fit_A1 + sim$freq_A2 * sim$fit_A2

# ------------------------------------------------------------
# 4.  PLOT
# ------------------------------------------------------------

col_A1    <- "#4E9AF1"  # blue for A1
col_A2    <- "#F17A4E"  # red for A2  
col_motX  <- "#A8D8A8"  # green for motif X
col_motY  <- "#DDA0DD"  # purple for motif Y
col_pop   <- "#2E2E2E"  # dark for population fitness

par(mfrow = c(3, 1),
    mar   = c(4, 6.5, 3.5, 2),
    mgp   = c(4.8, 1.2, 0),
    oma   = c(1, 0, 2, 0))

# ---- Panel 1: PRDM9 allele frequencies ----
plot(sim$time, sim$freq_A1,
     type = "l", lwd = 4, col = col_A1,
     ylim = c(0, 1),
     xlab = "", ylab = "PRDM9 allele frequency",
     main = "PRDM9 allele dynamics",
     cex.main = 1.4, cex.lab = 1.4, cex.axis = 1.2,
     axes = FALSE)
axis(1, lwd = 2.5, lwd.ticks = 2.5, cex.axis = 1.2)
axis(2, lwd = 2.5, lwd.ticks = 2.5, cex.axis = 1.2, las = 1)
box(lwd = 2.5)

lines(sim$time, sim$freq_A2, lwd = 4, col = col_A2)
abline(h = 0.5, lty = 3, lwd = 1.5, col = "grey70")

legend("right", 
       legend = c("PRDM9-A1 (binds motif X)", "PRDM9-A2 (binds motif Y)"),
       col = c(col_A1, col_A2), lwd = 4,
       bty = "n", cex = 1.1, xpd = NA)

# ---- Panel 2: motif availability ----
plot(sim$time, sim$motif_X,
     type = "l", lwd = 4, col = col_motX,
     ylim = c(0, 1),
     xlab = "", ylab = "Motif availability",
     main = "Motif depletion & recovery",
     cex.main = 1.4, cex.lab = 1.4, cex.axis = 1.2,
     axes = FALSE)
axis(1, lwd = 2.5, lwd.ticks = 2.5, cex.axis = 1.2)
axis(2, lwd = 2.5, lwd.ticks = 2.5, cex.axis = 1.2, las = 1)
box(lwd = 2.5)

lines(sim$time, sim$motif_Y, lwd = 4, col = col_motY)
abline(h = M_half, lty = 2, lwd = 2, col = "grey60")

legend("right", 
       legend = c("Motif X availability", "Motif Y availability", "Fitness threshold"),
       col = c(col_motX, col_motY, "grey60"), lty = c(1, 1, 2), lwd = c(4, 4, 2),
       bty = "n", cex = 1.1, xpd = NA)

# ---- Panel 3: fitness over time ----
plot(sim$time, sim$fit_A1,
     type = "l", lwd = 4, col = col_A1,
     ylim = c(0, 1),
     xlab = "Generation", ylab = "Fitness",
     main = "Meiotic success & population fitness",
     cex.main = 1.4, cex.lab = 1.4, cex.axis = 1.2,
     axes = FALSE)
axis(1, lwd = 2.5, lwd.ticks = 2.5, cex.axis = 1.2)
axis(2, lwd = 2.5, lwd.ticks = 2.5, cex.axis = 1.2, las = 1)
box(lwd = 2.5)

lines(sim$time, sim$fit_A2, lwd = 4, col = col_A2)
lines(sim$time, sim$pop_fitness, lwd = 4, col = col_pop, lty = 3)
abline(h = c(0.5, 1), lty = 3, lwd = 1.5, col = "grey70")

legend("right", 
       legend = c("Fitness PRDM9-A1", "Fitness PRDM9-A2", "Population avg fitness"),
       col = c(col_A1, col_A2, col_pop), lty = c(1, 1, 3), lwd = 4,
       bty = "n", cex = 1.1, xpd = NA)

mtext("PRDM9 Red Queen: Hotspot Self-Destruction Model",
      outer = TRUE, cex = 1.4, font = 2, line = 0.5)

par(mfrow = c(1,1), mar = c(5.1,4.1,4.1,2.1),
    mgp = c(3,1,0), oma = c(0,0,0,0))

# ------------------------------------------------------------
# 5.  SUMMARY
# ------------------------------------------------------------

cat("\n=== PRDM9 Red Queen Analysis ===\n")
cat("Mechanism: Each PRDM9 allele depletes its own binding motifs\n")
cat(sprintf("Motif depletion rate: %.3f per generation\n", depl_rate_X))
cat(sprintf("Motif recovery rate:  %.3f per generation\n", recovery_rate))
cat(sprintf("Fitness threshold:    %.2f motif availability\n", M_half))

cat("\nDynamics:\n")
cat(sprintf("Initial: A1 = %.2f, motif X = %.2f, motif Y = %.2f\n", x0, X0, Y0))
cat(sprintf("Final:   A1 = %.3f, motif X = %.3f, motif Y = %.3f\n", 
            tail(sim$freq_A1, 1), tail(sim$motif_X, 1), tail(sim$motif_Y, 1)))
cat(sprintf("Min population fitness: %.3f (due to meiotic failure)\n", min(sim$pop_fitness)))
cat(sprintf("Max population fitness: %.3f\n", max(sim$pop_fitness)))

# Find cycles
motif_X_range <- max(sim$motif_X) - min(sim$motif_X)
cat(sprintf("Motif X oscillation range: %.3f\n", motif_X_range))