
library(fields)
library(multiApply)
library(RColorBrewer)

plot_probabilistic_forecast <- function(obs_probs, fcst_probs, skill = NULL, skill_threshold = 0, thresholds = NULL, units = NULL, title = '', title_line = 0,
                                        cols = c("Blues", "YlOrBr", "Reds"), category_names = c('Below normal','Normal','Above normal'),
                                        xlabs = NULL, style = 'probabilities', cat_dim = 'bin', time_dim = 'year', auxiliar_lines = TRUE,
                                        legend_cex = 0.6, legend_title = 'Probability of the probabilistic categories', margins = c(0, 4, 3, 0),
                                        width = 2000, height = 1000, res = 200, cut_figure = TRUE, fileout = NULL){
  
  n_categories <- as.numeric(dim(obs_probs)[cat_dim])
  n_obs <- as.numeric(dim(obs_probs)[time_dim])
  n_fcst <- as.numeric(dim(fcst_probs)[time_dim])
  if (is.null(xlabs)){xlabs <- 1:n_fcst}
  
  ## Checks
  stopifnot(identical(n_categories,as.numeric(dim(obs_probs)[cat_dim])))
  stopifnot(identical(n_categories,as.numeric(dim(fcst_probs)[cat_dim])))
  stopifnot(identical(n_categories,as.numeric(length(cols))))
  stopifnot(identical(n_categories,as.numeric(length(category_names))))
  if (!is.null(thresholds)){stopifnot(identical(n_categories,as.numeric(length(thresholds)+1)))}
  
  ## Opening the output file
  if (!is.null(fileout)){png(filename = fileout, width = width, height = height, res = res)}
  
  ## Colours
  .tercileBrewerColorRamp <- function(ncolors, cols){
    ncolors <- ncolors-2
    tmp <- brewer.pal(ncolors, cols[1])
    rval <- data.frame(a = c(tmp[1],tmp[1],tmp), stringsAsFactors=F)
    for (i in 2:length(cols)){
      tmp <- brewer.pal(ncolors, cols[i])
      rval[[letters[i]]] <- c(tmp[1],tmp[1],tmp)
    }
    return(rval)
  }
  
  t.color <- .tercileBrewerColorRamp(ncolors = 10, cols = cols)     
  brks <- c(seq(0,1,length=nrow(t.color)+1))
  
  ## Figure
  # opar <- par(no.readonly=TRUE)
  # par(mar = c(4, 5, 0, 3)) #change lines of the margins
  # par(oma = c(0, 0, 0, 3)) #change size of the outer margins (ROCSS on the right)
  par(oma = margins) #change size of the outer margins (ROCSS on the left)
  
  ## Ensemble spread (hindcast and forecast)
  if (identical(style,'probabilities')){
    
    for (i in 1:n_categories){
      ## Color boxes
      if (i == 1){add <- FALSE} else {add <- TRUE}
      image(1:n_fcst, c(i-mean(1:n_categories)-0.5, i-mean(1:n_categories)+0.5), matrix(ClimProjDiags::Subset(x = fcst_probs, along = cat_dim, indices = i, drop = 'selected')), breaks=brks, col=t.color[[letters[i]]], ylab="", xlab="", asp = 1.5, yaxt="n", bty = "n", axes = FALSE, add = add)      
      ## Legend
      if (i == 1){
        fields::image.plot(add = TRUE, horizontal = T, smallplot = c(0.20, 0.85, seq(from=0.22,by=0.02,length.out=n_categories)[i], seq(from=0.24,by=0.02,length.out=n_categories)[i]), legend.only = TRUE, breaks = brks, col = t.color[,i], zlim=c(0,1), legend.lab = legend_title, legend.cex = legend_cex)
      } else {
        fields::image.plot(add = TRUE, horizontal = T, smallplot = c(0.20, 0.85, seq(from=0.22,by=0.02,length.out=n_categories)[i], seq(from=0.24,by=0.02,length.out=n_categories)[i]), legend.only = TRUE, breaks = brks, lab.breaks=c(rep("", length(brks))), col = t.color[,i], zlim=c(0,1))
      }
    }
    
  } else if (identical(style,'most-likely')){
    
    opacity <- 0.8
    
    ## Most likely tercile
    most_likely <- multiApply::Apply(data = fcst_probs, target_dims = cat_dim, output_dims = cat_dim, 
                                     fun = function(x, n){y <- rep(-1,n); y[which(x == max(x))] <- opacity; return(y)}, 
                                     n = n_categories, ncores = 1)$output1
    ## Color boxes
    for (i in 1:n_categories){
      if (i == 1){add <- FALSE} else {add <- TRUE}
      image(1:n_fcst, c(i-mean(1:n_categories)-0.5, i-mean(1:n_categories)+0.5), matrix(ClimProjDiags::Subset(x = most_likely, along = cat_dim, indices = i, drop = 'selected')), breaks=brks, col=t.color[[letters[i]]], ylab="", xlab="", asp = 1.5, yaxt="n", bty = "n", axes = FALSE, add = add)      
    }
    
  } else if (identical(style,'most-likely-skill')){
    
    opacity_skill <- 0.9
    opacity_noskill <- 0.6
    skill <- array(data = skill, dim = dim(fcst_probs)[time_dim])
    
    ## Most likely tercile (with different opacity depending on the skill)
    most_likely <- multiApply::Apply(data = list(fcst_probs, skill), target_dims = list(cat_dim,NULL), output_dims = cat_dim, 
                                     fun = function(x, n, skill, skill_threshold, opacity_skill, opacity_noskill){
                                       y <- rep(-1,n)
                                       if (skill > skill_threshold & !is.na(skill)){
                                         y[which(x == max(x))] <- opacity_skill
                                       } else if (skill <= skill_threshold & !is.na(skill)) {
                                         y[which(x == max(x))] <- opacity_noskill
                                       }
                                       return(y)
                                       }, 
                                     opacity_skill = opacity_skill, opacity_noskill = opacity_noskill,
                                     skill_threshold = skill_threshold, n = n_categories, ncores = 1)$output1
    
    ## Color boxes
    for (i in 1:n_categories){
      if (i == 1){add <- FALSE} else {add <- TRUE}
      image(1:n_fcst, c(i-mean(1:n_categories)-0.5, i-mean(1:n_categories)+0.5), matrix(ClimProjDiags::Subset(x = most_likely, along = cat_dim, indices = i, drop = 'selected')), breaks=brks, col=t.color[[letters[i]]], ylab="", xlab="", asp = 1, yaxt="n", bty = "n", axes = FALSE, add = add)      
    }
    
  } else {stop('style must be either probabilities or most-likely')}
  
  axis(1, at = 1:n_fcst, labels = xlabs, pos = -(n_categories-mean(1:n_categories)+0.5), cex.axis = 0.6, las = 2)
  points(x = 1:n_obs, y = -(1+n_categories-mean(1:n_categories)) + multiApply::Apply(data = obs_probs, target_dims = cat_dim, fun = function(data){ifelse(all(is.na(data)), NA, which.max(data))})$output1, pch = 21, cex = 0.5, bg = "black")
  # axis(2, at=-1:1, labels=c(paste0("Below (ROCSS = ",round(rocss$cat1, 2),')'), paste0("Normal (ROCSS = ",round(rocss$cat2, 2),')'), paste0("Above (ROCSS = ",round(rocss$cat3, 2),')')), cex.axis=0.6, las="2")
  
  labels <- c()
  for (i in 1:n_categories){
    if (i != mean(1:n_categories)){
      if (!is.null(thresholds)){
        if (i < mean(1:n_categories)){
          sign <- '<'
          threshold <- round(thresholds[i],2)
        } else {
          sign <- '>'
          threshold <- round(thresholds[i-1],2)
        }
        labels <- c(labels, paste0(category_names[i],' (',sign,' ',threshold,units,')'))  
      } else {
        labels <- c(labels, category_names[i])  
      }
    } else {
      labels <- c(labels, category_names[i])
    }
  }
  
  axis(2, at=-(1-mean(1:n_categories)):(1-mean(1:n_categories)), labels=labels, cex.axis=0.6, las="2")
  
  # ## Adding ROCSS (on the right)
  # axis(4, at=c(-1:1,2.3), labels=c(rocss$cat1, rocss$cat2, rocss$cat3, "ROCSS"), las="2", tick = FALSE, cex.axis=0.8)
  
  ## Auxiliary lines
  if (isTRUE(auxiliar_lines)){
    abline(h = seq(from = -(n_categories-mean(1:n_categories)+0.5), to = n_categories-mean(1:n_categories)+0.5, by = 1), lw = 1) ## Horizontal lines
    for (i in seq(from = 0.5, to = n_fcst+1, by = 1)){segments(x0 = i, y0 = -(n_categories-mean(1:n_categories)+0.5), x1 = i, y1 = n_categories-mean(1:n_categories)+0.5, lwd = 1)} ## Vertical lines
  } else if (!isFALSE(auxiliar_lines)){stop('auxiliar_lines must be TRUE or FALSE')}
  
  ## Title
  mtext(title, side=3, line=-5+title_line, adj=0, cex=0.9, font=2)
  
  ## Closing the output file
  if (!is.null(fileout)){
    dev.off()
    # Cropping the figure to remove unnecessary margins
    if (isTRUE(cut_figure)){
      if (identical(style,'probabilities')){
        system(paste0('convert -crop 100%x65%+0%+280% +repage ',fileout,' ',fileout,'_aux'))
      } else if (style %in% c('most-likely','most-likely-skill')){
        system(paste0('convert -crop 100%x75%+80%+55% +repage ',fileout,' ',fileout,'_aux'))
      }
      system(paste0('mv ',fileout,'_aux ',fileout))
    }
  }
 observed <- multiApply::Apply(data = obs_probs, target_dims = cat_dim, fun = function(data){ifelse(all(is.na(data)), NA, which.max(data))})$output1
 return(list(most_likely = most_likely, observed = observed)) 
}

########## Dummy data for testing ##########
# fcst_probs <- array(data = runif(500, min=0, max=1), dim = c(year = length(1960:2020), bin = 3))
# obs_probs <- array(data = runif(500, min=0, max=1), dim = c(year = length(1961:2015), bin = 3))
# rocss <- list(cat1 = 0.5, cat2 = 0.7, cat3 = 0.9)
# title <- paste('Title1','Title2', sep = '\n')
# xlabels <- rep('2000-3000', length(1960:2020))
# plot_probabilistic_forecast(obs = obs_probs, fcst = fcst_probs, title = title, xlabs = xlabels, rocss = rocss, style = 'probabilities', fileout = '~/Desktop/probabilities.png')
# plot_probabilistic_forecast(obs = obs_probs, fcst = fcst_probs, title = title, xlabs = xlabels, rocss = rocss, style = 'most-likely', fileout = '~/Desktop/most-likely.png')

