library(tidyverse)
library(stringr)
library(scales)
library(svglite)
library(ggplot2)
library(jsonlite)
library(xtable)

data_baseline <- (
  read_csv('./data/synthetic.baseline.csv', quote = "'") %>%
  mutate(collection = 'synthetic')
) %>%
  mutate(
    total_dups = sapply(
      map(duplicates, function(dups) { sum(fromJSON(dups)$cls) }),
     "[[",
     1
    ),
    len = n_axioms + n_derived,
    mus = str_detect(formula, 'mus'),
    instance_class = sub('-.*', '', formula),
    clique = str_detect(formula, 'clique'),
    full_instance_class = ifelse(
      instance_class == 'coloring',
                          ifelse(clique, 'coloring-clique', 'coloring'),
      instance_class
    )
  )

stopifnot(data_baseline %>% tally(is.na(len)) == 0)
stopifnot(data_baseline %>% tally(n_rat > 0) == 0)

dir.create('./svg', showWarnings = FALSE)

p <- ggplot(
  data_baseline %>%
  filter(trim & total_dups > 0 & !(instance_class %in% c('parity', 'php', 'op'))),
  aes(x = total_dups / (len + total_dups), fill = full_instance_class)
) +
  geom_histogram(color = 'white', binwidth = 5e-3) +
  scale_x_continuous(limits = c(0, 0.15), labels = scales::percent) +
  labs(x = 'Relative number of duplicates',
       y = 'Number of samples',
       fill = 'Instance class') +
  theme(legend.position = c(.8, .7), legend.background=element_blank())
ggsave('svg/relative-dups.svg', plot = p, width = 4, height = 4)

data_baseline_clean <- data_baseline %>% filter(trim) %>%
  select(-trim, -n_rat, -duplicates, -total_dups) |>
  left_join(
    data_baseline %>% filter(!trim) %>% mutate(solver_len = len) %>%
      select(formula, collection, solver_len),
  )

# Synthetic data report

data <- read_csv('./data/synthetic.csv') |>
  left_join(data_baseline_clean) %>%
  mutate(gap_to_mus_bound = ifelse(mus & best_proof == best_bound,
                                   best_proof - (2 * n_axioms - 1), NA))

## Parametric formulae

smu_report <- data %>% filter(instance_class == 'smu' & solver != 'proofmin-len') %>%
  mutate(n_vars = as.numeric(sub('-.*', '', sub('smu-', '', formula)))) %>%
  group_by(solver, n_vars) %>%
  summarize(
    n_formulas = n(),
    n_solved = sum(best_proof == best_bound),
    res_number = max(best_proof),
    total_time = sum(time)
  ) %>% ungroup
stopifnot(smu_report %>% tally(n_formulas != n_solved) == 0)

dir.create('./tex', showWarnings = FALSE)

smu_report %>%
  pivot_wider(names_from = n_vars, values_from = total_time, id_cols = c('solver')) %>%
  mutate(solver = ifelse(solver == 'proofmin-opt',
                         'Total duration, our approach (s)',
                         'Total duration, baseline (s)')) %>%
  rename("Number of clauses $m$" = solver,
         "9*" = `9`) %>%
  xtable(type = 'latex', label = "table:hardness-number-time",
         caption = "Total execution time for solving the proof length minimization problem for $m$-clause saturated minimally unsatisfiable formulas. The column with an asterisk corresponds to an \\emph{incomplete} set of formulas.") %>%
  print(file = 'tex/hardness-number-time.tex', include.rownames = FALSE)

data %>% filter(instance_class == 'php' & solver == 'proofmin-len') %>% # HACK proofmin-len
  select(formula, best_proof, best_bound, len, peak_memory) %>%
  mutate(param = as.numeric(sub('\\..*', '', sub('php-', '', formula)))) %>%
  arrange(param) %>%
  mutate(best_proof = ifelse(!is.na(best_bound) & best_proof == best_bound,
                        paste(best_proof, "*", sep = ""), best_proof),
         len = paste(len),
	 peak_memory = paste(floor(peak_memory / 1024))) %>%
  pivot_longer(c(best_proof, len, peak_memory)) %>%
  select(param, name, value) %>%
  pivot_wider(names_from = param, values_from = value) %>%
  mutate(name = ifelse(name == 'best_proof',
                       'Proof length, our approach',
		ifelse(name == 'len',
                       'Proof length, CaDiCaL + DRAT-trim', 'Peak memory usage, megabytes'))) %>%
  rename("Number of pigeons $n$" = name) %>%
  xtable(type = 'latex', label = "table:parametric-time-php",
         caption = "Proof lengths for pigeon-hole principle formulas. Values with asterisks are additionally proven to be optimal.") %>%
  print(file = 'tex/php.tex', include.rownames = FALSE)

data %>% filter(instance_class == 'parity' & solver == 'proofmin-len') %>% # HACK proofmin-len
  select(formula, best_proof, best_bound, len, peak_memory) %>%
  mutate(param = as.numeric(sub('\\..*', '', sub('parity-', '', formula)))) %>%
  arrange(param) %>%
  mutate(best_proof = ifelse(!is.na(best_bound) & best_proof == best_bound,
                        paste(best_proof, "*", sep = ""), best_proof),
         len = paste(len),
	 peak_memory = paste(floor(peak_memory / 1024))) %>%
  pivot_longer(c(best_proof, len, peak_memory)) %>%
  select(param, name, value) %>%
  pivot_wider(names_from = param, values_from = value) %>%
  mutate(name = ifelse(name == 'best_proof',
                       'Proof length, our approach',
		ifelse(name == 'len',
                       'Proof length, CaDiCaL + DRAT-trim', 'Peak memory usage, megabytes'))) %>%
  rename("Number of elements $n$" = name) %>%
  xtable(type = 'latex', label = "table:parametric-time-parity",
         caption = "Proof lengths for parity principle formulas. Values with asterisks are additionally proven to be optimal.") %>%
  print(file = 'tex/parity.tex', include.rownames = FALSE)

data %>% filter(instance_class == 'op' & (!mus | formula == 'op-1.mus.cnf') & solver == 'proofmin-len') %>% # HACK proofmin-len
  select(formula, best_proof, best_bound, len, peak_memory) %>%
  mutate(param = as.numeric(sub('\\..*', '', sub('op-', '', formula)))) %>%
  arrange(param) %>%
  mutate(best_proof = ifelse(!is.na(best_bound) & best_proof == best_bound,
                        paste(best_proof, "*", sep = ""), best_proof),
         len = paste(len),
	 peak_memory = paste(floor(peak_memory / 1024))) %>%
  pivot_longer(c(best_proof, len, peak_memory)) %>%
  select(param, name, value) %>%
  pivot_wider(names_from = param, values_from = value) %>%
  mutate(name = ifelse(name == 'best_proof',
                       'Proof length, our approach',
		ifelse(name == 'len',
                       'Proof length, CaDiCaL + DRAT-trim', 'Peak memory usage, megabytes'))) %>%
  rename("Number of elements $n$" = name) %>%
  xtable(type = 'latex', label = "table:parametric-time-op",
         caption = "Proof lengths for ordering principle formulas. Values with asterisks are additionally proven to be optimal.") %>%
  print(file = 'tex/op.tex', include.rownames = FALSE)

data %>% filter(instance_class == 'op' & (mus | formula == 'op-1.mus.cnf') & solver == 'proofmin-len') %>% # HACK proofmin-len
  select(formula, best_proof, best_bound, len, peak_memory) %>%
  mutate(param = as.numeric(sub('\\..*', '', sub('op-', '', formula)))) %>%
  arrange(param) %>%
  mutate(best_proof = ifelse(!is.na(best_bound) & best_proof == best_bound,
                        paste(best_proof, "*", sep = ""), best_proof),
         len = paste(len),
	 peak_memory = paste(floor(peak_memory / 1024))) %>%
  pivot_longer(c(best_proof, len, peak_memory)) %>%
  select(param, name, value) %>%
  pivot_wider(names_from = param, values_from = value) %>%
  mutate(name = ifelse(name == 'best_proof',
                       'Proof length, our approach',
		ifelse(name == 'len',
                       'Proof length, CaDiCaL + DRAT-trim', 'Peak memory usage, megabytes'))) %>%
  rename("Number of elements $n$" = name) %>%
  xtable(type = 'latex', label = "table:parametric-time-op",
         caption = "Proof lengths for MUSes of ordering principle formulas. Values with asterisks are additionally proven to be optimal.") %>%
  print(file = 'tex/op-mus.tex', include.rownames = FALSE)

## Introduction figure

# TODO Color should be number of clauses
# TODO Double-check the ylim
p <- ggplot(data %>% filter(solver == 'proofmin-opt' & best_proof == best_bound),
            aes(x = best_proof, y = len, color = n_axioms)) +
  geom_point(shape = 4) +
  geom_abline(color = 'red', linetype = 'dashed') +
  annotate('text', x=35, y=25, label = 'Optimal proof length',
           angle = 45, color = 'navy') +
  scale_color_gradient(low = "#FFDC00", high = "#FF0000",) +
  scale_x_continuous(expand = c(0, NA)) + scale_y_continuous(expand = c(0, NA)) +
  theme(legend.position = c(.25, .8), legend.background=element_blank()) +
  labs(color = 'Formula size',
       x = 'Optimal proof length', y = 'Proof length from CaDiCaL') +
  coord_equal(xlim = c(0, 60), ylim = c(0, 120))
ggsave('svg/cadical-vs-optimal.svg', plot = p, width = 3, height = 6)

data %>% filter(solver == 'proofmin-opt' & best_proof == best_bound) %>%
  group_by(instance_class, len >= 1.5 * best_proof) %>% count

data %>% filter(solver == 'proofmin-opt' & !is.na(best_bound) & len > 0 & best_proof == best_bound) %>%
  group_by(1) %>%
  reframe(enframe(quantile(len / best_proof, c(0, 0.25, 0.5, 0.75, 1)), "quantile", "rel_len")) %>%
  select(quantile, rel_len)

p <- ggplot(data %>% filter(solver == 'proofmin-opt' & instance_class == '3cnf' & !is.na(gap_to_mus_bound)),
            aes(x = factor(gap_to_mus_bound), y = time)) +
  geom_boxplot(show.legend = FALSE, color='darkgray', outliers = FALSE) +
  geom_jitter(size = 0.33, shape = 4, color = '#00000088') +
  scale_y_log10("Time to solve with\nour approach, s",
                breaks = trans_breaks("log10", function(x) 10^x),
                labels = trans_format("log10", math_format(10^.x))) +
  labs(
    x = 'Gap to MUS bound',
  )
ggsave('svg/bnb-time-gap.svg', plot = p, width = 4, height = 2)

p <- ggplot(data %>% filter(solver == 'proofmin-opt' & instance_class == '3cnf' & gap_to_mus_bound > 0),
            aes(x = factor(n_axioms), y = time, color = factor(gap_to_mus_bound))) +
  geom_boxplot(show.legend = FALSE, color='darkgray', outliers = FALSE) +
  geom_jitter(size = 0.75, shape = 4) +
  scale_y_log10("Time to solve with branch-and-bound, s",
                breaks = trans_breaks("log10", function(x) 10^x),
                labels = trans_format("log10", math_format(10^.x))) +
  labs(
    x = 'Formula size',
    y = 'Time to solve with branch-and-bound, s',
    color = 'Gap to MUS bound'
  ) + theme(legend.position='top')
ggsave('svg/bnb-time-clauses-nonzero.svg', plot = p, width = 5, height = 3)

p <- ggplot(data %>% filter(solver == 'proofmin-opt' & instance_class == '3cnf' & gap_to_mus_bound == 0),
            aes(x = factor(n_axioms), y = time)) +
  geom_boxplot(show.legend = FALSE, color='darkgray', outliers = FALSE) +
  geom_jitter(size = 0.75, shape = 4) +
  scale_y_log10("Time to solve with branch-and-bound, s",
                breaks = trans_breaks("log10", function(x) 10^x),
                labels = trans_format("log10", math_format(10^.x))) +
  labs(
    x = 'Formula size',
    y = 'Time to solve with branch-and-bound, s',
  ) +   theme(legend.position='top')
ggsave('svg/bnb-time-clauses-zero.svg', plot = p, width = 5, height = 3)

# TODO proofmin-len
p <- ggplot(data %>% filter(solver == 'proofmin-opt'),
            aes(x = best_proof, y = len, color = instance_class)) +
  geom_point(shape = 4) +
  scale_x_log10("Discovered proof length",
                breaks = trans_breaks("log10", function(x) 10^x),
                labels = trans_format("log10", math_format(10^.x))) +
  scale_y_log10("Reference proof length",
                breaks = trans_breaks("log10", function(x) 10^x),
                labels = trans_format("log10", math_format(10^.x))) +
  geom_abline(slope = 1, color = 'red', linetype = 'dashed') +
  labs(color = 'Instance class') +
  theme(legend.position = c(.2, .7),
        legend.background=element_blank())
ggsave('svg/length-scatter.svg', plot = p, width = 4, height = 4)

p <- ggplot(data %>% filter(solver == 'proofmin-len' & instance_class %in% c('subsetcard', 'coloring', '3cnf', 'smu')),
            aes(y = best_proof / len, x = instance_class, fill = instance_class)) +
  geom_boxplot() +
  labs(y = 'Proof length relative to CaDiCaL', x = 'Instance class') +
  scale_y_log10(labels = scales::percent_format()) +
  geom_hline(yintercept = 1, color = 'black') +
  theme(legend.position="none",
	axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1))
ggsave('svg/length-box.svg', plot = p, width = 2, height = 3.5)

rel_time <- (data %>% filter(solver == 'proofmin-opt' & best_proof == best_bound & mus)) |>
  inner_join(data %>% filter(solver == 'baseline' & time <= 60 * 60),
             by = c('formula'),
             suffix = c('.bnb', '.baseline')) %>%
  mutate(rel_time = time.bnb / time.baseline,
         formula_size = n_axioms.bnb) %>%
  select(formula, rel_time, formula_size)

p <- ggplot(rel_time %>% filter(!is.na(rel_time)),
            aes(y = rel_time, x = formula_size)) +
  geom_point(shape = 4) +
  geom_hline(yintercept=1, linetype="dashed", color = "red") +
  labs(x = 'Formula size') +
  scale_y_log10("Relative solving time",
                breaks = trans_breaks("log10", function(x) 10^x),
                labels = trans_format("log10", math_format(10^.x))) +
  theme(legend.position = c(.8, .9),
        legend.background=element_blank())
ggsave('svg/relative-time.svg', plot = p, width = 3, height = 4)


data %>% filter(solver == 'proofmin-opt' & best_proof == best_bound) |>
  full_join(data %>% filter(solver == 'baseline' & time <= 60 * 60),
             by = c('formula', 'mus', 'instance_class'),
             suffix = c('.bnb', '.baseline')) %>%
  mutate(opt_bnb = !is.na(time.bnb),
	 opt_baseline = !is.na(time.baseline)) %>%
  group_by(mus, instance_class, opt_bnb, opt_baseline) %>%
  count()
