library(tidyverse)
library(ggplot2)
library(worldfootballR)
library(pbapply)
library(scico)

clip_prob <- function(x) {
  pmax(0.001, pmin(0.999, x))
}

####################
# Data collection
season_years <- seq(2010, 2024, by = 1)
matches_list <- pblapply(season_years, function(year) {
  Sys.sleep(0.4) # be polite
  understat_league_match_results(
    league = "EPL",
    season_start_year = year
  )
})

matches_df <- matches_list %>%
  bind_rows() %>%
  as_tibble() %>%
  mutate(url = paste0("https://understat.com/match/", match_id))
matches_df %>%
  distinct() %>%
  write_csv("./data/football/raw/matches.csv") # remove duplicates due to API error

shots_list <- pblapply(matches_df$url, function(url) {
  Sys.sleep(0.4) # be polite
  understat_match_shots(match_url = url) %>% mutate(match_url = url)
})
shots_df <- shots_list %>%
  bind_rows() %>%
  as_tibble()
shots_df %>%
  distinct() %>%
  write_csv("./data/football/raw/shots.csv")


###############
# Dataset processing

# load the datasets
matches <- read_csv("./data/football/raw/matches.csv")
shots <- read_csv("./data/football/raw/shots.csv")


# Some basic EDA
shots_agg <- shots %>%
  mutate(X_round = round(105 * X), Y_round = round(68 * Y), shot_team = case_when(
    home_away == "h" ~ home_team,
    home_away == "a" ~ away_team
  )) %>%
  select(x = X_round, y = Y_round, shot_team, xG) %>%
  group_by(x, y) %>%
  summarise(metric = mean(xG))

response_vars <- read_csv("./data/football/processed/response_vars.csv")
shots_agg <- response_vars %>%
  mutate(X_round = round(105 * lat), Y_round = round(68 * long), Y = log(clip_prob(xG) / (1 - clip_prob(xG)))) %>%
  select(x = X_round, y = Y_round, xG, Y) %>%
  group_by(x, y) %>%
  summarise(metric = mean(Y))

ggplot() +
  geom_rect(aes(xmin = 0, xmax = 105, ymin = 0, ymax = 68), fill = "palegreen4", color = "white", linewidth = 1.2) +
  geom_tile(aes(x = x, y = y, fill = metric), data = shots_agg, alpha = 0.8) + # actual prob
  scale_fill_scico(palette = "vik", direction = 1, midpoint = 0, name = "summary") +
  geom_segment(aes(x = 52.5, y = 0, xend = 52.5, yend = 68), color = "white", linewidth = 1) + # halfway line
  ggforce::geom_circle(aes(x0 = 52.5, y0 = 34, r = 9.15), color = "white", linewidth = 1) + # center circle
  geom_rect(aes(xmin = 0, xmax = 16.5, ymin = 13.84, ymax = 54.16), fill = NA, color = "white", linewidth = 1) + # penalty boxes
  geom_rect(aes(xmin = 105 - 16.5, xmax = 105, ymin = 13.84, ymax = 54.16), fill = NA, color = "white", linewidth = 1) +
  geom_rect(aes(xmin = 0, xmax = 5.5, ymin = 27.84, ymax = 40.16), fill = NA, color = "white", linewidth = 1) + # goal area
  geom_rect(aes(xmin = 105 - 5.5, xmax = 105, ymin = 27.84, ymax = 40.16), fill = NA, color = "white", linewidth = 1) +
  coord_fixed() +
  theme(
    legend.position = "right",
    panel.background = element_rect(fill = "palegreen4", color = NA),
    axis.title = element_blank()
  )

# create motivating plot
chelsea_opp <- response_vars %>% filter(season == 2014 & opponent == "Chelsea")
chelsea_matches <- chelsea_opp %>%
  distinct(datetime, team) %>%
  filter(team %in% c("Manchester City", "Arsenal", "Liverpool")) %>%
  slice(c(1, 3, 5)) %>%
  mutate(match_num = row_number())

shot_3_matches <- chelsea_matches %>%
  inner_join(chelsea_opp, by = c("datetime", "team")) %>%
  mutate(X_round = round(105 * lat), Y_round = round(68 * long), Y = log(clip_prob(xG) / (1 - clip_prob(xG)))) %>%
  select(x = X_round, y = Y_round, xG, team)

ggplot(shot_3_matches) +
  geom_rect(aes(xmin = 0, xmax = 105, ymin = 0, ymax = 68), fill = "palegreen4", color = "white", linewidth = 1.2) +
  geom_point(aes(x = x, y = y, color = xG), size = 2) + # actual prob
  scale_color_scico(palette = "vik", direction = 1, midpoint = 0, name = "xG") +
  geom_segment(x = 52.5, y = 0, xend = 52.5, yend = 68, color = "white", linewidth = 1) + # halfway line
  ggforce::geom_circle(aes(x0 = 52.5, y0 = 34, r = 9.15), color = "white", linewidth = 1) + # center circle
  geom_rect(xmin = 0, xmax = 16.5, ymin = 13.84, ymax = 54.16, fill = NA, color = "white", linewidth = 1) + # penalty boxes
  geom_rect(xmin = 105 - 16.5, xmax = 105, ymin = 13.84, ymax = 54.16, fill = NA, color = "white", linewidth = 1) +
  geom_rect(xmin = 0, xmax = 5.5, ymin = 27.84, ymax = 40.16, fill = NA, color = "white", linewidth = 1) + # goal area
  geom_rect(xmin = 105 - 5.5, xmax = 105, ymin = 27.84, ymax = 40.16, fill = NA, color = "white", linewidth = 1) +
  coord_fixed() +
  coord_flip() +
  facet_wrap(~team) +
  theme(
    legend.position = "right",
    panel.background = element_rect(fill = "palegreen4", color = NA),
    axis.title = element_blank()
  )
ggsave("./figures/football_eda_plot.pdf", width = 8, height = 3.5)




# Model description & Details
#   Yt(s) = At match number t for teamA, at location s, logit probability of goal
#   Xt    = teamA's profile at match number t, opponent team's profile before the match

# Build every team's profile as function of season match index
team_profile <- matches %>%
  select(match_id, datetime, season, home_team, away_team, home_goals, away_goals) %>%
  pivot_longer(cols = c(home_team, away_team), names_to = "role", values_to = "team") %>%
  mutate(
    goals_for = ifelse(role == "home_team", home_goals, away_goals),
    goals_against = ifelse(role == "home_team", away_goals, home_goals),
    result = case_when(
      goals_for > goals_against ~ "win",
      goals_for < goals_against ~ "loss",
      TRUE ~ "draw"
    )
  ) %>%
  select(match_id, datetime, team, goals_for, goals_against, result) %>%
  left_join(
    shots %>%
      mutate(
        player_team = case_when(
          home_away == "h" ~ home_team,
          home_away == "a" ~ away_team
        ),
        isgoal = as.numeric(result == "Goal"),
        isgoodshot = as.numeric(result %in% c("BlockedShot", "SavedShot", "Goal"))
      ) %>%
      group_by(player_team, season, match_id, date) %>%
      summarise(
        n_shots = n(),
        goalshotprop = mean(isgoal),
        goodshotprop = mean(isgoodshot),
        expected_goal = mean(xG)
      ) %>%
      ungroup() %>%
      select(player_team, match_id, season, n_shots, goalshotprop, goodshotprop, expected_goal),
    by = c("team" = "player_team", "match_id" = "match_id")
  ) %>%
  group_by(team, season) %>%
  arrange(datetime, .by_group = TRUE) %>%
  mutate(
    n_games = row_number(), # used for time index within the season
    n_games_prop = row_number() / n(), # proportion of games already played for the season
    win_prop = cumsum(result == "win") / row_number(),
    loss_prop = cumsum(result == "loss") / row_number(),
    avg_goals_scored = cumsum(goals_for) / row_number(),
    avg_goals_conceded = cumsum(goals_against) / row_number(),
    avg_shots = cumsum(n_shots) / row_number(),
    avg_goalshots = cumsum(goalshotprop * n_shots) / cumsum(n_shots),
    avg_goodshots = cumsum(goodshotprop * n_shots) / cumsum(n_shots),
    avg_expected_goal = cumsum(expected_goal * n_shots) / cumsum(n_shots)
  ) %>%
  ungroup() %>%
  select(
    team, season, datetime,
    n_games, n_games_prop, win_prop, loss_prop, avg_goals_scored, avg_goals_conceded,
    avg_shots, avg_goalshots, avg_goodshots, avg_expected_goal
  ) %>%
  arrange(season, team, datetime)

response_vars <- shots %>%
  mutate(
    player_team = ifelse(home_away == "h", home_team, away_team),
    opponent_team = ifelse(home_away == "h", away_team, home_team),
  ) %>%
  select(season, team = player_team, datetime = date, lat = X, long = Y, xG, opponent = opponent_team) %>%
  arrange(season, team, datetime)

# save these processed data
team_profile %>% write_csv("./data/football/processed/team_profile.csv")
response_vars %>% write_csv("./data/football/processed/response_vars.csv")


###########
# Run the analysis
# choice of basis functions (1, s1, s2, s1^2, s2^2, s1 * s2)
team_profile <- read_csv("./data/football/processed/team_profile.csv")
response_vars <- read_csv("./data/football/processed/response_vars.csv")

hist(response_vars$lat) # lat is mostly between 0.6 to 1.0
hist(response_vars$long) # long is mostly between 0.3 to 0.7

# get a default average shot map
shots_agg <- response_vars %>%
  mutate(lat = round(lat, 2), long = round(long, 2), Y = log(clip_prob(xG) / (1 - clip_prob(xG)))) %>%
  group_by(lat, long, team, season) %>%
  summarise(metric = mean(Y)) %>%
  ungroup() %>%
  select(lat, long, team, season, avg_Y = metric)

mean(shots_agg$avg_Y) # -2.6136
var(shots_agg$avg_Y) # 1.5312


# team filters
team_list <- c("Arsenal", "Liverpool", "Manchester City")
response_vars <- response_vars %>% filter(team %in% team_list)

# let's filter for those
lat_bounds <- c(0.0, 1.0)
long_bounds <- c(0.0, 1.0)
response_vars <- response_vars %>%
  filter(lat >= lat_bounds[1] & lat <= lat_bounds[2] & long >= long_bounds[1] & long <= long_bounds[2]) %>%
  mutate(
    lat = round((lat - lat_bounds[1]) / (lat_bounds[2] - lat_bounds[1]), 2),
    long = round((long - long_bounds[1]) / (long_bounds[2] - long_bounds[1]), 2) # transform to 0-1 coordinate
  )

# basiswise mean computation with imputations
basis_count <- 12
create_basis_functions <- function(df) {
  df %>% mutate(
    basis_1 = 1,
    basis_2 = 2 * sqrt(3) * (lat - 1 / 2),
    basis_3 = 2 * sqrt(3) * (long - 1 / 2),
    basis_4 = 12 * (lat - 1 / 2) * (long - 1 / 2),
    basis_5 = 6 * sqrt(5) * (lat^2 - lat + 1 / 6),
    basis_6 = 6 * sqrt(5) * (long^2 - long + 1 / 6),
    basis_7 = 6 * sqrt(5) * (lat^2 - lat + 1 / 6) * (2 * sqrt(3) * (long - 1 / 2)),
    basis_8 = 6 * sqrt(5) * (long^2 - long + 1 / 6) * (2 * sqrt(3) * (lat - 1 / 2)),
    basis_9 = 2 * sqrt(7) * (lat^3 - 1.5 * lat^2 + 0.6 * lat - 1 / 20),
    basis_10 = 2 * sqrt(7) * (long^3 - 1.5 * long^2 + 0.6 * long - 1 / 20),
    basis_11 = basis_9 * (2 * sqrt(3) * (long - 1 / 2)),
    basis_12 = basis_10 * (2 * sqrt(3) * (lat - 1 / 2)),
    basisnorm = sqrt(
      basis_1^2 + basis_2^2 + basis_3^2 + basis_4^2 +
        basis_5^2 + basis_6^2 + basis_7^2 + basis_8^2 +
        basis_9^2 + basis_10^2 + basis_11^2 + basis_12^2
    )
  )
}

compute_basis_means <- function(df, team_name, season_name, grid_size = 100) {
  # bin observed locations to grid
  df_binned <- df %>%
    group_by(lat, long) %>%
    summarise(Y = mean(Y), .groups = "drop")

  # full spatial field
  grid <- expand_grid(
    lat = round(seq(0, 1, length.out = grid_size), 2),
    long = round(seq(0, 1, length.out = grid_size), 2)
  ) %>%
    as_tibble() %>%
    left_join(
      shots_agg %>% filter((team == team_name) & (season == season_name)),
      by = c("lat", "long")
    ) %>%
    mutate(
      impute_Y = ifelse(is.na(avg_Y), log(clip_prob(0) / (1 - clip_prob(0))), avg_Y)
    ) %>%
    select(lat, long, impute_Y) %>%
    create_basis_functions()

  field <- grid %>%
    left_join(df_binned, by = c("lat", "long")) %>%
    mutate(Y = if_else(is.na(Y), impute_Y, Y))

  # field <- df_binned %>%
  #   create_basis_functions()

  # compute weighted means
  field %>%
    summarise(
      across(
        starts_with("basis_"),
        ~ sum(Y * .x) / sum(abs(.x)),
        .names = "{.col}"
      )
    )
}


# takes ~ 10 seconds
Ytk <- response_vars %>%
  mutate(Y = log(clip_prob(xG) / (1 - clip_prob(xG)))) %>% # logit transformation
  group_by(season, team, datetime, opponent) %>%
  nest() %>%
  mutate(
    coefs = purrr::map(
      data,
      ~ compute_basis_means(.x, team, season)
    )
  ) %>%
  unnest(coefs) %>%
  select(-data) %>%
  ungroup()

{
  team_profile_cols <- c(
    "n_games_prop", "win_prop", "loss_prop", "avg_goals_scored", "avg_goals_conceded",
    "avg_shots", "avg_goalshots", "avg_goodshots", "avg_expected_goal"
  )
  Ytk_full <- Ytk %>%
    left_join(team_profile, by = c("season", "team", "datetime")) %>%
    group_by(season, team) %>%
    arrange(datetime) %>%
    mutate(
      across(all_of(team_profile_cols), ~ lag(., 1), .names = "team_{.col}_lag1")
    ) %>%
    ungroup() %>%
    select(
      -all_of(team_profile_cols)
    ) %>%
    arrange(season, team, datetime) %>%
    left_join(team_profile, by = c("season" = "season", "opponent" = "team", "datetime" = "datetime")) %>%
    group_by(season, opponent) %>%
    arrange(datetime) %>%
    mutate(
      across(all_of(team_profile_cols), ~ lag(., 1), .names = "opponent_{.col}_lag1")
    ) %>%
    ungroup() %>%
    select(
      -all_of(team_profile_cols)
    ) %>%
    drop_na() %>%
    arrange(season, team, datetime)


  # now run the processing and get estimates
  # normalize the covariates to a z-score for same scaling while computing distances
  logit_columns <- paste0(
    rep(c("team_", "opponent_"), each = 6),
    rep(paste0(c("n_games_prop", "win_prop", "loss_prop", "avg_goalshots", "avg_goodshots", "avg_expected_goal"), "_lag1"), 2)
  )
  log_columns <- paste0(
    rep(c("team_", "opponent_"), each = 3),
    rep(paste0(c("avg_goals_scored", "avg_goals_conceded", "avg_shots"), "_lag1"), 2)
  )

  logit_col_means <- Ytk_full %>%
    mutate(across(all_of(logit_columns), ~ log(clip_prob(.) / (1 - clip_prob(.))))) %>%
    select(all_of(logit_columns)) %>%
    colMeans(na.rm = T)
  logit_col_sds <- Ytk_full %>%
    mutate(across(all_of(logit_columns), ~ log(clip_prob(.) / (1 - clip_prob(.))))) %>%
    select(all_of(logit_columns)) %>%
    apply(FUN = function(x) {
      sd(x, na.rm = T)
    }, 2)
  log_col_means <- Ytk_full %>%
    mutate(across(all_of(log_columns), ~ log1p(.))) %>%
    select(all_of(log_columns)) %>%
    colMeans(na.rm = T)
  log_col_sds <- Ytk_full %>%
    mutate(across(all_of(log_columns), ~ log1p(.))) %>%
    select(all_of(log_columns)) %>%
    apply(FUN = function(x) {
      sd(x, na.rm = T)
    }, 2)

  Ytk_full <- Ytk_full %>%
    mutate(
      across(all_of(logit_columns), ~ scale(log(clip_prob(.) / (1 - clip_prob(.))), center = TRUE, scale = TRUE))
    ) %>%
    mutate(
      across(all_of(log_columns), ~ scale(log1p(.), center = TRUE, scale = TRUE))
    ) %>%
    select(
      all_of(c(
        "season", "team", "datetime", "opponent",
        paste0("basis_", 1:basis_count), logit_columns, log_columns
      ))
    )
}

# function that creates and outputs the surface
get_mean_prob_surface <- function(Ytk_df, covar_item, plot = TRUE, resolution_level = 50, plot_half = FALSE, plot_horiz = TRUE, plot_limits = c(0, 1)) {
  # apply the logit / log transformation based on column names
  n <- nrow(Ytk_df)
  covar_item[logit_columns] <- (log(clip_prob(covar_item[logit_columns]) / (1 - clip_prob(covar_item[logit_columns]))) - logit_col_means) / logit_col_sds
  covar_item[log_columns] <- (log(covar_item[log_columns]) - log_col_means) / log_col_sds


  # now using the kernel, find the closest matching profiles
  kernel_weights <- pbsapply(1:n, FUN = function(i) {
    d1 <- (Ytk_df[i, logit_columns] - covar_item[logit_columns])
    d2 <- (Ytk_df[i, log_columns] - covar_item[log_columns])

    # do Gaussian kernel
    return(exp(-sum(d1^2) - sum(d2^2)))
  })

  # create the weighted mean, for basis vectors
  mu_k <- (crossprod(as.matrix(Ytk_df[, paste0("basis_", 1:basis_count)]), kernel_weights) / sum(kernel_weights))[, 1]
  mu_k

  # sigma computation
  Ek <- as.matrix(Ytk_df[, paste0("basis_", 1:basis_count)] - matrix(mu_k, nrow = nrow(Ytk_df), ncol = basis_count, byrow = TRUE))
  sig_kl <- (t(Ek) %*% diag(kernel_weights) %*% Ek) / sum(kernel_weights)

  # obtain the football grid, predict and get the probability surface
  grid <- expand.grid(
    lat = seq(0, 1, length.out = resolution_level),
    long = seq(0, 1, length.out = resolution_level)
  ) %>%
    as_tibble() %>%
    mutate(
      x = 105 * ((lat_bounds[2] - lat_bounds[1]) * lat + lat_bounds[1]),
      y = 68 * ((long_bounds[2] - long_bounds[1]) * long + long_bounds[1])
    ) %>% # map to real-world meters
    create_basis_functions()

  grid$mu_logit <- as.vector(as.matrix(grid[, paste0("basis_", 1:basis_count)]) %*% mu_k) # / grid$basisnorm

  # scale to proper range
  data_range_param <- c(-2.6136, 1.5312)
  grid$mu_logit <- (grid$mu_logit - mean(grid$mu_logit)) * sqrt(data_range_param[2]) / sd(grid$mu_logit) + data_range_param[1]

  grid$sigma_logit <- as.vector(diag(as.matrix(grid[, paste0("basis_", 1:basis_count)]) %*% sig_kl %*% t(as.matrix(grid[, paste0("basis_", 1:basis_count)]))))
  grid$pred <- exp(grid$mu_logit) / (1 + exp(grid$mu_logit))
  grid$tscore <- grid$mu_logit / sqrt(grid$sigma_logit)

  # if (min(grid$tscore) < -10) {
  #   # scale the z-scores for plotting
  #   grid$tscore <- grid$tscore / min(grid$tscore) * (-10)
  # }

  if (plot) {
    # make the plot

    p <- ggplot() +
      geom_rect(aes(xmin = 0, xmax = 105, ymin = 0, ymax = 68), fill = "palegreen4", color = "white", linewidth = 1.2) +
      geom_raster(aes(x = x, y = y, fill = pred), data = grid, alpha = 0.8, interpolate = TRUE) + # actual prob
      scale_fill_scico(palette = "vik", direction = 1, name = "Predicted xG", limits = plot_limits) +
      geom_segment(aes(x = 52.5, y = 0, xend = 52.5, yend = 68), color = "white", linewidth = 1) + # halfway line
      ggforce::geom_circle(aes(x0 = 52.5, y0 = 34, r = 9.15), color = "white", linewidth = 1) + # center circle
      geom_rect(aes(xmin = 0, xmax = 16.5, ymin = 13.84, ymax = 54.16), fill = NA, color = "white", linewidth = 1) + # penalty boxes
      geom_rect(aes(xmin = 105 - 16.5, xmax = 105, ymin = 13.84, ymax = 54.16), fill = NA, color = "white", linewidth = 1) +
      geom_rect(aes(xmin = 0, xmax = 5.5, ymin = 27.84, ymax = 40.16), fill = NA, color = "white", linewidth = 1) + # goal area
      geom_rect(aes(xmin = 105 - 5.5, xmax = 105, ymin = 27.84, ymax = 40.16), fill = NA, color = "white", linewidth = 1) +
      theme(
        legend.position = "right",
        panel.background = element_rect(fill = "palegreen4", color = NA),
        axis.title = element_blank()
      )
    if (plot_half) {
      p <- p + coord_fixed(xlim = c(52.5, 105))
    } else {
      p <- p + coord_fixed()
    }

    if (!plot_horiz) {
      p <- p + coord_flip()
    }

    return(list("surface" = grid, "plot" = p))
  } else {
    return(list("surface" = grid))
  }
}


# pick up specific situation, for that, obtain the nonparametric estimate
# Situation 1: Mid season, strong attack vs weak defense (general)
covar_item <- c(
  team_n_games_prop_lag1 = 0.50,
  team_win_prop_lag1 = 0.70, team_loss_prop_lag1 = 0.15, team_avg_goalshots_lag1 = 0.25, team_avg_goodshots_lag1 = 0.70,
  team_avg_expected_goal_lag1 = 0.18, team_avg_goals_scored_lag1 = 2.20, team_avg_goals_conceded_lag1 = 0.90, team_avg_shots_lag1 = 16.0,
  opponent_n_games_prop_lag1 = 0.50,
  opponent_win_prop_lag1 = 0.30, opponent_loss_prop_lag1 = 0.45, opponent_avg_goalshots_lag1 = 0.10, opponent_avg_goodshots_lag1 = 0.40,
  opponent_avg_expected_goal_lag1 = 0.08, opponent_avg_goals_scored_lag1 = 0.80, opponent_avg_goals_conceded_lag1 = 1.90, opponent_avg_shots_lag1 = 9.0
)

output <- get_mean_prob_surface(Ytk_full, covar_item, plot_half = FALSE, plot_limits = c(0, 0.5))
output$plot


# Situation 2: Mid season, weak attack vs strong defense
# covar_item <- c(
#   team_n_games_prop_lag1 = 0.50,
#   team_win_prop_lag1 = 0.30, team_loss_prop_lag1 = 0.45, team_avg_goalshots_lag1 = 0.10, team_avg_goodshots_lag1 = 0.40,
#   team_avg_expected_goal_lag1 = 0.08, team_avg_goals_scored_lag1 = 0.80, team_avg_goals_conceded_lag1 = 1.90, team_avg_shots_lag1 = 9.0,
#   opponent_n_games_prop_lag1 = 0.50,
#   opponent_win_prop_lag1 = 0.70, opponent_loss_prop_lag1 = 0.15, opponent_avg_goalshots_lag1 = 0.25, opponent_avg_goodshots_lag1 = 0.70,
#   opponent_avg_expected_goal_lag1 = 0.18, opponent_avg_goals_scored_lag1 = 2.20, opponent_avg_goals_conceded_lag1 = 0.90, opponent_avg_shots_lag1 = 16.0
# )


# filter for team, season, etc. and build the relevant plots
plot_list <- list()
surface_list <- list()
for (team_name in team_list) {
  plot_row_list <- list()
  surface_row_list <- list()
  for (i in 1:3) {
    season_bounds <- list(c(2014, 2017), c(2018, 2021), c(2022, 2025))[[i]]

    # build median profile for team in mid-season
    team_median_profile <- team_profile %>%
      filter(team == team_name & n_games_prop == 0.5 & season >= season_bounds[1] & season <= season_bounds[2]) %>%
      select(-c(team, season, datetime, n_games)) %>%
      apply(FUN = function(x) {
        median(x, na.rm = T)
      }, 2)
    attacker_profile <- team_median_profile

    # build median profile for opponent in mid-season
    opponent_median_profile <- team_profile %>%
      filter(team_name %in% team_list) %>%
      filter(team != team_name & n_games_prop == 0.5 & season >= season_bounds[1] & season <= season_bounds[2]) %>%
      select(-c(team, season, datetime, n_games)) %>%
      apply(FUN = function(x) {
        median(x, na.rm = T)
      }, 2)
    defender_profile <- opponent_median_profile

    names(attacker_profile) <- paste0("team_", names(attacker_profile), "_lag1")
    names(defender_profile) <- paste0("opponent_", names(defender_profile), "_lag1")
    covar_item <- c(attacker_profile, defender_profile)

    output <- get_mean_prob_surface(
      Ytk_full %>% filter(team == team_name & season >= season_bounds[1] & season <= season_bounds[2]),
      covar_item,
      plot_half = FALSE,
      plot_limits = c(0, 0.55)
    )

    season_name <- paste0(season_bounds[1], "-", season_bounds[2])
    plot_row_list[[season_name]] <- output$plot
    
    # add to the surface
    season_label <- c("2017/18", "2021/22", "2024/25")[i]
    surface_row_list[[season_label]] <- output$surface
    
  }
  plot_list[[team_name]] <- plot_row_list
  surface_list[[team_name]] <- surface_row_list
}


# make the plot using patchwork
library(patchwork)

teams <- names(plot_list)
seasons <- names(plot_list[[teams[1]]])

lapply(teams, function(team) {
  # combine all seasons for that team horizontally
  reduce(plot_list[[team]], `+`) +
    plot_layout(ncol = length(seasons))
}) %>% reduce(`/`) +
  plot_layout(nrow = length(teams), guides = "collect") & # collect legends
  plot_annotation(
    # title = "Shot Efficiency Surfaces by Team and Season",
    theme = theme(
      plot.title = element_text(hjust = 0.5),
      legend.position = "right"
    ),
    tag_levels = "A"
  )
ggsave("./figures/football_estimated_xG.pdf", width = 8, height = 7)


# make a facet like plot

pred_df_combined <- list()
i <- 0
for (team_name in names(surface_list)) {
  for (season in names(surface_list[[team_name]])) {
    dat <- surface_list[[team_name]][[season]]
    dat <- dat %>%
      select(lat, long, x, y, pred, tscore) %>%
      mutate(
        team = team_name,
        season = season
      )
    i <- i + 1
    pred_df_combined[[i]] <- dat
  }
}
preddf <- bind_rows(pred_df_combined)


ggplot(preddf) +
  geom_rect(aes(xmin = 0, xmax = 105, ymin = 0, ymax = 68), fill = "palegreen4", color = "white", linewidth = 1.2) +
  geom_raster(aes(x = x, y = y, fill = pred), alpha = 0.8, interpolate = TRUE) + # actual prob
  scale_fill_scico(palette = "vik", direction = 1, name = "Predicted xG", 
                   limits = c(0, 0.55)) +
  geom_segment(aes(x = 52.5, y = 0, xend = 52.5, yend = 68), color = "white", linewidth = 1) + # halfway line
  ggforce::geom_circle(aes(x0 = 52.5, y0 = 34, r = 9.15), color = "white", linewidth = 1) + # center circle
  geom_rect(aes(xmin = 0, xmax = 16.5, ymin = 13.84, ymax = 54.16), fill = NA, color = "white", linewidth = 1) + # penalty boxes
  geom_rect(aes(xmin = 105 - 16.5, xmax = 105, ymin = 13.84, ymax = 54.16), fill = NA, color = "white", linewidth = 1) +
  geom_rect(aes(xmin = 0, xmax = 5.5, ymin = 27.84, ymax = 40.16), fill = NA, color = "white", linewidth = 1) + # goal area
  geom_rect(aes(xmin = 105 - 5.5, xmax = 105, ymin = 27.84, ymax = 40.16), fill = NA, color = "white", linewidth = 1) +
  facet_grid(rows = vars(team), cols = vars(season), switch = "y") +
  coord_fixed() +
  theme(
    legend.position = "right",
    panel.background = element_rect(fill = "palegreen4", color = NA),
    axis.title = element_blank()
  )
ggsave("./figures/football_estimated_xG.pdf", width = 8, height = 7)






