## get redshift offsets from a data cube or stacked rss file

regrid <- function(lambda.out, lib) {
    lambda.in <- lib$lambda
    nc <- ncol(lib)
    lib.out <- matrix(NA, length(lambda.out), nc)
    lib.out[,1] <- lambda.out
    for (i in 2:nc) {
        lib.out[,i] <- approx(lambda.in, lib[,i], xout=lambda.out)$y
    }
    colnames(lib.out) <- colnames(lib)
    data.frame(lib.out)
}

getdz <- function(gdat, lib, snrthresh=5, nlthresh=2000, dzlim=0.003, searchinterval=1e-4) {
  dims <- dim(gdat$flux)
  nr <- dims[1]
  nc <- dims[2]
  dz <- matrix(NA, nr, nc)
  dz.err <- matrix(NA, nr, nc)
  z <- gdat$meta$z
  lambda <- gdat$lambda
  fmla <- as.formula(paste("f ~", paste(names(lib)[-1], collapse="+")))
  fn <- function(x) {
    lambda.out <- lambda/(1+z+x)
    lib.out <- regrid(lambda.out, lib)
    lfit <- lm(fmla, data=lib.out, weights=iv)
    deviance(lfit)
  }
  fn_v <- Vectorize(fn)
  z_search <- seq(-dzlim, dzlim, by=searchinterval)

  for (i in 1:nr) {
    for (j in 1:nc) {
      if (is.na(gdat$snr[i,j]) || gdat$snr[i,j]<=snrthresh) next
        f <- gdat$flux[i, j, ]
        if (length(which(!is.na(f))) < nlthresh) next
          iv <- gdat$ivar[i, j, ]
          dev_grid <- fn_v(z_search)
          z0 <- z_search[which.min(dev_grid)]
          bestz <- Rsolnp::solnp(pars=z0, fun=fn,
                                 LB=z0-searchinterval, UB=z0+searchinterval,
                                 control=list(trace=0))
          dz[i,j] <- bestz$pars
          dz.err[i,j] <- sqrt(2./bestz$hessian)
    }
  }
  list(dz=dz, dz.err=dz.err)
}

## peculiar velocity

vrel <- function(z, zbar) {
  299792.458*(z-zbar)/(1+zbar)
}

## simple 2D spatial interpolation

fillpoly <- function(ra, dec, zvals, dxy=0.5) {
    ny <- max(round((max(dec)-min(dec))*3600/dxy), 100)
    nx <- ny
    allok <- complete.cases(ra, dec, zvals)
    zsurf <- akima::interp(x=ra[allok], y=dec[allok], z=zvals[allok],
                           linear=FALSE, nx=nx, ny=ny)
    list(ra=zsurf$x, dec=zsurf$y, z=zsurf$z)
}

## an image function using ggplot2

ggimage <- function(zmat, x=NULL, y=NULL, col=viridis::viridis(256),
                    xlab=NULL, ylab=NULL, legend=NULL, title=NULL,
                    xrev=FALSE, yrev=FALSE, asp=1,
                    addContour=FALSE, binwidth=NULL
                   ) {
    require(ggplot2)
    if (is.null(x)) x <- 1:nrow(zmat)
    if (is.null(y)) y <- 1:ncol(zmat)
    if (!is.null(dim(x))) {
      x <- colMeans(x)
      y <- rowMeans(y)
    }
    xy <- expand.grid(x,y)
    df <- data.frame(cbind(xy, as.vector(zmat)))
    names(df) <- c("x","y","z")
    g1 <- ggplot(df, aes(x=x, y=y, z=z)) +
        geom_raster(aes(fill=z)) +
        scale_fill_gradientn(colors=col, na.value="#FFFFFF00") +
        coord_fixed(ratio = asp)
    if (addContour) {
        if (!is.null(binwidth)) {
            g1 <- g1 + geom_contour(binwidth=binwidth, na.rm=TRUE)
        } else {
            g1 <- g1 + geom_contour(na.rm=TRUE)
        }
    }
    if (xrev) g1 <- g1 + scale_x_reverse()
    if (yrev) g1 <- g1 + scale_y_reverse()
    if (!is.null(xlab)) {
        g1 <- g1 + xlab(xlab)
    }
    if (!is.null(ylab)) {
        g1 <- g1 + ylab(ylab)
    }
    if (!is.null(legend)) {
        g1 <- g1 + labs(fill=legend)
    }
    if (!is.null(title)) {
        g1 <- g1 + ggtitle(title)
    }
    g1
}

rainbow <- c("blue", "cyan", "green", "yellow", "red")

## disk galaxy rotation curve model

vrot <- function(gdat.stack, dz.stack, phi.guess, sd.phi.guess=10,
                      ci.guess=0.7, sd.ci.guess=0.1,
                      sd.kc0=0.5,
                      order=3, N_r=100,
                      r_eff=1, remax=1.5, v_norm=100, PLOTS=TRUE,
                      smodel=NULL,
                      stanfile="vrot.stan", stanfiledir=".",
                      seed=220755, iter=2000, warmup=1000,
                      chains=4, cores=4,
                      control=list(adapt_delta=0.975, max_treedepth=15)) {
  require(rstan)
  require(ggplot2)
  v <- vrel(dz.stack$dz+gdat.stack$meta$z, gdat.stack$meta$z)
  v_sample <- v[!is.na(v)]
  dv <- (dz.stack$dz.err/dz.stack$dz)*v
  df <- data.frame(x=gdat.stack$xpos, y=gdat.stack$ypos, v=v, dv=dv)
  df <- df[complete.cases(df),]
  rmax <- max(sqrt(gdat.stack$xpos^2+gdat.stack$ypos^2)[!is.na(dz.stack$dz)])
  if (rmax %% 2 != 0) rmax = rmax+1
  r_norm <- r_eff*remax
  r_post <- seq(0, 0.99, length=N_r)*rmax/r_norm

  vlos_data <- list(order=order, N=nrow(df), x=df$x/r_norm, y=df$y/r_norm,
                    v=df$v/v_norm, dv=df$dv/v_norm,
                    phi0=cos(phi.guess*pi/180), sd_phi0=sin(phi.guess*pi/180)*sd.phi.guess*pi/180,
                    si0 = sqrt(1-ci.guess^2), sd_si0=sd.ci.guess,
                    sd_kc0 = sd.kc0,
                    r_norm=r_norm, v_norm=v_norm,
                    N_r=N_r, r_post=r_post)

  inits <- function(chain_id) {
    p <- phi.guess*pi/180
    si <- sqrt(1-ci.guess^2)
    list(phi=p, si=si, x_c=0., y_c=0., v_sys=0.)
  }
  if (!is.null(smodel)) {
    stanfit <- sampling(smodel, data=vlos_data, init=inits, seed=seed,
                        iter=iter, warmup=warmup, control=control, chains=chains, cores=cores)
  } else {
    stanfit <- stan(file=file.path(stanfiledir, stanfile), data=vlos_data, init=inits,
                  seed=seed, iter=iter, warmup=warmup, control=control, chains=chains, cores=cores)
  }
  post <- extract(stanfit)
  vf_obs <- fillpoly(gdat.stack$ra.f, gdat.stack$dec.f, v)
  v_model <- rep(NA, length(v))
  v_model[!is.na(v)] <- colMeans(post$v_model)*v_norm
  vf_model <- fillpoly(gdat.stack$ra.f, gdat.stack$dec.f, v_model)
  r_q <- apply(post$r*remax, 2, quantile, probs=c(0.025, 0.5, 0.975))
  rownames(r_q) <- paste("r", rownames(r_q), sep="_")
  v_r_q <- apply(post$v_rot*v_norm, 2, quantile, probs=c(0.025, 0.5, 0.975))
  rownames(v_r_q) <- paste("v_rot", rownames(v_r_q), sep="_")
  v_e_q <- apply(post$v_exp*v_norm, 2, quantile, probs=c(0.025, 0.5, 0.975))
  rownames(v_e_q) <- paste("v_exp", rownames(v_e_q), sep="_")
  df_v <- data.frame(cbind(t(r_q), t(v_r_q), t(v_e_q)))
  graphs <- list()
  i <- 1
  df <- data.frame(x=gdat.stack$xpos, y=gdat.stack$ypos, v=v)
  graphs[[i]] <- ggplot(df) + geom_point(aes(x=x, y=y, color=v), size=7) +
                              scale_colour_gradientn(colors=rainbow, na.value="gray95") +
                              coord_fixed(ratio=1) +
                              ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": observed fiber velocities"))
  i <- i+1
  graphs[[i]] <- ggimage(vf_obs$z, vf_obs$ra, vf_obs$dec, col=rainbow, asp=1/cos(pi*gdat.stack$meta$dec/180), addContour=T) +
                          xlim(rev(range(gdat.stack$ra.f))) + ylim(range(gdat.stack$dec.f)) +
                          xlab(expression(alpha)) + ylab(expression(delta)) +
                          ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": observed velocity field"))
  i <- i+1
  graphs[[i]] <- ggimage(vf_model$z, vf_model$ra, vf_model$dec, col=rainbow, asp=1/cos(pi*gdat.stack$meta$dec/180), addContour=T) +
                          xlim(rev(range(gdat.stack$ra.f))) + ylim(range(gdat.stack$dec.f)) +
                          xlab(expression(alpha)) + ylab(expression(delta)) +
                          ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": model velocity field"))
  if (exists("v_res", post)) {
    v_res <- rep(NA, length(v))
    v_res[!is.na(v)] <- colMeans(post$v_res)*v_norm
    vf_res <- fillpoly(gdat.stack$ra.f, gdat.stack$dec.f, v_res)
  } else {
    vf_res <- list(z=vf_obs$z-vf_model$z, ra=vf_obs$ra, dec=vf_obs$dec)
  }
  i <- i+1
  graphs[[i]] <- ggimage(vf_res$z, vf_res$ra, vf_res$dec, asp=1/cos(pi*gdat.stack$meta$dec/180), addContour=T) +
                        xlim(rev(range(gdat.stack$ra.f))) + ylim(range(gdat.stack$dec.f)) +
                        xlab(expression(alpha)) + ylab(expression(delta)) +
                        ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": residual velocity field"))
  df1 <- data.frame(r=as.vector(abs(post$r))*remax, v_rot=as.vector(post$v_rot)*v_norm, v_exp=as.vector(post$v_exp)*v_norm)
  i <- i+1
  graphs[[i]] <- ggplot(df_v) + geom_point(aes(x=r_50., y=v_rot_50.)) +
          geom_errorbar(aes(x=r_50., ymin=v_rot_2.5., ymax=v_rot_97.5.)) +
          geom_errorbarh(aes(x=r_50., y=v_rot_50., xmin=r_2.5., xmax=r_97.5.)) +
          stat_density_2d(aes(x=r, y=v_rot, fill=..level..), data=df1, geom="polygon", alpha=0.5, show.legend=FALSE) +
          ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": joint distribution of r, v_rot")) +
          xlab(expression(r/r[eff]))
  i <- i+1
  graphs[[i]] <- ggplot(df_v) + geom_point(aes(x=r_50., y=v_exp_50.)) +
          geom_errorbar(aes(x=r_50., ymin=v_exp_2.5., ymax=v_exp_97.5.)) +
          geom_errorbarh(aes(x=r_50., y=v_exp_50., xmin=r_2.5., xmax=r_97.5.)) +
          stat_density_2d(aes(x=r, y=v_exp, fill=..level..), data=df1, geom="polygon", alpha=0.5, show.legend=FALSE) +
          ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": joint distribution of r, v_exp")) +
          xlab(expression(r/r[eff]))
  if(exists("vrot_post", post)) {
  	  vq <- t(apply(v_norm * post$vrot_post, 2, quantile, probs=c(0.025, 0.5, 0.975)))
  	  df2 <- data.frame(r=r_post*remax, ymin=vq[,"2.5%"], ymed=vq[,"50%"], ymax=vq[,"97.5%"])
          i <- i+1
  	  graphs[[i]] <- ggplot(df2) + geom_line(aes(x=r, y=ymed)) + geom_ribbon(aes(x=r, ymin=ymin, ymax=ymax), alpha=0.5) +
          xlab(expression(r/r[eff]))
  } else {
  	  df2 <- NULL
  }
  if (PLOTS) {
  	  for (i in seq_along(graphs)) {
  	  	  x11()
  	  	  plot(graphs[[i]])
  	  }
  }
  list(stanfit=stanfit, r_eff=r_eff, remax=remax, v_norm=v_norm, vlos_data=vlos_data,
       v_sample=v_sample, vf_obs=vf_obs, vf_model=vf_model, vf_res=vf_res, df_v=df_v, df1=df1, df2=df2,
       graphs=graphs
  )
}

## disk galaxy rotation curve model -- gaussian process version

vrot_gp <- function(gdat.stack, dz.stack, phi.guess, sd.phi.guess=10,
                      ci.guess=0.7, sd.ci.guess=0.1,
                      sd.kc0=0.5,
                      order=3, N_xy=25^2,
                      r_eff=1, remax=1.5, v_norm=100, PLOTS=TRUE,
                      smodel=NULL,
                      stanfile="vrot_gp.stan", stanfiledir="~/vrot_stanmodels",
                      seed=220755, iter=1000, warmup=500,
                      chains=4, cores=4,
                      control=list(adapt_delta=0.9, max_treedepth=10)) {
  require(rstan)
  require(spmutils)
  require(ggplot2)
  require(akima)
  v <- vrel(dz.stack$dz+gdat.stack$meta$z, gdat.stack$meta$z)
  v_sample <- v[!is.na(v)]
  dv <- (dz.stack$dz.err/dz.stack$dz)*v
  df <- data.frame(x=gdat.stack$xpos, y=gdat.stack$ypos, v=v, dv=dv)
  allok <- complete.cases(df)
  df <- df[allok,]

  n_r <- round(sqrt(N_xy))
  N_xy <- n_r^2
  rho <- seq(1/n_r, 1, length=n_r)
  theta <- seq(0, (n_r-1)*2*pi/n_r, length=n_r)
  rt <- expand.grid(theta, rho)
  x_pred <- rt[,2]*cos(rt[,1])
  y_pred <- rt[,2]*sin(rt[,1])
  r_norm <- r_eff*remax

  vlos_data <- list(order=order, N=nrow(df), x=df$x/r_norm, y=df$y/r_norm,
                    v=df$v/v_norm, dv=df$dv/v_norm,
                    phi0=phi.guess*pi/180, sd_phi0=sd.phi.guess*pi/180,
                    ci0 = ci.guess, sd_ci0=sd.ci.guess,
                    sd_kc0 = sd.kc0,
                    r_norm=r_norm, v_norm=v_norm,
                    N_xy=N_xy, N_r=n_r, x_pred=x_pred, y_pred=y_pred
                    )

  inits <- function(chain_id) {
    p <- phi.guess*pi/180
    ci <- ci.guess
    list(phi=p, cos_i=ci, x_c=0., y_c=0., v_sys=0.)
  }
  if (!is.null(smodel)) {
    stanfit <- sampling(smodel, data=vlos_data, init=inits, seed=seed,
                        iter=iter, warmup=warmup, control=control, chains=chains, cores=cores)
  } else {
    stanfit <- stan(file=file.path(stanfiledir, stanfile), data=vlos_data, init=inits,
                  seed=seed, iter=iter, warmup=warmup, control=control, chains=chains, cores=cores)
  }
  post <- extract(stanfit)
  graphs <- list()
  vf_model <- rep(NA, length(gdat.stack$xpos))
  vf_res <- rep(NA, length(gdat.stack$ypos))
  vf_model[allok] <- colMeans(post$vf_model)*v_norm
  vf_res[allok] <- colMeans(post$vf_res)*v_norm
  df <- data.frame(x=gdat.stack$xpos, y=gdat.stack$ypos, v=v, vf_model=vf_model, vf_res=vf_res)
  g <- 1
  graphs[[g]] <- ggplot(df) + geom_point(aes(x=x, y=y, color=v), size=7) +
                              scale_colour_gradientn(colors=rainbow, na.value="gray95") +
                              coord_fixed(ratio=1) +
                              ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": observed fiber velocities"))
  g <- g+1
  df <- df[allok,]
  df1 <- akima::interp(x=df$x, y=df$y, z=df$vf_model, nx=200, ny=200)
  graphs[[g]] <- ggimage(df1$z, df1$x, df1$y, col=rainbow, xlab="x", ylab="y", legend="vf_model", addContour=TRUE)
  g <- g+1
  df1 <- akima::interp(x=df$x, y=df$y, z=df$vf_res, nx=200, ny=200)
  graphs[[g]] <- ggimage(df1$z, df1$x, df1$y, xlab="x", ylab="y", legend="vf_res", addContour=TRUE)

  v_pred <- colMeans(post$v_pred)
  sd_v_pred <- apply(post$v_pred, 2, sd)
  v_model <- colMeans(post$v_model)
  v_res <- colMeans(post$v_res)
  x <- x_pred*r_norm
  y <- y_pred*r_norm
  rho <- c(0, rho)*remax
  df <- akima::interp(x=x, y=y, z=v_pred, nx=200, ny=200)
  g <- g+1
  graphs[[g]] <- ggimage(df$z, df$x, df$y, col=rainbow, xlab="x", ylab="y", legend="v_pred", addContour=TRUE)
  df <- akima::interp(x=x, y=y, z=v_model, nx=200, ny=200)
  g <- g+1
  graphs[[g]] <- ggimage(df$z, df$x, df$y, col=rainbow, xlab="x", ylab="y", legend="v_model", addContour=TRUE)
  df <- akima::interp(x=x, y=y, z=v_res, nx=200, ny=200)
  g <- g+1
  graphs[[g]] <- ggimage(df$z, df$x, df$y, xlab="x", ylab="y", legend="v_res", addContour=TRUE)
  df <- akima::interp(x=x, y=y, z=sd_v_pred, nx=200, ny=200)
  g <- g+1
  graphs[[g]] <- ggimage(df$z, df$x, df$y, xlab="x", ylab="y", legend="sd_v_pred", addContour=TRUE)
  vrot_pred <- post$vrot_pred
  vexp_pred <- post$vexp_pred
  v_r_q <- data.frame(cbind(rho, t(apply(vrot_pred, 2, quantile, probs=c(.025, 0.5, 0.975)))))
  names(v_r_q) <- c("rho", "vr_025", "vr_500", "vr_975")
  g <- g+1
  graphs[[g]] <- ggplot(v_r_q) + geom_line(aes(x=rho, y= vr_500)) + geom_ribbon(aes(x=rho, ymin=vr_025, ymax=vr_975), alpha=0.5) +
                  xlab(expression(r/r[eff]))+ylab(expression(V[rot]))
  v_e_q <- data.frame(cbind(rho, t(apply(vexp_pred, 2, quantile, probs=c(.025, 0.5, 0.975)))))
  names(v_e_q) <- c("rho", "ve_025", "ve_500", "ve_975")
  g <- g+1
  graphs[[g]] <- ggplot(v_e_q) + geom_line(aes(x=rho, y= ve_500)) + geom_ribbon(aes(x=rho, ymin=ve_025, ymax=ve_975), alpha=0.5) +
                  xlab(expression(r/r[eff]))+ylab(expression(V[exp]))

  r_q <- apply(post$r*remax, 2, quantile, probs=c(0.025, 0.5, 0.975))
  rownames(r_q) <- paste("r", rownames(r_q), sep="_")
  v_r_q <- apply(post$v_rot*v_norm, 2, quantile, probs=c(0.025, 0.5, 0.975))
  rownames(v_r_q) <- paste("v_rot", rownames(v_r_q), sep="_")
  df_v <- data.frame(cbind(t(r_q), t(v_r_q)))
  df <- data.frame(r=as.vector(abs(post$r))*remax, v_rot=as.vector(post$v_rot)*v_norm)

  g <- g+1
  graphs[[g]] <- ggplot(df_v) + geom_point(aes(x=r_50., y=v_rot_50.)) +
          geom_errorbar(aes(x=r_50., ymin=v_rot_2.5., ymax=v_rot_97.5.)) +
          geom_errorbarh(aes(x=r_50., y=v_rot_50., xmin=r_2.5., xmax=r_97.5.)) +
          ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": joint distribution of r, v_rot")) +
          xlab(expression(r/r[eff])) + ylab("v (km/sec)")
  if (PLOTS) {
    for (i in seq_along(graphs)) {
      x11()
      plot(graphs[[i]])
    }
  }

  list(stanfit=stanfit, r_eff=r_eff, remax=remax, v_norm=v_norm, vlos_data=vlos_data,
       v_sample=v_sample, x_pred=x_pred, y_pred=y_pred,
       graphs=graphs
  )
}


library(tidyr)
library(dplyr)
library(rstan)
library(data.table)
library(lubridate)
library(gdata)
library(EnvStats)
library(matrixStats)
library(scales)
library(gridExtra)
library(ggpubr)
library(bayesplot)
library(cowplot)
library(svglite)
library(ggplot2)

source("geom-stepribbon.r")
#---------------------------------------------------------------------------
make_three_pannel_plot <- function(){

  args <- commandArgs(trailingOnly = TRUE)

  filename2 = args[1]
  percent_pop = FALSE

  load(paste0("combinedTDI/results/", filename2))
  print(sprintf("loading: %s",paste0("results/",filename2)))

  filename2 <- strsplit(filename2, "-stanfit.Rdata")[[1]][1]

  # covariates = read.csv('data/interventions.csv', stringsAsFactors = FALSE)
  #
  # # ADD SLOVENIA:
  # covariates_slo <- covariates %>% filter(Country=='Sweden')
  # covariates_slo$Country <- 'Slovenia'
  # covariates_slo$Event <- ''
  # covariates_slo$Date.effective <- c("16.03.2020", "10.03.2020",
  #                                    "20.03.2020", "16.03.2020",
  #                                    "09.03.2020")
  # covariates <- rbind(covariates, covariates_slo)
  #
  # names_covariates = c('Schools + Universities','Self-isolating if ill', 'Public events',
  #                      'Lockdown', 'Social distancing encouraged')
  # covariates <- covariates %>%
  #   filter((Type %in% names_covariates))
  # covariates <- covariates[,c(1,2,4)]
  # covariates <- spread(covariates, Type, Date.effective)
  # names(covariates) <- c('Country','lockdown', 'public_events', 'schools_universities','self_isolating_if_ill', 'social_distancing_encouraged')
  # covariates <- covariates[c('Country','schools_universities', 'self_isolating_if_ill', 'public_events', 'lockdown', 'social_distancing_encouraged')]
  # covariates$schools_universities <- as.Date(covariates$schools_universities, format = "%d.%m.%Y")
  # covariates$lockdown <- as.Date(covariates$lockdown, format = "%d.%m.%Y")
  # covariates$public_events <- as.Date(covariates$public_events, format = "%d.%m.%Y")
  # covariates$self_isolating_if_ill <- as.Date(covariates$self_isolating_if_ill, format = "%d.%m.%Y")
  # covariates$social_distancing_encouraged <- as.Date(covariates$social_distancing_encouraged, format = "%d.%m.%Y")

  all_data <- data.frame()
  intervention_data <- data.frame()
  for(i in 1:length(countries)){
    print(i)
    N <- length(dates[[i]])
    country <- countries[[i]]

    predicted_cases <- colMeans(prediction[,1:N,i])
    predicted_cases_li <- colQuantiles(prediction[,1:N,i], probs=.025)
    predicted_cases_ui <- colQuantiles(prediction[,1:N,i], probs=.975)
    predicted_cases_li2 <- colQuantiles(prediction[,1:N,i], probs=.25)
    predicted_cases_ui2 <- colQuantiles(prediction[,1:N,i], probs=.75)

    estimated_deaths <- colMeans(estimated.deaths[,1:N,i])
    estimated_deaths_li <- colQuantiles(estimated.deaths[,1:N,i], probs=.025)
    estimated_deaths_ui <- colQuantiles(estimated.deaths[,1:N,i], probs=.975)
    estimated_deaths_li2 <- colQuantiles(estimated.deaths[,1:N,i], probs=.25)
    estimated_deaths_ui2 <- colQuantiles(estimated.deaths[,1:N,i], probs=.75)

    estimated_deathsh <- colMeans(estimated.deathsh[,1:N,i])
    estimated_deathsh_li <- colQuantiles(estimated.deathsh[,1:N,i], probs=.025)
    estimated_deathsh_ui <- colQuantiles(estimated.deathsh[,1:N,i], probs=.975)
    estimated_deathsh_li2 <- colQuantiles(estimated.deathsh[,1:N,i], probs=.25)
    estimated_deathsh_ui2 <- colQuantiles(estimated.deathsh[,1:N,i], probs=.75)

    estimated_deathsc <- colMeans(estimated.deathsc[,1:N,i])
    estimated_deathsc_li <- colQuantiles(estimated.deathsc[,1:N,i], probs=.025)
    estimated_deathsc_ui <- colQuantiles(estimated.deathsc[,1:N,i], probs=.975)
    estimated_deathsc_li2 <- colQuantiles(estimated.deathsc[,1:N,i], probs=.25)
    estimated_deathsc_ui2 <- colQuantiles(estimated.deathsc[,1:N,i], probs=.75)

    estimated_cases <- colMeans(estimated.cases[,1:N,i])
    estimated_cases_li <- colQuantiles(estimated.cases[,1:N,i], probs=.025)
    estimated_cases_ui <- colQuantiles(estimated.cases[,1:N,i], probs=.975)
    estimated_cases_li2 <- colQuantiles(estimated.cases[,1:N,i], probs=.25)
    estimated_cases_ui2 <- colQuantiles(estimated.cases[,1:N,i], probs=.75)

    estimated_hosp <- colMeans(estimated.hospitals[,1:N,i])
    estimated_hosp_li <- colQuantiles(estimated.hospitals[,1:N,i], probs=.025)
    estimated_hosp_ui <- colQuantiles(estimated.hospitals[,1:N,i], probs=.975)
    estimated_hosp_li2 <- colQuantiles(estimated.hospitals[,1:N,i], probs=.25)
    estimated_hosp_ui2 <- colQuantiles(estimated.hospitals[,1:N,i], probs=.75)

    estimated_hospi <- colMeans(estimated.hospitalsi[,1:N,i])
    estimated_hospi_li <- colQuantiles(estimated.hospitalsi[,1:N,i], probs=.025)
    estimated_hospi_ui <- colQuantiles(estimated.hospitalsi[,1:N,i], probs=.975)
    estimated_hospi_li2 <- colQuantiles(estimated.hospitalsi[,1:N,i], probs=.25)
    estimated_hospi_ui2 <- colQuantiles(estimated.hospitalsi[,1:N,i], probs=.75)

    estimated_hospo <- colMeans(estimated.hospitalso[,1:N,i])
    estimated_hospo_li <- colQuantiles(estimated.hospitalso[,1:N,i], probs=.025)
    estimated_hospo_ui <- colQuantiles(estimated.hospitalso[,1:N,i], probs=.975)
    estimated_hospo_li2 <- colQuantiles(estimated.hospitalso[,1:N,i], probs=.25)
    estimated_hospo_ui2 <- colQuantiles(estimated.hospitalso[,1:N,i], probs=.75)

    estimated_icu <- colMeans(estimated.icus[,1:N,i])
    estimated_icu_li <- colQuantiles(estimated.icus[,1:N,i], probs=.025)
    estimated_icu_ui <- colQuantiles(estimated.icus[,1:N,i], probs=.975)
    estimated_icu_li2 <- colQuantiles(estimated.icus[,1:N,i], probs=.25)
    estimated_icu_ui2 <- colQuantiles(estimated.icus[,1:N,i], probs=.75)

    estimated_icui <- colMeans(estimated.icusi[,1:N,i])
    estimated_icui_li <- colQuantiles(estimated.icusi[,1:N,i], probs=.025)
    estimated_icui_ui <- colQuantiles(estimated.icusi[,1:N,i], probs=.975)
    estimated_icui_li2 <- colQuantiles(estimated.icusi[,1:N,i], probs=.25)
    estimated_icui_ui2 <- colQuantiles(estimated.icusi[,1:N,i], probs=.75)

    estimated_icuo <- colMeans(estimated.icuso[,1:N,i])
    estimated_icuo_li <- colQuantiles(estimated.icuso[,1:N,i], probs=.025)
    estimated_icuo_ui <- colQuantiles(estimated.icuso[,1:N,i], probs=.975)
    estimated_icuo_li2 <- colQuantiles(estimated.icuso[,1:N,i], probs=.25)
    estimated_icuo_ui2 <- colQuantiles(estimated.icuso[,1:N,i], probs=.75)

    rt <- colMeans(out$Rt_adj[,1:N,i])
    rt_li <- colQuantiles(out$Rt_adj[,1:N,i],probs=.025)
    rt_ui <- colQuantiles(out$Rt_adj[,1:N,i],probs=.975)
    rt_li2 <- colQuantiles(out$Rt_adj[,1:N,i],probs=.25)
    rt_ui2 <- colQuantiles(out$Rt_adj[,1:N,i],probs=.75)

    covariates_country <- as.data.frame(matrix(NA, nrow=1, ncol=length(interventions)))
    colnames(covariates_country) <- paste0("interventions",1:length(interventions))
    for(iter in 1:ncol(covariates_country)){
        class(covariates_country[,iter]) <- "Date"
        covariates_country[1,iter] <- interventions[iter]
    }
    covariates_country_long <- gather(covariates_country, key = "key",
                                      value = "value")
    covariates_country_long$x <- rep(NULL, length(covariates_country_long$key))
    un_dates <- unique(covariates_country_long$value)

    for (k in 1:length(un_dates)){
      idxs <- which(covariates_country_long$value == un_dates[k])
      max_val <- round(max(rt_ui)) + 0.3
      for (j in idxs){
        covariates_country_long$x[j] <- max_val
        max_val <- max_val - 0.3
      }
    }

    covariates_country_long$value <- as_date(covariates_country_long$value)
    covariates_country_long$country <- rep(country,
                                           length(covariates_country_long$value))

    data_country <- data.frame("time" = as_date(as.character(dates[[i]])),
                               "country" = rep(country, length(dates[[i]])),
                               "reported_cases" = reported_cases[[i]],
                               "reported_cases_c" = cumsum(reported_cases[[i]]),
                               "predicted_cases_c" = cumsum(predicted_cases),
                               "predicted_min_c" = cumsum(predicted_cases_li),
                               "predicted_max_c" = cumsum(predicted_cases_ui),
                               "predicted_cases" = predicted_cases,
                               "predicted_min" = predicted_cases_li,
                               "predicted_max" = predicted_cases_ui,
                               "predicted_min2" = predicted_cases_li2,
                               "predicted_max2" = predicted_cases_ui2,
                               "deaths" = deaths_by_country[[i]],
                               "deaths_c" = cumsum(deaths_by_country[[i]]),
                               "estimated_deaths_c" =  cumsum(estimated_deaths),
                               "death_min_c" = cumsum(estimated_deaths_li),
                               "death_max_c"= cumsum(estimated_deaths_ui),
                               "estimated_deaths" = estimated_deaths,
                               "death_min" = estimated_deaths_li,
                               "death_max"= estimated_deaths_ui,
                               "death_min2" = estimated_deaths_li2,
                               "death_max2"= estimated_deaths_ui2,
                               "rt" = rt,
                               "rt_min" = rt_li,
                               "rt_max" = rt_ui,
                               "rt_min2" = rt_li2,
                               "rt_max2" = rt_ui2
    )

    data_country2 <- data.frame("estimated_cases_c" =  cumsum(estimated_cases),
                                "cases_min_c" = cumsum(estimated_cases_li),
                                "cases_max_c"= cumsum(estimated_cases_ui),
                                "estimated_cases" = estimated_cases,
                                "cases_min" = estimated_cases_li,
                                "cases_max"= estimated_cases_ui,
                                "cases_min2" = estimated_cases_li2,
                                "cases_max2"= estimated_cases_ui2)

    data_deathsh <- data.frame("time" = as_date(as.character(dates[[i]])),
                               "deaths" = deathsh_by_country[[i]],
                               "deaths_c" = cumsum(deathsh_by_country[[i]]),
                               "estimated_deaths_c" =  cumsum(estimated_deathsh),
                               "death_min_c" = cumsum(estimated_deathsh_li),
                               "death_max_c"= cumsum(estimated_deathsh_ui),
                               "estimated_deaths" = estimated_deathsh,
                               "death_min" = estimated_deathsh_li,
                               "death_max"= estimated_deathsh_ui,
                               "death_min2" = estimated_deathsh_li2,
                               "death_max2"= estimated_deathsh_ui2)

    data_deathsc <- data.frame("time" = as_date(as.character(dates[[i]])),
                               "deaths" = deathsc_by_country[[i]],
                               "deaths_c" = cumsum(deathsc_by_country[[i]]),
                               "estimated_deaths_c" =  cumsum(estimated_deathsc),
                               "death_min_c" = cumsum(estimated_deathsc_li),
                               "death_max_c"= cumsum(estimated_deathsc_ui),
                               "estimated_deaths" = estimated_deathsc,
                               "death_min" = estimated_deathsc_li,
                               "death_max"= estimated_deathsc_ui,
                               "death_min2" = estimated_deathsc_li2,
                               "death_max2"= estimated_deathsc_ui2)

    data_hosp <- data.frame("time" = as_date(as.character(dates[[i]])),
                            "estimated_hosp_c" =  cumsum(estimated_hosp),
                            "hosp_min_c" = cumsum(estimated_hosp_li),
                            "hosp_max_c"= cumsum(estimated_hosp_ui),
                            "estimated_hosp" = estimated_hosp,
                            "hosp_min" = estimated_hosp_li,
                            "hosp_max"= estimated_hosp_ui,
                            "hosp_min2" = estimated_hosp_li2,
                            "hosp_max2"= estimated_hosp_ui2,
                            "reported_hosp" = hosps_by_country[[i]])

    data_hospi <- data.frame("time" = as_date(as.character(dates[[i]])),
                            "estimated_hosp_c" =  cumsum(estimated_hospi),
                            "hosp_min_c" = cumsum(estimated_hospi_li),
                            "hosp_max_c"= cumsum(estimated_hospi_ui),
                            "estimated_hosp" = estimated_hospi,
                            "hosp_min" = estimated_hospi_li,
                            "hosp_max"= estimated_hospi_ui,
                            "hosp_min2" = estimated_hospi_li2,
                            "hosp_max2"= estimated_hospi_ui2,
                            "reported_hosp" = hospsin_by_country[[i]])

    data_hospo <- data.frame("time" = as_date(as.character(dates[[i]])),
                             "estimated_hosp_c" =  cumsum(estimated_hospo),
                             "hosp_min_c" = cumsum(estimated_hospo_li),
                             "hosp_max_c"= cumsum(estimated_hospo_ui),
                             "estimated_hosp" = estimated_hospo,
                             "hosp_min" = estimated_hospo_li,
                             "hosp_max"= estimated_hospo_ui,
                             "hosp_min2" = estimated_hospo_li2,
                             "hosp_max2"= estimated_hospo_ui2,
                             "reported_hosp" = hospsout_by_country[[i]])

    data_icu <- data.frame("time" = as_date(as.character(dates[[i]])),
                            "estimated_hosp_c" =  cumsum(estimated_icu),
                            "icu_min_c" = cumsum(estimated_icu_li),
                            "icu_max_c"= cumsum(estimated_icu_ui),
                            "estimated_icu" = estimated_icu,
                            "icu_min" = estimated_icu_li,
                            "icu_max"= estimated_icu_ui,
                            "icu_min2" = estimated_icu_li2,
                            "icu_max2"= estimated_icu_ui2,
                            "reported_icu" = icus_by_country[[i]])

    data_icui <- data.frame("time" = as_date(as.character(dates[[i]])),
                           "estimated_hosp_c" =  cumsum(estimated_icui),
                           "icu_min_c" = cumsum(estimated_icui_li),
                           "icu_max_c"= cumsum(estimated_icui_ui),
                           "estimated_icu" = estimated_icui,
                           "icu_min" = estimated_icui_li,
                           "icu_max"= estimated_icui_ui,
                           "icu_min2" = estimated_icui_li2,
                           "icu_max2"= estimated_icui_ui2,
                           "reported_icu" = icusin_by_country[[i]])

    data_icuo <- data.frame("time" = as_date(as.character(dates[[i]])),
                            "estimated_hosp_c" =  cumsum(estimated_icuo),
                            "icu_min_c" = cumsum(estimated_icuo_li),
                            "icu_max_c"= cumsum(estimated_icuo_ui),
                            "estimated_icu" = estimated_icuo,
                            "icu_min" = estimated_icuo_li,
                            "icu_max"= estimated_icuo_ui,
                            "icu_min2" = estimated_icuo_li2,
                            "icu_max2"= estimated_icuo_ui2,
                            "reported_icu" = icusout_by_country[[i]])

    all_data <- rbind(all_data, data_country)
    intervention_data <- rbind(intervention_data, covariates_country_long)

    make_plots(data_country = data_country,
               data_country2 = data_country2,
               data_deathsh = data_deathsh,
               data_deathsc = data_deathsc,
               data_hosp = data_hosp,
               data_hospi = data_hospi,
               data_hospo = data_hospo,
               data_icu = data_icu,
               data_icui = data_icui,
               data_icuo = data_icuo,
               covariates_country_long = covariates_country_long,
               filename2 = filename2,
               country = country,
               percent_pop = percent_pop,
               fit.spline = fit.spline)

  }
  # write.csv(all_data, paste0("results/", "base-plot.csv"))
  # write.csv(intervention_data, paste0("results/", "base-intervention.csv"))
}

#---------------------------------------------------------------------------
make_plots <- function(data_country, data_country2,
                       data_deathsh, data_deathsc,
                       data_hosp, data_hospi, data_hospo,
                       data_icu, data_icui, data_icuo,
                       covariates_country_long,
                       filename2, country, percent_pop, fit.spline){

  language_english <- TRUE

  if (country == 'United_Kingdom')
    country = 'United Kingdom'
  data_cases_95 <- data.frame(data_country$time, data_country$predicted_min,
                              data_country$predicted_max)
  names(data_cases_95) <- c("time", "cases_min", "cases_max")
  data_cases_95$key <- rep("nintyfive", length(data_cases_95$time))
  data_cases_50 <- data.frame(data_country$time, data_country$predicted_min2,
                              data_country$predicted_max2)
  names(data_cases_50) <- c("time", "cases_min", "cases_max")
  data_cases_50$key <- rep("fifty", length(data_cases_50$time))
  data_cases <- rbind(data_cases_95, data_cases_50)
  levels(data_cases$key) <- c("ninetyfive", "fifty")

  if(language_english) ylab_text <- "Daily number of infections"
  else ylab_text <- "Dnevno stevilo okuzb"

  p1 <- ggplot(data_country) +
    geom_bar(data = data_country, aes(x = time, y = reported_cases),
             fill = "coral4", stat='identity', alpha=0.5) +
    geom_ribbon(data = data_cases,
                aes(x = time, ymin = cases_min, ymax = cases_max, fill = key)) +
    # xlab("") +
    # ylab("Daily number of infections\n") +
    labs(x="",
         y=ylab_text)+
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b")) +
    scale_y_continuous(expand = c(0, 0), labels = comma) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("deepskyblue4", 0.55),
                                 alpha("deepskyblue4", 0.45))) +
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "None") + ggtitle(country) +
    guides(fill=guide_legend(ncol=1))

  ##
  # Deaths:
  data_deaths_95 <- data.frame(data_country$time, data_country$death_min,
                               data_country$death_max)
  names(data_deaths_95) <- c("time", "death_min", "death_max")
  data_deaths_95$key <- rep("nintyfive", length(data_deaths_95$time))
  data_deaths_50 <- data.frame(data_country$time, data_country$death_min2,
                               data_country$death_max2)
  names(data_deaths_50) <- c("time", "death_min", "death_max")
  data_deaths_50$key <- rep("fifty", length(data_deaths_50$time))
  data_deaths <- rbind(data_deaths_95, data_deaths_50)
  levels(data_deaths$key) <- c("ninetyfive", "fifty")

  if(language_english) ylab_text <- "Daily number of deaths"
  else ylab_text <- "Dnevno stevilo umrlih\n"

  if(language_english) xlab_text <- "Date"
  else xlab_text <- "Datum"

  p2 <-   ggplot(data_country, aes(x = time)) +
    geom_bar(data = data_country, aes(y = deaths, fill = "reported"),
             fill = "coral4", stat='identity', alpha=0.5) +
    geom_ribbon(
      data = data_deaths,
      aes(ymin = death_min, ymax = death_max, fill = key)) +
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b")) +
    scale_y_continuous(expand = c(0, 0), labels = comma) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("deepskyblue4", 0.55),
                                 alpha("deepskyblue4", 0.45))) +
    xlab(xlab_text)+
    ylab(ylab_text)+
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "None") +
    guides(fill=guide_legend(ncol=1))

  # Deaths H:
  data_deaths_95 <- data.frame(data_country$time, data_deathsh$death_min,
                               data_deathsh$death_max)
  names(data_deaths_95) <- c("time", "death_min", "death_max")
  data_deaths_95$key <- rep("nintyfive", length(data_deaths_95$time))
  data_deaths_50 <- data.frame(data_country$time, data_deathsh$death_min2,
                               data_deathsh$death_max2)
  names(data_deaths_50) <- c("time", "death_min", "death_max")
  data_deaths_50$key <- rep("fifty", length(data_deaths_50$time))
  data_deaths <- rbind(data_deaths_95, data_deaths_50)
  levels(data_deaths$key) <- c("ninetyfive", "fifty")

  if(language_english) ylab_text <- "Deaths H"
  else ylab_text <- "Dnevno stevilo umrlih\n"

  if(language_english) xlab_text <- "Date"
  else xlab_text <- "Datum"

  p2_h <-   ggplot(data_deathsh, aes(x = time)) +
    geom_bar(data = data_deathsh, aes(y = deaths, fill = "reported"),
             fill = "coral4", stat='identity', alpha=0.5) +
    geom_ribbon(
      data = data_deaths,
      aes(ymin = death_min, ymax = death_max, fill = key)) +
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b")) +
    scale_y_continuous(expand = c(0, 0), labels = comma) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("deepskyblue4", 0.55),
                                 alpha("deepskyblue4", 0.45))) +
    xlab(xlab_text)+
    ylab(ylab_text)+
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "None") +
    guides(fill=guide_legend(ncol=1))

  # Deaths C:
  data_deaths_95 <- data.frame(data_country$time, data_deathsc$death_min,
                               data_deathsc$death_max)
  names(data_deaths_95) <- c("time", "death_min", "death_max")
  data_deaths_95$key <- rep("nintyfive", length(data_deaths_95$time))
  data_deaths_50 <- data.frame(data_country$time, data_deathsc$death_min2,
                               data_deathsc$death_max2)
  names(data_deaths_50) <- c("time", "death_min", "death_max")
  data_deaths_50$key <- rep("fifty", length(data_deaths_50$time))
  data_deaths <- rbind(data_deaths_95, data_deaths_50)
  levels(data_deaths$key) <- c("ninetyfive", "fifty")

  if(language_english) ylab_text <- "Deaths C"
  else ylab_text <- "Dnevno stevilo umrlih\n"

  if(language_english) xlab_text <- "Date"
  else xlab_text <- "Datum"

  p2_c <-   ggplot(data_deathsc, aes(x = time)) +
    geom_bar(data = data_deathsc, aes(y = deaths, fill = "reported"),
             fill = "coral4", stat='identity', alpha=0.5) +
    geom_ribbon(
      data = data_deaths,
      aes(ymin = death_min, ymax = death_max, fill = key)) +
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b")) +
    scale_y_continuous(expand = c(0, 0), labels = comma) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("deepskyblue4", 0.55),
                                 alpha("deepskyblue4", 0.45))) +
    xlab(xlab_text)+
    ylab(ylab_text)+
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "None") +
    guides(fill=guide_legend(ncol=1))

  # plot_labels <- c(
  #   "Popolna zapora",
  #   # "Zaprtje obcin",
  #   "Prepoved javnih dogodkov",
  #   "Zaprtje sol",
  #   "Samoizolacija",
  #   "Socialno distanciranje")

  # Rt:
  plot_labels <- covariates_country_long$key
  plot_values <- c(21, 22, 23, 24, 25, 12:20,26:40)[1:length(plot_labels)]

  # Plotting interventions
  data_rt_95 <- data.frame(data_country$time,
                           data_country$rt_min, data_country$rt_max)
  names(data_rt_95) <- c("time", "rt_min", "rt_max")
  data_rt_95$key <- rep("nintyfive", length(data_rt_95$time))
  data_rt_50 <- data.frame(data_country$time, data_country$rt_min2,
                           data_country$rt_max2)
  names(data_rt_50) <- c("time", "rt_min", "rt_max")
  data_rt_50$key <- rep("fifty", length(data_rt_50$time))
  data_rt <- rbind(data_rt_95, data_rt_50)
  levels(data_rt$key) <- c("ninetyfive", "fifth")

  if(language_english) interven <- "Interventions"
  else interven <- "Intervencije"

  p3 <- ggplot(data_country)

  if(!fit.spline){
    p3 <- p3 +
      geom_stepribbon(data = data_rt, aes(x = time, ymin = rt_min, ymax = rt_max,
                                          group = key,
                                          fill = key))
  } else{
    p3 <- p3 +
      geom_ribbon(data = data_rt, aes(x = time, ymin = rt_min, ymax = rt_max,
                                          group = key,
                                          fill = key))
  }

  p3 <- p3 +
    geom_hline(yintercept = 1, color = 'black', size = 0.1) +
    geom_segment(data = covariates_country_long,
                 aes(x = value, y = 0, xend = value, yend = max(x)),
                 linetype = "dashed", colour = "grey", alpha = 0.75) +
    geom_point(data = covariates_country_long, aes(x = value,
                                                   y = x,
                                                   group = key,
                                                   shape = key,
                                                   col = key), size = 2) +
    xlab("") +
    ylab(expression(R[t])) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("seagreen", 0.75), alpha("seagreen", 0.5))) +
    scale_shape_manual(name = interven, labels = plot_labels
                       ,values = plot_values) +
    scale_colour_discrete(name = interven, labels = plot_labels) +
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b"),
                 limits = c(data_country$time[1],
                            data_country$time[length(data_country$time)])) +
    scale_y_continuous(expand = expansion(mult=c(0,0.1))) +
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1)) +
    theme(legend.position="right")
  if (country == 'United Kingdom')
    country = 'United_Kingdom'

  ## Confirmed cases:
  data_cases_952 <- data.frame(data_country$time, data_country2$cases_min,
                              data_country2$cases_max)
  names(data_cases_952) <- c("time", "cases_min", "cases_max")
  data_cases_952$key <- rep("nintyfive", length(data_cases_952$time))
  data_cases_502 <- data.frame(data_country$time, data_country2$cases_min2,
                              data_country2$cases_max2)
  names(data_cases_502) <- c("time", "cases_min", "cases_max")
  data_cases_502$key <- rep("fifty", length(data_cases_502$time))
  data_cases2 <- rbind(data_cases_952, data_cases_502)
  levels(data_cases2$key) <- c("ninetyfive", "fifty")

  if(language_english) ylab_text <- "Daily number of confirmed infections\n"
  else ylab_text <- "Dnevno stevilo potrjeno okuzenih\n"

  p4 <-   ggplot(data_country, aes(x = time)) +
    geom_bar(data = data_country, aes(y = reported_cases, fill = "reported"),
             fill = "coral4", stat='identity', alpha=0.5) +
    geom_ribbon(
      data = data_cases2,
      aes(ymin = cases_min, ymax = cases_max, fill = key)) +
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b")) +
    scale_y_continuous(expand = c(0, 0), labels = comma) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("deepskyblue4", 0.55),
                                 alpha("deepskyblue4", 0.45))) +
    xlab(xlab_text)+
    ylab(ylab_text)+                                 ### Nina spremenila
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "None") +
    guides(fill=guide_legend(ncol=1))

  ## Hospitalization:
  data_hosp_95 <- data.frame(data_country$time, data_hosp$hosp_min,
                               data_hosp$hosp_max)
  names(data_hosp_95) <- c("time", "hosp_min", "hosp_max")
  data_hosp_95$key <- rep("nintyfive", length(data_hosp_95$time))
  data_hosp_50 <- data.frame(data_country$time, data_hosp$hosp_min2,
                               data_hosp$hosp_max2)
  names(data_hosp_50) <- c("time", "hosp_min", "hosp_max")
  data_hosp_50$key <- rep("fifty", length(data_hosp_50$time))
  data_hosp_ci <- rbind(data_hosp_95, data_hosp_50)
  levels(data_hosp_ci$key) <- c("ninetyfive", "fifty")

  if(language_english) ylab_text <- "Number of hospitalized patients\n"
  else ylab_text <- "Stevilo hospitaliziranih\n"

  p5 <- ggplot(data_country, aes(x = time)) +
    geom_bar(data = data_hosp, aes(x=time, y = reported_hosp, fill = "reported"),
             fill = "coral4", stat='identity', alpha=0.5) +
    geom_ribbon(
      data = data_hosp_ci,
      aes(ymin = hosp_min, ymax = hosp_max, fill = key)) +
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b")) +
    scale_y_continuous(expand = c(0, 0), labels = comma) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("deepskyblue4", 0.55),
                                 alpha("deepskyblue4", 0.45))) +
    xlab(xlab_text)+
    ylab(ylab_text)+                                 ### Nina spremenila
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "None") +
    guides(fill=guide_legend(ncol=1))

  ## Hospitalization IN:
  data_hosp_95 <- data.frame(data_country$time, data_hospo$hosp_min,
                             data_hospo$hosp_max)
  names(data_hosp_95) <- c("time", "hosp_min", "hosp_max")
  data_hosp_95$key <- rep("nintyfive", length(data_hosp_95$time))
  data_hosp_50 <- data.frame(data_country$time, data_hospo$hosp_min2,
                             data_hospo$hosp_max2)
  names(data_hosp_50) <- c("time", "hosp_min", "hosp_max")
  data_hosp_50$key <- rep("fifty", length(data_hosp_50$time))
  data_hosp_ci <- rbind(data_hosp_95, data_hosp_50)
  levels(data_hosp_ci$key) <- c("ninetyfive", "fifty")

  if(language_english) ylab_text <- "Hospitalizations IN\n"
  else ylab_text <- "Stevilo hospitaliziranih\n"

  p5_IN <- ggplot(data_country, aes(x = time)) +
    geom_bar(data = data_hospo, aes(x=time, y = reported_hosp, fill = "reported"),
             fill = "coral4", stat='identity', alpha=0.5) +
    geom_ribbon(
      data = data_hosp_ci,
      aes(ymin = hosp_min, ymax = hosp_max, fill = key)) +
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b")) +
    scale_y_continuous(expand = c(0, 0), labels = comma) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("deepskyblue4", 0.55),
                                 alpha("deepskyblue4", 0.45))) +
    xlab(xlab_text)+
    ylab(ylab_text)+                                 ### Nina spremenila
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "None") +
    guides(fill=guide_legend(ncol=1))

  ## Hospitalization OUT:
  data_hosp_95 <- data.frame(data_country$time, data_hospi$hosp_min,
                             data_hospi$hosp_max)
  names(data_hosp_95) <- c("time", "hosp_min", "hosp_max")
  data_hosp_95$key <- rep("nintyfive", length(data_hosp_95$time))
  data_hosp_50 <- data.frame(data_country$time, data_hospi$hosp_min2,
                             data_hospi$hosp_max2)
  names(data_hosp_50) <- c("time", "hosp_min", "hosp_max")
  data_hosp_50$key <- rep("fifty", length(data_hosp_50$time))
  data_hosp_ci <- rbind(data_hosp_95, data_hosp_50)
  levels(data_hosp_ci$key) <- c("ninetyfive", "fifty")

  if(language_english) ylab_text <- "Hospitalizations OUT\n"
  else ylab_text <- "Stevilo hospitaliziranih\n"

  p5_OUT <- ggplot(data_country, aes(x = time)) +
    geom_bar(data = data_hospi, aes(x=time, y = reported_hosp, fill = "reported"),
             fill = "coral4", stat='identity', alpha=0.5) +
    geom_ribbon(
      data = data_hosp_ci,
      aes(ymin = hosp_min, ymax = hosp_max, fill = key)) +
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b")) +
    scale_y_continuous(expand = c(0, 0), labels = comma) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("deepskyblue4", 0.55),
                                 alpha("deepskyblue4", 0.45))) +
    xlab(xlab_text)+
    ylab(ylab_text)+                                 ### Nina spremenila
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "None") +
    guides(fill=guide_legend(ncol=1))

  ## ICUs:
  data_icu_95 <- data.frame(data_country$time, data_icu$icu_min,
                             data_icu$icu_max)
  names(data_icu_95) <- c("time", "icu_min", "icu_max")
  data_icu_95$key <- rep("nintyfive", length(data_icu_95$time))
  data_icu_50 <- data.frame(data_country$time, data_icu$icu_min2,
                             data_icu$icu_max2)
  names(data_icu_50) <- c("time", "icu_min", "icu_max")
  data_icu_50$key <- rep("fifty", length(data_icu_50$time))
  data_icu_ci <- rbind(data_icu_95, data_icu_50)
  levels(data_icu_ci$key) <- c("ninetyfive", "fifty")

  if(language_english) ylab_text <- "Number of patients in ICU\n"
  else ylab_text <- "Stevilo bolnikov na intenzivni negi\n"

  p6 <- ggplot(data_country, aes(x = time)) +
    geom_bar(data = data_icu, aes(x=time, y = reported_icu, fill = "reported"),
             fill = "coral4", stat='identity', alpha=0.5) +
    geom_ribbon(
      data = data_icu_ci,
      aes(ymin = icu_min, ymax = icu_max, fill = key)) +
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b")) +
    scale_y_continuous(expand = c(0, 0), labels = comma) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("deepskyblue4", 0.55),
                                 alpha("deepskyblue4", 0.45))) +
    xlab(xlab_text)+
    ylab(ylab_text)+                                 ### Nina spremenila
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "None") +
    guides(fill=guide_legend(ncol=1))

  ## ICUs IN:
  data_icu_95 <- data.frame(data_country$time, data_icui$icu_min,
                            data_icui$icu_max)
  names(data_icu_95) <- c("time", "icu_min", "icu_max")
  data_icu_95$key <- rep("nintyfive", length(data_icu_95$time))
  data_icu_50 <- data.frame(data_country$time, data_icui$icu_min2,
                            data_icui$icu_max2)
  names(data_icu_50) <- c("time", "icu_min", "icu_max")
  data_icu_50$key <- rep("fifty", length(data_icu_50$time))
  data_icu_ci <- rbind(data_icu_95, data_icu_50)
  levels(data_icu_ci$key) <- c("ninetyfive", "fifty")

  if(language_english) ylab_text <- "ICU IN\n"
  else ylab_text <- "Stevilo bolnikov na intenzivni negi\n"

  p6_IN <- ggplot(data_country, aes(x = time)) +
    geom_bar(data = data_icui, aes(x=time, y = reported_icu, fill = "reported"),
             fill = "coral4", stat='identity', alpha=0.5) +
    geom_ribbon(
      data = data_icu_ci,
      aes(ymin = icu_min, ymax = icu_max, fill = key)) +
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b")) +
    scale_y_continuous(expand = c(0, 0), labels = comma) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("deepskyblue4", 0.55),
                                 alpha("deepskyblue4", 0.45))) +
    xlab(xlab_text)+
    ylab(ylab_text)+                                 ### Nina spremenila
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "None") +
    guides(fill=guide_legend(ncol=1))

  ## ICUs OUT:
  data_icu_95 <- data.frame(data_country$time, data_icuo$icu_min,
                            data_icuo$icu_max)
  names(data_icu_95) <- c("time", "icu_min", "icu_max")
  data_icu_95$key <- rep("nintyfive", length(data_icu_95$time))
  data_icu_50 <- data.frame(data_country$time, data_icuo$icu_min2,
                            data_icuo$icu_max2)
  names(data_icu_50) <- c("time", "icu_min", "icu_max")
  data_icu_50$key <- rep("fifty", length(data_icu_50$time))
  data_icu_ci <- rbind(data_icu_95, data_icu_50)
  levels(data_icu_ci$key) <- c("ninetyfive", "fifty")

  if(language_english) ylab_text <- "ICU OUT\n"
  else ylab_text <- "Stevilo bolnikov na intenzivni negi\n"

  p6_OUT <- ggplot(data_country, aes(x = time)) +
    geom_bar(data = data_icuo, aes(x=time, y = reported_icu, fill = "reported"),
             fill = "coral4", stat='identity', alpha=0.5) +
    geom_ribbon(
      data = data_icu_ci,
      aes(ymin = icu_min, ymax = icu_max, fill = key)) +
    scale_x_date(date_breaks = "weeks", labels = date_format("%e %b")) +
    scale_y_continuous(expand = c(0, 0), labels = comma) +
    scale_fill_manual(name = "", labels = c("50%", "95%"),
                      values = c(alpha("deepskyblue4", 0.55),
                                 alpha("deepskyblue4", 0.45))) +
    xlab(xlab_text)+
    ylab(ylab_text)+                                 ### Nina spremenila
    theme_pubr(base_family="sans") +
    theme(axis.text.x = element_text(angle = 45, hjust = 1),
          legend.position = "None") +
    guides(fill=guide_legend(ncol=1))

  # Special plot settings for mobile
  p3_mobile <- p3
  # +
  #   theme(legend.position="below")

  # Plots for Web, Desktop version
  dir.create("combinedTDI/web/figures/desktop/", showWarnings = FALSE, recursive = TRUE)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/desktop/", country, '_', filename2, "_infections", ".svg"),
            p1, base_height = 4, base_asp = 1.618)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/desktop/", country, '_', filename2, "_deaths", ".svg"),
            p2, base_height = 4, base_asp = 1.618)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/desktop/", country, '_', filename2, "_rt", ".svg"),
            p3, base_height = 4, base_asp = 1.618 * 2)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/desktop/", country, '_', filename2, "_cases", ".svg"),
            p4, base_height = 4, base_asp = 1.618)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/desktop/", country, '_', filename2, "_hosp", ".svg"),
            p5, base_height = 4, base_asp = 1.618)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/desktop/", country, '_', filename2, "_icu", ".svg"),
            p6, base_height = 4, base_asp = 1.618)

  # Plots for Web, Mobile version
  dir.create("combinedTDI/web/figures/mobile/", showWarnings = FALSE, recursive = TRUE)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/mobile/", country, '_', filename2, "_infections", ".svg"),
            p1, base_height = 4, base_asp = 1.1)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/mobile/", country, '_', filename2, "_deaths", ".svg"),
            p2, base_height = 4, base_asp = 1.1)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/mobile/", country, '_', filename2, "_rt", ".svg"),
            p3_mobile, base_height = 4, base_asp = 1.5)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/mobile/", country, '_', filename2, "_cases", ".svg"),
            p4, base_height = 4, base_asp = 1.1)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/mobile/", country, '_', filename2, "_hosp", ".svg"),
            p5, base_height = 4, base_asp = 1.1)
  cowplot::save_plot(filename = paste0("combinedTDI/web/figures/mobile/", country, '_', filename2, "_icu", ".svg"),
            p6, base_height = 4, base_asp = 1.1)

  # Save this for plot-forecast:
  save(p1, p3_mobile, file = paste0("combinedTDI/results/", filename2, '_', "graphs.RData"))

  #
  # p <- plot_grid(p1, p3,
  #                ncol = 2, rel_widths = c(1, 1))
  # p2 <- plot_grid(p2, p4,
  #           ncol = 2, rel_widths = c(1, 1))
  # p3 <- plot_grid(p5, p6,
  #                 ncol = 2, rel_widths = c(1, 1))
  # p4 <- plot_grid(p5_IN, p5_OUT,
  #                 ncol = 2, rel_widths = c(1, 1))
  # p5 <- plot_grid(p6_IN, p6_OUT,
  #                 ncol = 2, rel_widths = c(1, 1))
  # p <- plot_grid(p, p2, p3, p4, p5, nrow=5)
  # cowplot::save_plot(filename = paste0("figures/", country, "_three_pannel_", filename2, ".pdf"),
                     # p, base_width = 14*2/3, base_height=7.42*5/2)
  p <- plot_grid(p1, p4, p3,
                 ncol = 3)
  p2 <- plot_grid(p2, p2_h, p2_c,
                 ncol = 3)
  p3 <- plot_grid(p5, p5_IN, p5_OUT,
                  ncol = 3)
  p4 <- plot_grid(p6, p6_IN, p6_OUT,
                  ncol = 3)

  p <- plot_grid(p, p2, p3, p4, nrow=4)
  cowplot::save_plot(filename = paste0("combinedTDI/figures/", country, "_three_pannel_", filename2, ".pdf"),
                     p, base_width = 14, base_height=7.42*2)

}

make_three_pannel_plot()


# author:       Stephanie Hyland (sh985@cornell.edu)
# date:         January 2014
# description:  script to analyse clusters manually. Find mean distance to centroid in each cluster, take top 10 clusters by this metric, see members
# usage:        source('analyse_clusters.r') in R

#args<-commandArgs(TRUE)
#filepath<-args[1]

print(list.files())
filepath<-readline("File? ")
#filepath = 'm_vecs.clusters.95.txt'

cwords<-function(data,c){
    clust<-da[da$"clust"==c,]
    print(clust$"word")
    browser()
}

da<-read.table(filepath,as.is=TRUE)
names(da)<-c("word","clust","dist")

n_clust<-range(da$"clust")[2]
clust_means<-rep(0,n_clust)
clust_sd<-rep(0,n_clust)

for (c in seq(n_clust)){
    clust<-da[da$"clust"==c,]
    mean_dist<-mean(clust$"dist")
    sd_dist<-sd(clust$"dist")
    clust_means[c]<-mean_dist
    clust_sd[c]<-sd_dist
}

top_10<-order(clust_means)[1:10]
for (c in top_10){
    cat("Cluster: ",c,", mean distance: ",clust_means[c],",sd: ",clust_sd[c],"\n",sep="")
    cwords(da,c)
}

# Now, let's plot some histograms.
cat("Plotting histograms!\n")
pdf(paste("hists_",n_clust,".pdf",sep=""))
par(mfrow=c(2,1))
hist(da$"dist",breaks=50,col="palegreen3",main=paste("Histogram of distance to centroid (n_clust =",n_clust,")"),xlab="Point-centroid distance")
hist(clust_means,breaks=50,col="deepskyblue2",main=paste("Histogram of mean point-centroid distance in cluster (n_clust =",n_clust,")"),xlab="Mean point-centroid distance in cluster")
dev.off()


###################################################################
######      Implementation of multiple matrix regression     ######
######     with randomization via lgrMMRR of PopGenReport    ######
###################################################################

##### Load required packages #####
library(geosphere)
library(PopGenReport)
library(devtools)
library(VPLandscapeGenetics)
data(VPLandscapeGenetics)

##### Calculate genetic distances #####
##SS
nei.SS <- as.matrix(pairwise_Gst_Nei(genind.SS, linearized = TRUE))
nei.SS <- as.matrix(nei.SS)
dimnames(nei.SS) <- list(1:90, 1:90)

##WF
nei.WF <- as.matrix(pairwise_Gst_Nei(genind.WF, linearized = TRUE))
nei.WF <- as.matrix(nei.WF)
dimnames(nei.WF) <- list(1:87, 1:87)

##### Calculate geographic distances #####
##SS
geoDist.SS <- as.matrix(distm(latLong.SS, fun=distGeo))
geoDist.SS <- geoDist.SS/1000
dim(geoDist.SS) <- c(90, 90)
dimnames(geoDist.SS) <- list(1:90, 1:90)

##WF
geoDist.WF <- as.matrix(distm(latLong.WF, fun=distGeo))
geoDist.WF <- geoDist.WF/1000
dim(geoDist.WF) <- c(87, 87)
dimnames(geoDist.WF) <- list(1:87, 1:87)

##### Calculate distance matricies for environmental metrics #####
##SS
ltRds.SS <- as.numeric(as.matrix(dist(envData.SS$SS_lightroads, method="euclidean")))
secRds.SS <- as.numeric(as.matrix(dist(envData.SS$SS_secondaryroads, method="euclidean")))
prmRds.SS <- as.numeric(as.matrix(dist(envData.SS$SS_primaryroads, method="euclidean")))

dim(ltRds.SS) <- c(90, 90)
dimnames(ltRds.SS) <- list(1:90, 1:90)
dim(secRds.SS) <- c(90, 90)
dimnames(secRds.SS) <- list(1:90, 1:90)
dim(prmRds.SS) <- c(90, 90)
dimnames(prmRds.SS) <- list(1:90, 1:90)

envDist_list.SS <- list(ltRds.SS, secRds.SS, prmRds.SS)

##WF
ltRds.WF <- as.numeric(as.matrix(dist(envData.WF$WF_lightroads, method="euclidean")))
secRds.WF <- as.numeric(as.matrix(dist(envData.WF$WF_secondaryroads, method="euclidean")))
prmRds.WF <- as.numeric(as.matrix(dist(envData.WF$WF_primaryroads, method="euclidean")))

dim(ltRds.WF) <- c(87, 87)
dimnames(ltRds.WF) <- list(1:87, 1:87)
dim(secRds.WF) <- c(87, 87)
dimnames(secRds.WF) <- list(1:87, 1:87)
dim(prmRds.WF) <- c(87, 87)
dimnames(prmRds.WF) <- list(1:87, 1:87)

envDist_list.WF <- list(ltRds.WF, secRds.WF, prmRds.WF)

##### Perform MMRR analyses #####
##SS
mmrr.ss <- lgrMMRR(nei.SS, envDist_list.SS, geoDist.SS, nperm=10000)

##WF
mmrr.WF <- lgrMMRR(nei.WF, envDist_list.WF, geoDist.WF, nperm=10000)


# Copyright 2020 The MuLT Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

library(ggplot2)
library(scales)

mic <- read.csv('data/output/correlation_clinical_x_response_vars_mic.csv', stringsAsFactors = FALSE)
mic[is.na(mic)] <- 0.0
mic['METHOD'] <- 'Maximal Information Coefficient'
colnames(mic) <- c('IndependentVariable', 'DependentVariable', 'Value', 'Method')

dis <- read.csv('data/output/correlation_clinical_x_response_vars_distcorr.csv', stringsAsFactors = FALSE)
dis[is.na(dis)] <- 0.0
dis['METHOD'] <- 'Distance Correlation'
colnames(dis) <- c('IndependentVariable', 'DependentVariable', 'Value', 'Method')

pea <- read.csv('data/output/correlation_clinical_x_response_vars_pearson.csv', stringsAsFactors = FALSE)
pea[is.na(pea)] <- 0.0
pea['METHOD'] <- 'Pearson Correlation Coefficient'
colnames(pea) <- c('IndependentVariable', 'DependentVariable', 'Value', 'Method')

df <- rbind(mic, dis, pea)

df[df == 'First-line-Therapy'] <- 'Therapy'

#
#
#
df[df == "Maximal Information Coefficient"] <- "MIC"
df[df == "Distance Correlation"] <- "DCR"
df[df == "Pearson Correlation Coefficient"] <- "PCC"

#
#
#
df[df == "Days-to-Progression"] <- "M2DP"
df[df == "Greater18Months-Days-to-Progression"] <- "M2DP_GT18M"
df[df == "Group1-From-Best-Response-FirstLine"] <- "BRFT_G1"
df[df == "Group2-From-Best-Response-FirstLine"] <- "BRFT_G2"
df[df == "Group3-From-Best-Response-FirstLine"] <- "BRFT_G3"
df[df == "Group4-From-Best-Response-FirstLine"] <- "BRFT_G4"
df[df == "Group5-From-Best-Response-FirstLine"] <- "BRFT_G5"
df[df == "Best-Response-FirstLine-ID"] <- "BRFT"

p <-
  ggplot(df) +
  geom_bar(stat='identity', aes(x=IndependentVariable, y=Value)) +
  coord_flip() +
  facet_grid(DependentVariable ~ Method,  scales='free_x') +
  ylab('Correlation Metric') + xlab('Independent Variable') +
  theme(text=element_text(size=11, family='serif'),
        axis.text.x=element_text(angle = 45, hjust = 1))

ggsave('images/histogram_clinical_correlation_all_dep_vars.pdf', p, units='cm', width=16, height=24)


p_a1 <- cbind(x = c(0, 0, 0.75, 1, 0.5, 0.8, 0.69, 0),
              y = c(0, 1, 1, 0.8, 0.7, 0.6, 0, 0))
p_a2 <- cbind(x = c(0.2, 0.2, 0.4, 0.2),
              y = c(0.25, 0.45, 0.4, 0.25))
p_b1 <- cbind(x = c(0.69, 0.8, 1.1, 1.23, 0.69),
              y = c(0, 0.6, 0.63, 0.3, 0))

library(raster)
spoly <- spPolygons(list(p_a1, p_a2[rev(seq(nrow(p_a2))), ]), p_b1, attr = data.frame(num = c(1, 2), name = c("poly1", "poly2"), property = c(0.1, 3), stringsAsFactors = FALSE))
library(spbabel)
tab <- sptable(spoly)

plot(y~x, tab, asp = 1)
plot(spoly)
text(y~x, tab, lab = seq(nrow(tab)))
axis(1);axis(2)
tab$order <- seq(nrow(tab))
#gtab <- tab  %>% arrange(cump, desc(order))
library(grid)
grid.newpage()
grid.circle(0.3, 0.38, r = 0.02, gp = gpar(fill = "red"))

grid.polygon(tab$x, tab$y, id = tab$cump, gp = gpar(fill  =scales::alpha(c("grey", "NA", "brown"), c(1, 1, 1))))

str(polygonGrob(tab$x, tab$y, id = tab$cump, gp = gpar(fill  =scales::alpha(c("grey", "NA", "brown"), c(1, 1, 1)))))

library(ggplot2)
g <- ggplot(tab) + aes(x = x, y = y) +  geom_polygon(aes(fill = object, group = object))

ggplot(tab) + aes(x = x, y = y) + geom_point(data = data.frame(x = 0.3, y = 0.38) + geom_polygon(aes(fill = object, group = object))

ggplot(data.frame(x = 0.3, y = 0.38)) + aes(x = x, y = y) + geom_point()

ggplot(tab) + aes(x = x, y = y) + geom_polygon(aes(fill = part, group = object))

partord <- function(x) {
  unlist(lapply(split(x, x), function(y) seq_along(y)))
}

library(raster)
spoly <- spPolygons(list(p_a1, p_a2[rev(seq(nrow(p_a2))), ]), p_b1, attr = data.frame(num = c(1, 2), name = c("poly1", "poly2"), property = c(0.1, 3), stringsAsFactors = FALSE))

library(spbabel)
tab <- sptable(spoly)

d <- as.data.frame(geom(spoly))
library(ggplot2)
ggplot(d) + aes(x = x, y = y, fill = object, group = object) + geom_polygon()

a <- as.data.frame(head(geom(spoly[1,], sepNA = TRUE), -1))
library(grid)
grid.newpage()
grid.polygon(rev(d$x), rev(d$y), id = rev(d$part),  gp = gpar(fill = c("grey")))
?grid.polygon

polyr <- spPolygons(list(p_a1[nrow(p_a1):1, ], p_a2), p_b1, attr = data.frame(num = c(1, 2), name = c("poly1", "poly2"), property = c(0.1, 3), stringsAsFactors = FALSE))

poly1 <- spPolygons(list(p_a1, p_a2), p_b1, attr = data.frame(num = c(1, 2), name = c("poly1", "poly2"), property = c(0.1, 3), stringsAsFactors = FALSE))
poly2 <- spPolygons(list(p_a1[nrow(p_a1):1, ], p_a2[nrow(p_a2):1, ]), p_b1, attr = data.frame(num = c(1, 2), name = c("poly1", "poly2"), property = c(0.1, 3), stringsAsFactors = FALSE))

ggplot(as.data.frame(head(geom(poly1[1,], sepNA = TRUE), -1))) + aes(x = x, y = y, fill = object, group = object) + geom_polygon()

library(ggplot2)
library(maptools)
d <- as.data.frame(head(geom(poly[1,], sepNA = TRUE), -1))
plot(d$x, d$y)
polypath(d$x, d$y, col = "grey", rule = "evenodd")
ggplot(d) + aes(x = x, y = y, fill = object, group = object) + geom_polygon()

library(gris)
g <- gris(poly)

p <- p_a1
pr <- p[nrow(p):1, ]

str(Polygon(p))
str(Polygon(pr))

str(Polygons(list(Polygon(p)), "1"))
str(Polygons(list(Polygon(pr)), "1"))

str(Polygons(list(Polygon(p)), "1"))
str(Polygons(list(Polygon(pr)), "1"))

str(SpatialPolygons(list(Polygons(list(Polygon(p)), "1"))))
str(SpatialPolygons(list(Polygons(list(Polygon(pr)), "1"))))

SpatialPolygons(list(Polygons(list(Polygon(p)),  "1")))@polygons[[1]]@Polygons[[1]]@coords
SpatialPolygons(list(Polygons(list(Polygon(pr)), "1")))@polygons[[1]]@Polygons[[1]]@coords


#' sweep
#'
#' Sweep a vector through a float matrix.
#'
#' @details
#' Note that if the length of \code{STATS} does not recycle exactly across
#' \code{MARGIN}, the results here will differ slightly from the results of
#' base R.
#'
#' @param x
#' A float vector/matrix.
#' @param MARGIN
#' 1 (rows) or 2 (columns)
#' @param STATS
#' Vector to sweep out.
#' @param FUN
#' Sweeping function; must be one of \code{"+"}, \code{"-"}, \code{"*"},  or
#' \code{"/"}.
#' @param check.margin
#' Should x/STATS margin lengths be checked?
#' @param ...
#' Theoretically these are additional arguments passed to an arbitrary function.
#' However, we only support basic arithmetic, so they are ignored.
#'
#' @return
#' A matrix of the same type as the highest precision input.
#'
#' @examples
#' library(float)
#'
#' s = flrunif(10, 3)
#' sweep(s, 2, fl(1))
#'
#' @useDynLib float R_sweep_spm
#' @name sweep
#' @rdname sweep
NULL

sweep_float32 = function(x, MARGIN, STATS, FUN="-", check.margin=TRUE, ...)
{
  if (is.double(STATS))
    return(sweep(dbl(x), MARGIN, STATS, FUN, check.margin, ...))
  else if (!is.logical(STATS) && !is.integer(STATS) && !is.float(STATS))
    stop("non-numeric argument to binary operator")

  if (isavec(x))
    stop("Error in array(STATS, dims[perm]) : 'dims' cannot be of length 0", call.=FALSE)

  if (!isTRUE(all.equal(MARGIN, 1)) && !isTRUE(all.equal(MARGIN, 2)))
    stop("missing value where TRUE/FALSE needed")

  if (!is.character(FUN))
    stop("only '+', '-', '*', and '/' are implemented for floats")

  FUN = match.arg(FUN, c('+', '-', '*', '/'))

  if (isTRUE(check.margin))
  {
    if ((MARGIN == 1 && NROW(x)%%length(STATS)) || (MARGIN == 2 && NCOL(x)%%length(STATS)))
      warning("STATS does not recycle exactly across MARGIN")
  }

  ret = .Call(R_sweep_spm, DATA(x), as.integer(MARGIN), STATS, FUN)
  float32(ret)
}

#' @rdname sweep
#' @export
setMethod("sweep", signature(x="float32"), sweep_float32)


library(dplyr)
library(readr)
library(tidyr)
library(caret)
library(caTools)
library(ROCit)
library(caretEnsemble)
source('./scr/draw_confusion_plot.r')

data_test <- read.csv('./data/processed/processed_test.csv')
data_train <- read.csv('./data/processed/processed_train.csv')

# Loading models
model_c5 <- readRDS('./models/C5.0.rds')
model_egb <- readRDS('./models/EGB.rds')
model_rf <- readRDS('./models/RF.rds')
model_c5_full <- readRDS('./models/C5.0_full.rds')

# Comparing models

# TODO how to incorporate ensemble models?
cv_results <-
    resamples(list(
        C5 = model_c5,
        EGB = model_egb,
        RF = model_rf
    ))

# Random Forest and Extreme Gradient boosting are clearly the best models
summary(cv_results)
png(filename='./figures/dotplot_comparison.png')
dotplot(cv_results)
dev.off()

# Model correlations are low so combining models makes sense
modelCor(cv_results)
png(filename='./figures/model_prediction_correlations.png')
splom(cv_results)
dev.off()

predictions_c5 <-
    predict(model_c5, newdata = data_test, type = 'prob')
predictions_egb <-
    predict(model_egb, newdata = data_test, type = 'prob')
predictions_rf <-
    predict(model_rf, newdata = data_test, type = 'prob')
# In the ensemble model the predictions are for the wrong category
predictions_c5_rf <-
    1 - predict(model_c5_rf, newdata = data_test, type = 'prob')
png(filename='./figures/AUC_comparison.png')
rf_rocit <- rocit(predictions_rf$Sony, data_test$brand)
c5_rocit <- rocit(predictions_c5$Sony, data_test$brand)
egb_rocit <- rocit(predictions_egb$Sony, data_test$brand)
plot(rf_rocit, col=1, YIndex = FALSE, legend = FALSE)
lines(c5_rocit$TPR ~ c5_rocit$FPR, col=2, lwd=2)
lines(egb_rocit$TPR ~ egb_rocit$FPR, col=4, lwd=2)
legend('bottomright', col = c(1,2,4), c('RF', 'C5.0', 'EGB'), lwd=2)
dev.off()

# plotting confusion matrices
predictions_c5_classes <-
    predict(model_c5, newdata = data_test)
predictions_egb_classes <-
    predict(model_egb, newdata = data_test)
predictions_rf_classes <-
    predict(model_rf, newdata = data_test)

predictions_c5_classes_train <-
    predict(model_c5, newdata = data_train)
predictions_egb_classes_train <-
    predict(model_egb, newdata = data_train)
predictions_rf_classes_train <-
    predict(model_rf, newdata = data_train)

cm_c5 <- confusionMatrix(predictions_c5_classes, data_test$brand)
cm_egb <- confusionMatrix(predictions_egb_classes, data_test$brand)
cm_rf <- confusionMatrix(predictions_rf_classes, data_test$brand)

cm_c5_train <- confusionMatrix(predictions_c5_classes_train, data_train$brand)
cm_egb_train <- confusionMatrix(predictions_egb_classes_train, data_train$brand)
cm_rf_train <- confusionMatrix(predictions_rf_classes_train, data_train$brand)

png(filename='./figures/cm_c5.png')
cm_c5_plot <- draw_confusion_plot(cm_c5, title='C5.0 Confusion Matrix: Test')
dev.off()
png(filename='./figures/cm_egb.png')
cm_egb_plot <- draw_confusion_plot(cm_egb, title='EGB Confusion Matrix: Test')
dev.off()
png(filename='./figures/cm_rf.png')
cm_rf_plot <- draw_confusion_plot(cm_rf, title='RF Confusion Matrix: Test')
dev.off()
png(filename='./figures/cm_c5_train.png')
cm_c5_train_plot <- draw_confusion_plot(cm_c5_train, title='C5.0 Confusion Matrix: Train')
dev.off()
png(filename='./figures/cm_egb_train.png')
cm_egb_train_plot <- draw_confusion_plot(cm_egb_train, title='EGB Confusion Matrix: Train')
dev.off()
png(filename='./figures/cm_rf_train.png')
cm_rf_train_plot <- draw_confusion_plot(cm_rf_train, title='RF Confusion Matrix: Train')
dev.off()

# Model correlations are low there ensembling is a good idea
modelCor(cv_results)
png(filename='./figures/model_prediction_correlations.png')
splom(cv_results)
dev.off()

# Model variable importance
png(filename='./figures/c5_feature_importance.png')
plot(varImp(model_c5, useModel = FALSE), top = 15)
dev.off()
png(filename='./figures/egb_feature_importance.png')
plot(varImp(model_egb, useModel = FALSE), top = 15)
dev.off()
png(filename='./figures/rf_feature_importance.png')
plot(varImp(model_rf, useModel = FALSE), top = 15)
dev.off()
png(filename='./figures/c5_full_feature_importance.png')
plot(varImp(model_c5_full, useModel = FALSE), top = 15)
dev.off()

# Salary is in all models the most important predictor
# Lets draw the distribution of wage and brand preference with the full data
data_full <- bind_rows(data_train, data_test)
png(filename='./figures/brand_salary_distribution.png')
ggplot(data_full, aes(x=salary, color=brand, group=brand)) +
    geom_freqpoly(data=subset(data_full, brand=='Sony'), bins=60, aes(y=(..count..)/sum(..count..))) +
    geom_freqpoly(data=subset(data_full, brand=='Acer'), bins=60, aes(y=(..count..)/sum(..count..))) +
    xlab('Salary') +
    ylab('%') +
    scale_y_continuous(labels=scales::percent)
dev.off()


#' Method of RCTtoolbox: \code{power()}
#'
#' When you create R6 object "RCTtoolbox"
#' via \code{link{create_RCTtoolbox}()}
#' and run its method \code{power()},
#' the method internally implements a function
#' that performs power analysis.
#' The method \code{power()} provides some arguments.
#' See the section "Arguments."
#'
#' @param ctrl character. New control arm.
#' If NULL (default),
#' the first element of \code{treat_levels} is control arm.
#' @param subset subset condition.
#' If NULL (default),
#' full observations of \code{data} is used to implement power analysis.
#' @param sd numeric. standard deviation (default is 1).
#' Absolute mean difference between two groups
#' is calculated by effect size times \code{sd}.
#' @param \dots arguments passed to \code{\link{ttest_power}()}.
#' This function has five arguments: \code{n0}, \code{n1}, \code{d},
#' \code{alpha}, and \code{beta}.
#' This method automatically specify the first two arguments.
#' Users must specify two of
#' \code{d} (effect size), \code{alpha} (significant level),
#' and \code{beta} (power) in dot-dot-dot.
#' This method calculate one missing argument using \code{\link{uniroot}}.
#'
#' @return R6 object with "RCTtoolbox.power.analysis" class.
#' The returned object has following field and methods:
#' \describe{
#'   \item{\code{result}}{Field.
#'     Data frame including estimated result.
#'     See the section "Result Field."}
#'   \item{\code{print()}}{Method.
#'     Print information about the returned object.
#'     Run \code{$print()}.}
#'   \item{\code{table()}}{Method.
#'     Create output table of result.
#'     Run \code{$plot()}.}
#'   \item{\code{summary()}}{Method.
#'     Print result in console.
#'     Run \code{$summary()}.}
#' }
#'
#' @section Result Field:
#'
#' The R6 object with "RCTtoolbox.power.analysis" class has
#' \code{result} field.
#' This field has a data frame with 8 variables:
#' \describe{
#'   \item{n0}{Number of observations of control arm}
#'   \item{n1}{Number of observations of each experimental arm}
#'   \item{arms}{Factor. Experimental arms.
#'     The first level is control arm.}
#'   \item{d}{Effect size}
#'   \item{alpha}{Significant level}
#'   \item{power}{Power}
#'   \item{sd}{Standard deviation}
#'   \item{diff_mean}{Absolute mean difference of two groups}
#' }
#'
#' @section Developer Note:
#'
#' A method \code{power()} provided by R6 object RCTtoolbox
#' implements \code{RCTtoolbox.power.analysis$new()}
#' which generates R6 object with "RCTtoolbox.power.analysis" class.
#' Initialization of R6 object "RCTtoolbox.ttest" run
#' \code{power_calculation(private$dvec, self$data, private$dvec.levels,
#' private$dvec.labels, ...)}
#' where \code{...} accepts arguments explained in the section "Arguments."
#' The first four arguments passed to \code{ttest_multi_mod_arm()} are
#' \describe{
#'   \item{\code{treat}}{Character. Treatment Variable.
#'     The method \code{power()} automatically
#'     passes the private field \code{dvec} of R6 object "RCTtoolbox"
#'     to this argument.}
#'   \item{\code{data}}{data.frame/tibble object that you want to use.
#'     The method \code{power()} automatically
#'     passes a public field \code{data} of R6 object "RCTtoolbox"
#'     to this argument.}
#'   \item{\code{treat_levels}}{character vector. Level of experimental arms.
#'     The first element is control arm.
#'     The method \code{power()} automatically
#'     passes the private field \code{dvec.levels} of R6 object "RCTtoolbox"
#'     to this argument.}
#'   \item{\code{treat_labels}}{character vector. Label of experimental arms
#'     corresponding to \code{treat_levels}.
#'     The method \code{power()} automatically
#'     passes the private field \code{dvec.labels} of R6 object "RCTtoolbox"
#'     to this argument.}
#' }
#'
#' @examples
#' \dontrun{
#' data(RubellaNudge)
#' rct <- create_RCTtoolbox(
#'   atest + avacc ~ treat,
#'   data = RubellaNudge,
#'   treat_levels = LETTERS[1:7]
#' )
#'
#' # calculate effect size
#' rct$power(alpha = 0.05, power = 0.8)$summary()
#' # calculate effect size with control arm "C"
#' rct$power(alpha = 0.05, power = 0.8, ctrl = "C")$result
#' }
#'
#' @name power
#'
NULL


source("code/fonctionsThermo.r")

energies <- read.table("data/energiesPartLibres.txt", header=TRUE)
energies <- energies$Energies

temperatures <- 1:298

partitions <- lapply(temperatures, calcul.partition, valeurs=energies)
partitions <- unlist(partitions)

write.table(partitions, "data/partitionsTemperatureFreePart.txt", row.names=FALSE, col.names=FALSE)