message("Beginning simulation plots...")
library(dplyr)
library(ggplot2)
library(tidyr)

# Load in data
load("output_montes_small.RData")
res = res[!sapply(res, is.null)]
out = as.data.frame(do.call(rbind, res))

# Process data (just renaming things for easier plotting)
names(out) = tolower(gsub("[(|)]", "", names(out))) # cleanup intercept name
names(out) = gsub("uni1.ct_", "uni.ct1_", names(out)) # make univariate name consistent with others
names(out) = gsub("uni1.zi_", "uni.zi1_", names(out))
names(out) = gsub("uni2.ct_", "uni.ct2_", names(out))
names(out) = gsub("uni2.zi_", "uni.zi2_", names(out))
names(out) = gsub("(.*)(\\d)[.](\\d)", "\\1\\2\\3", names(out))

# reshape estimates from wide to long
ests = out %>%
  pivot_longer(
    cols = starts_with("est"),
    names_to = c("model", "param"),
    names_pattern = "est[.](.*)[.](.*)",
    values_to = "est"
  ) %>%
  dplyr::select(psi1_tru, psi2_tru, dep_true, model, param, est) %>%
  mutate(id = row_number())

# reshape ses from wide to long
ses = out %>%
  pivot_longer(
    cols = starts_with("se"),
    names_to = c("model", "param"),
    names_pattern = "se[.](.*)[.](.*)",
    values_to = "se"
  ) %>%
  dplyr::select(psi1_tru, psi2_tru, dep_true, model, param, se) %>%
  mutate(id = row_number())

clean = left_join(ests, ses, by = "id")
clean = clean %>% dplyr::select(-ends_with("x"))
names(clean) = sub(".y", "", names(clean))
clean = clean %>% 
  dplyr::select(psi1_tru, psi2_tru, dep_true, model, param, est, se) %>% 
  filter(!is.na(se))

temp = filter(clean, param == "ct1_x11", psi2_tru == .1)

n = 500
dep = c(.15, .85)

b1 = c(1, 3.25, -2.3)
b2 = c(2, -1.75, 3.5)

psi1 = psi2 = c(.1, .6)

temp$b_true = b1[2]

# subtract of true value of beta from all estimates
temp$est.c = temp$est - temp$b_true
temp = ungroup(filter(temp, between(est.c, -2.25, 2.25)))

clrs = RColorBrewer::brewer.pal(3, "Set1")[c(3,1,2)]

# density for beta
p=ggplot(temp, aes(x = est.c, color = model, fill = model)) + 
  geom_density(trim=TRUE, alpha=.15, bw="sj") +
  geom_vline(xintercept=0, color="black", linetype=2) + 
  labs(x = "Estimate", y = "Density", title = "Count Slope 1, Margin 1") + 
  facet_grid(dep_true ~ psi1_tru,
             labeller = label_bquote(
               rows = rho == .(dep_true),
               cols = psi == .(psi1_tru)
              )
             ) +
  coord_cartesian(xlim = c(-.5,1.5)) +
  theme_bw() +
  theme(
    panel.grid.major = element_blank(),
    panel.grid.minor = element_blank(),
    axis.text.x = element_text(size = 8),
    panel.spacing = unit(.25, "lines"),
    legend.position = "bottom",
    legend.background = element_rect(colour = "black"),
    plot.title = element_text(hjust = .5)
  ) +
  scale_colour_manual(
    labels =  c("Correct Model", "Biv Omits ZI", "Univariate ZI"),
    values = clrs,
    name = "Model"
  ) +
  scale_fill_manual(
    labels =  c("Correct Model", "Biv Omits ZI", "Univariate ZI"),
    values = clrs,
    name = "Model"
  )+
  labs(
    x = expression(hat(beta)~-~beta),
    y = "Kernel Density Estimate",
    title = "" #expression(paste("Density for ", hat(beta)[1]~-~beta, " in Count Margin 1, holding", psi[2] == 0.1))
  ) +
  guides(colour = guide_legend(
    title.position = "top",
    title.hjust = 0.5,
    override.aes = list(size = 2)
  ))

pdf("Figures/figure_2.pdf")
print(p)
dev.off()
print(p)


# Overconfidence plot
sums = clean %>% filter(param == "ct1_x11", psi2_tru == .1) %>% 
  group_by(model, psi1_tru, dep_true) %>%
  mutate(b_true =  b1[2]) %>% 
  summarise(overconf = mad(est)/median(se))


p=ggplot(sums, aes(x=model, y=overconf, fill=model)) + 
  geom_col(alpha = .75) + 
  facet_grid(dep_true ~ psi1_tru,
             labeller = label_bquote(
               rows = rho == .(dep_true),
               cols = psi == .(psi1_tru)
             )) +
theme_bw() +
  theme(
    panel.grid.major = element_blank(),
    panel.grid.minor = element_blank(),
    axis.text.x = element_text(
      size = 8
    ),
    panel.spacing = unit(.25, "lines"),
    legend.position = "none",
    legend.background = element_rect(colour = "black"),
    plot.title = element_text(hjust = .5)
  ) +
  scale_fill_manual(
    labels =  c("Correct Model", "Biv Omits ZI", "Univariate ZI"),
    values = clrs,
    name = "Model"
  ) +
  scale_x_discrete(
    labels = c("Correct Model", "Biv Omits ZI", "Univariate ZI")
    ) +
  labs(
    x = "Model",
    y = "Overconfidence",
    title = ""#expression(paste("Overconfidence for ", hat(beta)[1], " in Count Margin 1"))
  ) +
  guides(colour = guide_legend(
    title.position = "top",
    title.hjust = 0.5,
    override.aes = list(size = 3)
  ))

pdf("Figures/figure_3.pdf")
print(p)
dev.off()
print(p)




clrs = clrs[-3]
# Dependence plot
dat = filter(clean, param=="dependence", psi2_tru==.1)
dat$est.c = dat$est - dat$dep_true
p=ggplot(data=dat, aes(x=est.c, color = model, fill=model)) + 
  geom_density(alpha = .1, trim=TRUE, bw="sj") +
  geom_vline(xintercept=0, color="black", linetype=2) + 
  labs(x = "Estimate", y = "Density", title = "Dependence when Psi_2=.6") + 
  facet_grid(dep_true ~ psi1_tru,
             labeller = label_bquote(
               rows = rho == .(dep_true),
               cols = psi[1] == .(psi1_tru)
             )) +
  coord_cartesian(xlim = c(-.5,.5)) +
  theme_bw()  +
  theme(
    panel.grid.major = element_blank(),
    panel.grid.minor = element_blank(),
    axis.text.x = element_text(
      size = 8
    ),
    panel.spacing = unit(.25, "lines"),
    legend.position = "bottom",
    legend.background = element_rect(colour = "black"),
    plot.title = element_text(hjust = .5)
  ) +
  scale_colour_manual(
    labels =  c("Correct Model", "Biv Omits ZI"),
    values = clrs,
    name = "Model"
  ) +
  scale_fill_manual(
    labels =  c("Correct Model", "Biv Omits ZI"),
    values = clrs,
    name = "Model"
  )+
  labs(
    x = expression(hat(rho)~-~rho),
    y = "Kernel Density Estimate",
    title = ""#expression(paste("Density for ", hat(rho)~-~rho, " holding ", psi[2] == .1))
  ) +
  guides(colour = guide_legend(
    title.position = "top",
    title.hjust = 0.5,
    override.aes = list(size = 3)
  ))

pdf("Figures/figure_4.pdf")
print(p)
dev.off()
print(p)

message("Simulation plots complete.")

