library(tidyverse)
library(httr2)
library(rvest)
library(knitr)

if (!str_ends(getwd(), "posts/food")) {
  setwd("posts/food")
}

download_files <- function() {
  url <- "https://data.food.gov.uk/catalog/datasets/0bfd916a-4e01-4cb8-ba16-763f0b36b50c"

  html_content <- request(url) |>
    req_perform() |>
    resp_body_string()

  # Assuming html_content is already defined
  doc <- read_html(html_content)
  xlsx_links <- doc %>% html_nodes("a[href$='.xlsx']") %>% html_attr("href")
  for (link in xlsx_links) {
    file_url <- link
    file_name <- basename(file_url)
    destfile <- sprintf("data/%s", file_name)
    if (!file.exists(destfile)) {
      download.file(url = file_url, destfile = destfile)
    }
  }
}

download_files()

library(tidyverse)
library(readxl)
library(knitr)

extract_month_year <- function(file) {
  decoded_filename <- URLdecode(file)

  month_year_regex <- "^.*(January|February|March|April|May|June|July|August|September|October|November|December)\\s?(\\d{2}|\\d{4}).*.xlsx$"

  if (grepl(month_year_regex, decoded_filename)) {
    extracted_month_year <- regmatches(
      decoded_filename,
      regexec(month_year_regex, decoded_filename)
    )

    # Extract month and year
    month <- extracted_month_year[[1]][[2]]
    year <- extracted_month_year[[1]][[3]]

    if (str_length(year) == 2) {
      year <- paste0("20", year)
    }

    return(list(filename = file, month = month, year = year))
  } else {
    print(file)
    return(NULL)
  }
}

extract_question <- function(file, question_prefix) {
  fmy <- extract_month_year(file)

  data <- read_excel(
    file,
    sheet = "Percents",
    col_names = FALSE
  )
  headers <- as.vector(data[5, ], mode = "character")
  headers[1] <- "QCategory"
  headers <- zoo::na.locf(headers)
  subheaders <- as.vector(data[6, ], mode = "character")
  subheaders <- zoo::na.fill(subheaders, "")
  combined_headers <- map2(
    headers,
    subheaders,
    \(x, y) str_c(x, y, sep = ":: ")
  )
  combined_headers[1] <- "QCategory"
  combined_headers[2] <- "Total"

  colnames(data) <- combined_headers

  has_question <- any(str_starts(pull(data, QCategory), question_prefix))

  if (!is.na(has_question)) {
    data |>
      filter(!is.na(QCategory)) |>
      slice(
        which(str_starts(QCategory, question_prefix)):n()
      ) |>
      head(10) |> # take the cats
      tail(-1) |> # drop the header
      mutate(
        filename = fmy$filename,
        monthname = fmy$month,
        year = as.integer(fmy$year),
        month = match(
          monthname,
          c(
            "January",
            "February",
            "March",
            "April",
            "May",
            "June",
            "July",
            "August",
            "September",
            "October",
            "November",
            "December"
          )
        )
      )
  }
}

question_df <- function(question_prefix) {
  df_filename <- sprintf("data/%s.rds", question_prefix)
  if (file.exists(df_filename)) {
    df <- read_rds(df_filename)
  } else {
    xlsx_files <- list.files(
      path = "data",
      pattern = "*.xlsx",
      full.names = TRUE
    )
    df <- xlsx_files |>
      sort(decreasing = TRUE) |>
      lapply(\(f) extract_question(f, question_prefix)) |>
      keep(\(x) !is.null(x)) |>
      bind_rows() |>
      mutate(Question = question_prefix)

    df |> write_rds(df_filename)
  }
  df
}

df <- question_df("Q12_14")

tidy_df <- df |>
  filter(QCategory != "Unweighted base" & QCategory != "Base: All") |>
  pivot_longer(
    c(Total, contains("::")),
    names_to = "Demographic",
    values_to = "Value"
  ) |>
  mutate(
    Period = as.Date(sprintf("%04d-%02d-01", year, month)),
    Response = QCategory,
    Value = as.numeric(Value)
  ) |>
  select(Question, Demographic, Period, Response, Value) |>
  arrange(Question, Period, Response, Demographic)

tidy_df |> head() |> kable()

library(patchwork)

RESPONSE_CATS <- c(
  "Highly concerned",
  "Somewhat concerned",
  "Not very concerned",
  "Not concerned at all",
  "Don't know"
)

net_concern <- tidy_df |>
  filter(Demographic == "Total" & str_starts(Response, "Net: "))

breakdown <- tidy_df |>
  filter(Demographic == "Total" & !str_starts(Response, "Net: ")) |>
  mutate(
    Response = if_else(Response %in% RESPONSE_CATS, Response, "Don't know"),
    Response = factor(Response, levels = RESPONSE_CATS, ordered = TRUE)
  ) |>
  # Re-normalise after removing the IDK cat
  group_by(Question, Demographic, Period) |>
  mutate(Value = Value / sum(Value)) |>
  ungroup()

plot_net_concern <- ggplot(net_concern) +
  aes(x = Period, y = Value) +
  scale_y_continuous(limits = c(0, 1)) +
  geom_point() +
  stat_smooth(method = "lm") +
  labs(title = "Net concern", x = NULL, y = "Proportion of respondents")

plot_breakdown <- ggplot(breakdown) +
  aes(x = Period, y = Value, group = Response, colour = Response) +
  geom_point() +
  stat_smooth(method = "lm") +
  labs(title = "Response breakdown", x = NULL, y = "Proportion of respondents")

plot_net_concern + plot_breakdown

source("articles.R")

upf_articles <- load_upf_articles()
domain_counts <- load_domain_counts()

ggplot(
  upf_articles |>
    mutate(
      mention_type = if_else(is_focus, "Focus article", "Mentions UPFs")
    )
) +
  aes(x = period, group = is_focus) +
  geom_histogram(stat = "count") +
  facet_wrap(~mention_type) +
  labs(
    title = "Articles mentioning ultra-processed food from top UK sources, by month.",
    x = NULL,
    y = "Article count"
  )

mention_counts <- ts(
  (group_by(upf_articles, period) |> count())$n,
  start = c(2024, 1),
  deltat = 1 / 12
)
afit <- forecast::auto.arima(mention_counts, d = 1, D = 1)
plot(forecast::forecast(afit, h = 5, level = 90))

ggplot(
  upf_articles |> group_by(domain, period) |> count()
) +
  aes(x = period, y = n, group = domain, colour = domain) +
  facet_wrap(~domain, nrow = 2) +
  geom_point(show.legend = F) +
  geom_line(linewidth = 1, show.legend = F)

ggplot(
  group_by(domain_counts, domain) |> summarise(n = sum(n_articles))
) +
  aes(
    y = fct_reorder(domain, n),
    x = n,
    group = domain,
    colour = domain,
    fill = domain
  ) +
  geom_col(linewidth = 1, show.legend = F)

readership <- tibble(
  domain = factor(DOMAINS, levels = DOMAINS, ordered = TRUE),
  monthly_readers = 1000000 *
    c(59, 19.7, 12.3, 8.3, 5.9, 5.7, 4.6, 2.6, 2.5, 1.3) # BBC estimated from OFCOM survey
) |>
  merge(
    domain_counts |>
      group_by(domain, period) |>
      summarise(n_total_articles = sum(n_articles)),
    all.x = TRUE
  ) |>
  merge(
    upf_articles |> group_by(domain, period) |> summarise(n_upf_articles = n()),
    all.x = TRUE
  ) |>
  mutate(
    n_upf_articles = coalesce(n_upf_articles, 0),
    prob_upf = n_upf_articles / n_total_articles,
    # Number of articles in the time period, calibrated such
    # that the expected number of BBC articles is approx 1/day.
    # This is an educated guess, not much research to back it I'm afraid. Pew found in 2015 that
    # "An overwhelming majority of both long-form readers (72%) and short-form readers (79%)
    # view just one article on a given site over the course of a month on their cellphone."
    # - https://www.pewresearch.org/journalism/2016/05/05/long-form-reading-shows-signs-of-life-in-our-mobile-news-world/
    # That was specific to the heady earlier days of the mobile web though.
    # I found various industry-backed studies suggesting it could be higher, such as this NewsWorks study that
    # found young people reading 6/day, however this is really suspicious and they have an
    # obvious incentive to inflate the numbers.
    # https://pressgazette.co.uk/media-audience-and-business-data/media_metrics/young-people-news/
    lambda = 30.4 * monthly_readers / max(monthly_readers),
  )

set.seed(42)

sample_article_counts <- function(month_readership, n_sim_readers) {
  result <- matrix(
    ncol = length(DOMAINS),
    nrow = n_sim_readers,
    dimnames = list(1:n_sim_readers, DOMAINS)
  )
  for (i in 1:nrow(month_readership)) {
    source <- as.list(month_readership[i, ])
    n <- rpois(n_sim_readers, source$lambda * source$prob_upf)
    result[, i] <- n
  }
  result
}

PERIODS <- seq.Date(
  from = as.Date("2024-01-01"),
  to = as.Date("2025-05-01"),
  by = "1 month"
)

N_SIM_READERS = 10000

sim_counts <- PERIODS |>
  map(
    function(p) {
      sample_article_counts(
        month_readership = filter(readership, period == p),
        n_sim_readers = N_SIM_READERS
      )
    }
  )

upf_readership_over_time <- map2(
  PERIODS,
  monthly_sim_count_totals,
  function(p, t) {
    tibble(
      period = p,
      pc_sim_readers_reading_at_least_one_article = 100 *
        (N_SIM_READERS - t[[1]]) /
        N_SIM_READERS
    )
  }
) |>
  bind_rows()

ggplot(upf_readership_over_time) +
  aes(x = period, y = pc_sim_readers_reading_at_least_one_article) +
  geom_line() +
  scale_y_continuous(
    limits = c(0, 20),
    breaks = seq(0, 20, 10),
    minor_breaks = seq(0, 20, 5)
  ) +
  labs(
    x = NULL,
    y = "Percent of simulated readers",
  )

#|
readership |>
  group_by(domain) |>
  summarise(
    lambda = mean(lambda),
    prob_upf = mean(prob_upf),
    monthly_readers = mean(monthly_readers),
    monthly_total_articles = mean(n_total_articles),
    monthly_upf_articles = mean(n_upf_articles),
    exp_monthly_upf_articles = mean(lambda * prob_upf)
  ) |>
  mutate(
    exp_monthly_upf_articles = exp_monthly_upf_articles /
      sum(exp_monthly_upf_articles)
  ) |>
  arrange(desc(exp_monthly_upf_articles)) |>
  kable(
    digits = c(0, 1, 4, 0, 0, 0, 2),
    format.args = list(decimal.mark = ".", big.mark = ",")
  )

sample_index <- expand.grid(
  period = 1:length(PERIODS),
  reader = 1:N_SIM_READERS,
  domain = 1:length(DOMAINS)
)

sampled_articles <- 1:nrow(sample_index) |>
  map(
    function(i) {
      p <- sample_index[i, 1]
      r <- sample_index[i, 2]
      s <- sample_index[i, 3]

      n <- sim_counts[[p]][r, s]

      if (n > 0) {
        upf_articles |>
          filter(period == PERIODS[p] & domain == DOMAINS[s]) |>
          slice_sample(n = n, replace = FALSE) |>
          mutate(period = PERIODS[p], reader = r, n = n)
      }
    },
    .progress = F
  ) |>
  bind_rows()

print(
  sprintf(
    "Percent of readers who have read at least one UPF article in the total period: %.1f%%",
    100 * n_distinct(sampled_articles$reader) / N_SIM_READERS
  )
)

[1] "Percent of readers who have read at least one UPF article in the total period: 46.8%"

# This is the system prompt, which initiates every conversation by giving the AI a persona and instructions.
SYSTEM_PROMPT <- "You are a member of the British public, with your own life experiences and opinions.

This is your identity: 
<identity>
{{simulated_identity}}
</identity>

You are being asked for your views as part of a survey. Respond to each question from the pollster."

# This is the question precisely as asked in the YouGov survey.
UPF_QUESTION <- "At the moment, how concerned, if at all, do you personally feel about ultra-processed, or over-processing of food?"

# This is how we wrap the question with relevant articles (and a date, for context).
QUESTION_ARTICLE_WRAPPER <- "Answer the question, recalling that you have seen the following articles that may or may not have influenced your opinion.

<articles>
{{headlines}}
</articles>

(Today's date is {{date_today}}.)

QUESTION: {{question}}"

library(ellmer)

create_respondent <- function(
  simulated_identity = "Average person"
) {
  chat_openai(
    base_url = "http://0.0.0.0:10240/v1/",
    # Gemma is good
    #model = "mlx-community/gemma-3-1b-it-4bit-DWQ",
    #model = "mlx-community/gemma-3-27b-it-4bit-DWQ",
    #model = "mlx-community/gemma-3-27b-it-qat-8bit",
    # Qwen is great
    #model = "mlx-community/Qwen3-0.6B-4bit-DWQ-053125",
    #model = "mlx-community/Qwen3-1.7B-4bit-DWQ-053125",
    model = "mlx-community/Qwen3-4B-4bit-DWQ-053125",
    #model = "mlx-community/Qwen3-14B-4bit-DWQ-053125",
    # Uncensored?
    #model = "mlx-community/Josiefied-Qwen3-0.6B-abliterated-v1-4bit",
    # GPT for debugging, do not use with data
    #model = "gpt-4.1-nano",
    api_key = "n/a",
    system_prompt = interpolate(SYSTEM_PROMPT),
    params = params(
      # Single output (all the local server supports)
      n = 1,
      # Get max available log probs
      log_probs = TRUE,
      top_logprobs = 5,
      # Keep the full distribution for sampling
      top_p = 1,
      # Do not adjust the distribution, so no cascading randomness
      temperature = 0
    ),
    api_args = list(
      # Ensure we disable reasoning from Qwen
      enable_thinking = FALSE
    )
  )
}

extract_answer_dist <- function(respondent, choices) {
  rc <- tibble(
    token = LETTERS[1:length(choices)],
    Response = factor(choices, levels = choices, ordered = TRUE)
  )

  token <- respondent$last_turn()@json$choices[[1]]$logprobs$content |>
    keep(~ str_detect(.x$token, paste(rc$token, collapse = "|"))) |>
    first()

  if (is.null(token)) {
    stop(
      respondent$last_turn()@json$choices[[1]]
    )
  }

  token$top_logprobs |>
    bind_rows() |>
    select(token, logprob) |>
    unique() |>
    merge(rc) |>
    mutate(Probability = exp(logprob)) |>
    select(Response, Probability, logprob) |>
    arrange(Response)
}

question_with_articles <- function(question, articles, date_today = today()) {
  headlines <- articles |>
    map(function(a) {
      with(
        a,
        {
          excerpts <- paste(
            map(highlights, ~ sprintf("> %s", .)),
            collapse = "\n"
          )
          sprintf(
            "HEADLINE: %s (%s, %s)\nEXCERPTS:\n%s\n\n",
            title,
            domain,
            publish_date,
            excerpts
          )
        }
      )
    })

  interpolate(QUESTION_ARTICLE_WRAPPER)
}


multiple_choice_question <- function(question, choices) {
  paste(
    c(
      question,
      "",
      imap(choices, ~ sprintf("%s) %s", LETTERS[.y], .x)),
      "",
      "Respond with the letter of your chosen answer."
    ),
    collapse = "\n"
  )
}

ask <- function(respondent, question, choices) {
  mcq <- multiple_choice_question(question, choices)
  tryCatch(
    {
      respondent$chat_structured(
        mcq,
        type = type_enum(
          description = "Answer letter",
          values = LETTERS[1:length(choices)]
        )
      )
      extract_answer_dist(respondent, choices)
    },
    error = function(e) {
      print(respondent$last_turn()@json)
      stop(e)
    }
  )
}

default_response <- ask(
  create_respondent("Average person"),
  UPF_QUESTION,
  RESPONSE_CATS
)
default_response |> kable(digits = 3)

fake_articles <- list(
  list(
    title = "Ultra-processed food linked to potential bowel cancer risk, scientists say",
    domain = "dailymail.co.uk",
    publish_date = as.Date("2025-05-01"),
    highlights = c(
      "Ultra-processed foods increase the cancer risk by 3%.",
      "Maybe we should be concerned about ultra-processed foods."
    )
  )
)

ask(
  create_respondent("Average person"),
  question_with_articles(UPF_QUESTION, fake_articles),
  RESPONSE_CATS
) |>
  kable(digits = 3)

ask(
  create_respondent("Someone very concerned about their health"),
  question_with_articles(UPF_QUESTION, fake_articles),
  RESPONSE_CATS
) |>
  kable(digits = 3)

gender_resp <- c("Female", "Male") |>
  map(
    ~ ask(
      create_respondent(.),
      UPF_QUESTION,
      RESPONSE_CATS
    ) |>
      mutate(simulated_identity = .)
  ) |>
  bind_rows()

# gender_resp |> mutate(Net = Response %in% c("Highly concerned", "Somewhat concerned")) |> group_by(simulated_identity, Net) |> summarise(p = sum(Probability))

gender_resp |>
  pivot_wider(
    id_cols = Response,
    names_from = simulated_identity,
    values_from = Probability
  ) |>
  kable(digits = 3)

gender_breakdown <- tidy_df |>
  filter(str_starts(Demographic, "Gender") & !str_starts(Response, "Net: ")) |>
  mutate(
    Response = if_else(Response %in% RESPONSE_CATS, Response, "Don't know"),
    Response = factor(Response, levels = RESPONSE_CATS, ordered = TRUE),
    Demographic = str_replace(Demographic, "Gender:: ", "")
  ) |>
  # Re-normalise after removing the IDK cat
  group_by(Question, Demographic, Period) |>
  mutate(Value = Value / sum(Value)) |>
  ungroup()

plot_gender_dist <- ggplot(
  gender_breakdown |>
    group_by(Demographic, Response) |>
    summarise(Value = mean(Value)) |>
    mutate(Value = if_else(Demographic == "Male", Value, -Value))
) +
  aes(x = Value, y = fct_rev(Response), fill = Demographic) +
  scale_fill_manual(values = c("Male" = "orange", "Female" = "purple")) +
  geom_col() +
  labs(
    title = "Gender response distribution",
    x = "Mean proportion of respondents",
    y = NULL
  )

plot_gender_time <- ggplot(gender_breakdown) +
  aes(x = Period, y = Value, group = Response, colour = Response) +
  facet_wrap(~Demographic) +
  geom_point() +
  stat_smooth(method = "lm") +
  labs(
    title = "Response breakdown over time",
    x = NULL,
    y = "Proportion of respondents"
  )

plot_gender_dist + plot_gender_time + plot_layout(nrow = 2, ncol = 1)

library(furrr)

samples_per_reader <- sampled_articles |>
  group_by(period, reader) |>
  summarize(
    sample_key = str_flatten_comma(sort(article_id)),
    article_ids = list(sort(article_id)),
    n_article = n_distinct(article_id),
  )

unique_samples <- samples_per_reader |>
  distinct(period, sample_key, article_ids)

sim_ident <- "Average person"

simulate_reader <- function(row) {
  tryCatch(
    {
      arts <- upf_articles |> filter(article_id %in% row$article_ids)
      dist <- ask(
        create_respondent(sim_ident),
        question_with_articles(
          UPF_QUESTION,
          pmap(arts, list),
          date_today = row$period
        ),
        RESPONSE_CATS
      )
      dist |>
        mutate(sample_key = row$sample_key, simulated_identity = sim_ident)
    },
    error = \(e) print(e, row)
  )
}

map_simulated_readers <- function(unique_samples, sim_ident, force = FALSE) {
  datafile <- sprintf("data/article-responses - %s.rds", sim_ident)

  if (!file.exists(datafile) || force) {
    with(plan(multisession, workers = 4), {
      tasks <- future_map(
        pmap(unique_samples, list),
        simulate_reader,
        .progress = TRUE
      )
    })

    article_responses <- bind_rows(tasks) |> as_tibble()
    article_responses |> write_rds(datafile)
  }
  read_rds(datafile)
}

article_responses <- map_simulated_readers(
  unique_samples,
  sim_ident,
  #force = TRUE
)

sim_results <- samples_per_reader |> merge(article_responses)

period_sim_dists <- matrix(nrow = length(PERIODS), ncol = length(RESPONSE_CATS))

for (i in seq_along(PERIODS)) {
  period_sim_results <- sim_results |> filter(period == PERIODS[i])
  n_non_readers <- N_SIM_READERS - n_distinct(period_sim_results$reader)

  non_reader_dist <- n_non_readers * default_response$Probability
  reader_dist <- (period_sim_results |>
    group_by(Response) |>
    summarise(p = sum(Probability)))$p
  total_dist <- non_reader_dist + reader_dist
  period_sim_dists[i, ] <- total_dist / sum(total_dist)
}

rownames(period_sim_dists) <- as.character(PERIODS)
colnames(period_sim_dists) <- RESPONSE_CATS

sim_breakdown <- as_tibble(period_sim_dists) |>
  mutate(Period = PERIODS) |>
  pivot_longer(-c(Period), names_to = "Response", values_to = "Value") |>
  mutate(Response = factor(Response, levels = RESPONSE_CATS, ordered = TRUE))

sim_net_concern <- sim_breakdown |>
  mutate(
    Concerned = Response %in% c("Highly concerned", "Somewhat concerned")
  ) |>
  group_by(Period, Concerned) |>
  summarise(Value = sum(Value)) |>
  filter(Concerned)

plot_actual_breakdown <- ggplot(breakdown) +
  aes(x = Period, y = Value, group = Response, colour = Response) +
  geom_point(show.legend = FALSE) +
  stat_smooth(method = "lm", show.legend = FALSE) +
  labs(
    title = "Actual response breakdown",
    x = NULL,
    y = "Proportion of respondents"
  )

plot_sim_breakdown <- ggplot(sim_breakdown) +
  aes(x = Period, y = Value, group = Response, colour = Response) +
  geom_point() +
  stat_smooth(method = "lm") +
  labs(
    title = "Simulated response breakdown",
    x = NULL,
    y = "Proportion of simulated respondents"
  )

plot_actual_breakdown + plot_sim_breakdown

ggplot(
  article_responses |>
    group_by(Response) |>
    summarise(Probability = mean(Probability))
) +
  aes(y = forcats::fct_rev(Response), x = Probability, fill = Response) +
  geom_col(show.legend = F) +
  scale_x_continuous(limits = c(0, 1)) +
  labs(x = "Mean probability across all articles", y = "Response")

ggplot(
  sim_results |>
    rename(Period = period) |>
    group_by(Period, Response) |>
    summarise(Value = sum(Probability) / N_SIM_READERS)
) +
  aes(x = Period, y = Value, group = Response, colour = Response) +
  geom_point() +
  stat_smooth() +
  labs(x = NULL, y = "Proportion of all respondents")