library(tidyverse)
library(gridExtra)
library(grid)
library(imputeTestbench)
library(PSF)
library(imputePSF)
library(imputeTS)
library(scales)

# total obs in each simulated time series
n <- length(nottem)

# average of nottem for centering simulated data
avenot <- mean(nottem)

# create simulated data
# wvs is periodic only, nrms is random only, sim is combo of the two
simdat <- crossing(
  mts = seq(0, 1, length = 25),
  sds = seq(0, 8, length = 25)
) %>%
  group_by(mts, sds) %>%
  nest %>%
  mutate(
    wvs = map(mts, function(mts) ((as.numeric(nottem) - avenot) * mts) + avenot),
    nrm = map(sds, function(sds) rnorm(n = n, sd = sds))
  ) %>%
  select(-data) %>%
  mutate(
    sim = pmap(list(wvs, nrm), function(wvs, nrm) wvs + nrm)
  )

# run imputePSF with imputeTestbench, compare with replacement by means
simres <- simdat %>%
  mutate(
    errs = map(sim, impute_errors,
               methods = c("impute_PSF", "na.mean"), # "na.approx", "na.locf", "na.mean"),
               smps = 'mar', missPercentFrom = 30, missPercentTo = 30, blck = 100,
               errorParameter = 'mape')
  )

# optionally save and reload, simres takes a while
save(simres, file = 'simres.RData', compress = 'xz')
# load(file = 'simres.RData')

##
# plot output

# example plot for simulated observations, data prep
toplo <- simdat %>%
  filter(mts %in% c(0, 0.5, 1) & sds %in% c(0, 4, 8)) %>%
  mutate(ind = list(1:n)) %>%
  unnest

# plot
p1 <- ggplot(toplo, aes(x = ind, y = sim)) +
  geom_line(size = 0.25) +
  geom_point(size = 0.5, alpha = 0.8, colour = 'black') +
  facet_grid(mts ~ sds) +
  theme_bw(base_family = 'serif', base_size = 14) +
  theme(
    strip.background = element_blank()
  ) +
  scale_y_continuous('Simulated values') +
  scale_x_continuous('Index')

# plot of imputation differences, data prep
pls <- simres %>%
  mutate(
    toplo = map(errs, as.data.frame)
  ) %>%
  select(mts, sds, toplo) %>%
  unnest %>%
  select(-Parameter) %>%
  mutate(
    errdf = na.mean - impute_PSF,
    dir = ifelse(errdf > 0, '+', '')
  )

# ggplot base theme
pbase <- theme_bw(base_family = 'serif') +
  theme(
    panel.grid.major = element_blank(),
    panel.grid.minor = element_blank(),
    axis.text.x = element_text(size = 8),
    axis.text.y = element_text(size = 8),
    legend.position = 'top',
    legend.direction = 'horizontal',
    plot.margin = unit(c(1,1,0,0), "lines"),
    strip.background = element_blank(),
    strip.text.y = element_text(angle = 0, hjust = 0, vjust = 0.5),
    panel.background = element_rect(fill = 'black')
  )

# plot
p2 <- ggplot(pls) +
  geom_tile(aes(y = mts, x = sds, fill = errdf), colour = 'black') +
  geom_text(aes(y = mts, x = sds, label = dir)) +
  pbase +
  scale_y_reverse('Increasing signal', expand = c(0, 0)) +
  scale_x_continuous('Increasing noise', expand = c(0, 0)) +
  scale_fill_gradient2('Error difference\n(mean replacement - imputePSF)', low = muted("red"), mid = "white", high = muted("green"), midpoint = 0) +
  guides(fill = guide_colourbar(barheight = 0.5, barwidth = 8))

gparg <- gpar(fontfamily = 'serif', cex = 1.2)

# first plot save as jpeg
jpeg('simex.jpg', height = 5, width = 6.5, units = 'in', res = 500, quality = 100)

grid.arrange(
  arrangeGrob(
    p1,
    top = textGrob('Increasing noise', gp = gparg),
    right = textGrob('Increasing periodicity', rot = 270, gp = gparg))
)
dev.off()

# second plot save as jpeg
jpeg('simres.jpg', height = 4, 4.5, units = 'in', res = 500, quality = 100)
p2
dev.off()