
## Utility

initialise_MQR <- function(nrow,quantiles){
  
  multipleQuantiles <- matrix(NA,nrow=nrow,ncol=length(quantiles))
  colnames(multipleQuantiles) <- paste0("q",100*quantiles)
  multipleQuantiles <- as.data.frame(multipleQuantiles)
  class(multipleQuantiles) <- c("MultiQR",class(multipleQuantiles))
  return(multipleQuantiles)
  
}

#############################
## Covariance Estimation ####
#############################

Weighted_Cov <- function(u_data,u_time,boundary_threshold=NA,method="vonmises",forcePD=F,covariate_x=NULL,covariate_value=NULL,...){
  
  na_data <- unique(which(is.na(u_data),arr.ind = T)[,1])
  
  u_data <- as.matrix(u_data)
  
  # Handle 0s and 1s: either avoir infty or produce NAs
  u_data <- ifelse(u_data==0,boundary_threshold,u_data)
  u_data <- ifelse(u_data==1,1-boundary_threshold,u_data)
  
  # Transform to Gaussian...
  g_data <- qnorm(u_data)
  rm(u_data)
  
  
  if(method=="vonmises"){
    require(circular)
    
    weights <- dvonmises(x=circular(as.numeric(format(u_time,"%j"))*2*pi/366,
                                    type="angles",units="radians",template = "none",
                                    modulo = "2pi",zero = 0,rotation = "clock"),...)
    data_out <- cov.wt(x = g_data[-na_data,], wt = weights[-na_data],center = F)$cov
    
    # plot(u_time,weights,type="l")
    
  }else if(method=="covariate"){
    if(is.null(covariate_value)){stop("covariate_x must be specified.")}
    if(is.null(covariate_x)){stop("covariate_x must be specified.")}
    if(length(covariate_x)!=length(u_time)){stop("covariate_x must be same length as u_time.")}
    
    data_out <- cov.wt(x = g_data[-na_data,],
                       wt = dnorm(x = covariate_x-covariate_value,sd = 0.2*sd(covariate_x,na.rm = T))[-na_data],
                       center = F)$cov
    
    
  }else{
    stop("Method not recognised...")
  }
  
  if(forcePD){

    data_out <- as.matrix(Matrix::nearPD(data_out,corr=F)$mat)
    
  }
  return(data_out)
}


Extract_Blockdiag <- function(cov,block_size,n_blocks,inverse=F){
  
  block <- matrix(1,block_size,block_size)
  index <- as.matrix(Matrix::bdiag(rep(list(block),n_blocks)))
  
  if(inverse){return(cov*(1-index))}else{return(cov*index)}
}

##################
## Evaluation ####
##################

all_reliability <- function(dt,Names,mfrow=c(1,length(Names)),...){
  
  par(mfrow=mfrow)
  for(N in Names){
    reliability(qrdata = dt[[paste0(N,"_qr")]][which(!dt[[N]]$BadData),],
                realisations = dt[[N]]$node_n[which(!dt[[N]]$BadData)],
                kfolds = dt[[N]]$kfold[which(!dt[[N]]$BadData)],main=N)
  }
  par(mfrow=c(1,1))
  
}


all_det <- function(dt,Names,n=2*7*48,ylim1=c(0,8),cv=NULL,ylim2=c(-1,1)){
  
  if(is.null(cv)){cv<-dt[[Names[1]]][,unique(kfold)]}
  
  par(mfrow=c(2,2))
  
  
  # MAE by lead-time
  for(N in Names){
    temp <- data.table(kfold = dt[[N]]$kfold,
                       LeadTime = dt[[N]]$LeadTime,
                       error = dt[[paste0(N,"_qr")]]$q50-dt[[N]]$node_n,
                       BadData = dt[[N]]$BadData)
    
    Normalise <- 100*dt[[paste0(N,"_norm")]][2]/dt[[N]][BadData==F,max(node)]
    temp2 <- temp[BadData==F & kfold %in% cv,.(NMAE=mean(abs(error),na.rm = T)*Normalise),by=LeadTime]
    setkey(temp2,LeadTime)
    
    if(N==Names[1]){
      plot(temp2$LeadTime,temp2$NMAE,type="l",col=which(Names==N),
           xlab="Lead-time [h]",ylab="NMAE [%]",ylim=ylim1,
           main="Normalised MAE")
    }else{
      lines(temp2$LeadTime,temp2$NMAE,col=which(Names==N),lty=which(Names==N))
    }
    
  }; rm(temp,temp2)
  
  # Bias by lead-time
  for(N in Names){
    temp <- data.table(kfold = dt[[N]]$kfold,
                       LeadTime = dt[[N]]$LeadTime,
                       error = dt[[paste0(N,"_qr")]]$q50-dt[[N]]$node_n,
                       BadData = dt[[N]]$BadData)
    
    temp2 <- temp[BadData==F & kfold %in% cv,.(Bias=mean(error,na.rm = T)),by=LeadTime]
    setkey(temp2,LeadTime)
    
    if(N==Names[1]){
      plot(temp2$LeadTime,temp2$Bias,type="l",col=which(Names==N),
           xlab="Lead-time [h]",ylab="Bias [standardised]",ylim=c(-.2,.2),
           main="Bias by Lead-time")
    }else{
      lines(temp2$LeadTime,temp2$Bias,col=which(Names==N),lty=which(Names==N))
    }
    
  }; rm(temp,temp2)
  
  ## Smoothed error time series
  for(doy in c(T,F)){
    for(i in 1:length(Names)){
      N <- Names[i]
      
      temp <- data.table(kfold = dt[[N]]$kfold,
                         dtm = dt[[N]]$targetTime,
                         BadData = dt[[N]]$BadData,
                         error = dt[[paste0(N,"_qr")]]$q50-dt[[N]]$node_n)
      # error = dt[[N]]$node_n)
      temp[BadData==T,error:=NA]
      temp <- temp[!duplicated(dtm) & kfold %in% cv,]
      if(doy){
        temp$error_smooth <- frollmean(temp[,error],n=n,align = "center",na.rm = T)  
        temp[,dtm:=as.numeric(format(dtm,"%j"))]
        temp <- temp[,.(error_smooth=mean(error_smooth)),by=dtm]
      }else{
        temp$error_smooth <- frollmean(temp[,error],n=n,align = "center",na.rm = T)  
        
      }
      
      
      if(i==1){
        plot(temp$dtm,temp$error_smooth,type="l",col=i,ylim = if(doy){c(-.2,.2)}else{ylim2},
             ylab=if(doy){"Bias [standardised]"}else{"Smoothed Error [standardised]"},
             xlab=if(doy){"Day of Year"}else{"Date/Time"},
             main=if(doy){"Bias by Time of Year"}else{"Rolling Mean Error"})
      }else{
        lines(temp$dtm,temp$error_smooth,col=i,lty=i)
      }
    }
    lines(range(temp$dtm),c(0,0),lty=2,col="red",lwd=2)
  }
  
  # legend("topleft",allNs,col=1:i,lty=1:i,bty="n",ncol = max(1,floor(i/4)))
  par(mfrow=c(1,1))
}

bias_by_covariate <- function(pred,obs,cova,kfold=NULL){
  
  
  
}


WormPlot <- function(U_data,cov_max_lag=NULL,add=F,xlim = c(-3.1,3.1),ylim=c(-0.3,0.3),...){
  
  
  GetConf <- function(x,COV=0){
    z = seq(-4, 4, length.out = 1000)
    n = length(x)
    sample_sd = sd(x) + 2*COV
    se = sample_sd * (sqrt(pnorm(z) * (1 - pnorm(z)) / n) / dnorm(z))
    line = mean(x) + sample_sd * z
    return(list(z=z,
                line = line,
                upper = line + 2*se,
                lower = line - 2*se))
  }
  
  d <- list()
  d$x_sample = qnorm(na.omit(U_data))
  d$x_sample <- d$x_sample[d$x_sample>-Inf & d$x_sample<Inf]
  d$quantile = ppoints(length(d$x_sample))
  q = (rank(d$x_sample) - .5) / length(d$x_sample)
  d$x_theoretical = qnorm(q)
  d$z_theoretical <- qnorm(d$quantile, 0, 1)
  d$line <- mean(d$x_sample) + d$z_theoretical * sd(d$x_sample)
  
  if(!is.null(cov_max_lag)){
    
    COV <- sum(ccf(d$x_sample,d$x_sample,type = "covariance",lag.max = cov_max_lag,plot = F)$acf[-1])
    ConfInt <- GetConf(d$x_sample,COV=COV)
    
  }else{
    ConfInt <- GetConf(d$x_sample,COV=0)  
  }
  
  # Worm Plot
  if(add){
    lines(d$z_theoretical,y=sort(d$x_sample)-d$line,...)
  }else{
    plot(d$z_theoretical,y=sort(d$x_sample)-d$line,type="l",
         xlim = xlim,ylim=ylim,
         xlab="Unit normal quantile",
         ylab = "Deviation",...)
    grid()
    
    ## Top Axis
    pr <- c(0.1,1,10,50,90,99,99.9)
    unQ <- qnorm(pr/100)
    axis(3,unQ,paste0(pr,"%"),las=2)
    mtext("Probability Level", 3, line=3)
  }
  
  ## Conf Intervals
  lines(ConfInt$z,ConfInt$upper-ConfInt$line,col=rgb(1,0,0,alpha = 0.3))
  lines(ConfInt$z,ConfInt$lower-ConfInt$line,col=rgb(1,0,0,alpha = 0.3))
  
  
  
}

################
## Sampling ####
################

draw_samples <- function(issue,data,names,n_samp = 5,use_kfold=T,cov_matrix=NULL,
                         cov_name=NULL, # name of list element containing list of cov matrices corresponding to kfolds
                         dim_names,
                         tail_from=NULL,
                         tail_data=NULL # List containing tail params/data for each "names"
                         ){
  
  # require(mvtnorm) <<< depriciated, see commented out lines using "rmvnorm()"
  
  require(mvnfast)
  
  if(!is.null(cov_name)){data[["Cov"]] <- data[[cov_name]]}
  
  if(!is.null(cov_matrix) & use_kfold){stop("use_kfold and cov_matrix!=NULL not supported.")}
  
  if(!use_kfold){
    
    if(is.null(cov_matrix)){cov_matrix <- data[["Cov"]][[1]]}
    # samp_u <- as.data.table(pnorm(rmvnorm(n=n_samp,mean = data[["Means"]][[1]],sigma = cov_matrix)))
    samp_u <- as.data.table(pnorm(rmvn(n=n_samp,mu = data[["Means"]][[1]],sigma = cov_matrix)))
  }else{
    kfold <- unique(data[[N]]$kfold[data[[N]]$issueTime==issue])
    if(length(kfold)!=1){stop("kfold not unique.")}
    # samp_u <- as.data.table(pnorm(rmvnorm(n=n_samp,mean = data[["Means"]][[kfold]],sigma = data[["Cov"]][[kfold]])))
    samp_u <- as.data.table(pnorm(rmvn(n=n_samp,mu = data[["Means"]][[kfold]],sigma = data[["Cov"]][[kfold]])))
    
  }
  
  if(is.null(tail_from)){
    tail_from_index <- 1:ncol(data[[paste0(names[1],"_qr")]])
  }else{
    tail_from_index <- which(as.numeric(gsub("q","",names(data[[paste0(names[1],"_qr")]])))>=tail_from & 
                               as.numeric(gsub("q","",names(data[[paste0(names[1],"_qr")]])))<=100-tail_from)
  }
  
  colnames(samp_u) <- dim_names
  
  Scenarios <- list()
  for(N in names){
    Scenarios[[N]] <- data.frame(targetTime=data[[N]][issueTime==issue,unique(targetTime)])
    temp <- matrix(NA,nrow=nrow(Scenarios[[N]]),ncol=n_samp)
    for(i in 1:nrow(Scenarios[[N]])){
      index <- which(data[[N]]$issueTime==issue & data[[N]]$targetTime==Scenarios[[N]]$targetTime[i])
      temp[i,] <- PIT(data[[paste0(N,"_qr")]][index,tail_from_index],
                      t(samp_u[,grep(paste0(N,"[0-9]{2}"),colnames(samp_u))[i],with=F]),
                      inverse = T,
                      tails=if(is.null(tail_data)){
                        data[[paste0("tail_params_",N)]]
                      }else{
                        list(method="gpd",
                             scale_r=tail_data[[N]]$scale_r[i],shape_r=tail_data[[N]]$shape_r[i],
                             scale_l=tail_data[[N]]$scale_l[i],shape_l=tail_data[[N]]$shape_l[i])
                      })
      
    }
    Scenarios[[N]] <- cbind(Scenarios[[N]],temp)
  }
  
  return(Scenarios)
  
}

get_obs <- function(dt,issue,names,n=T){
  
  out <- NULL
  for(N in names){
    if(n){
      out <- c(out,dt[[N]][issueTime==issue,node_n])
    }else{
      out <- c(out,dt[[N]][issueTime==issue,node])
    }
  }
  
  return(out)
}

###############
## GAM+MQR #### Now in probcast
###############
#' Multiple Quantile Regression Using Generalised Additive Models and Linear Quantile Regression
#'
#' This function fits multiple conditional linear quantile regression models to the residuals of
#' a generalised additive model using \code{mgcv} and with facilities for cross-validation.
#' 
#' @author Jethro Browell, \email{jethro.browell@@strath.ac.uk}
#' @param data A \code{data.frame} containing target and explanatory
#' variables. May optionally contain a collumn called "kfold" with
#' numbered/labeled folds and "Test" for test data.
#' @param formala A \code{formula} object with the response on the left
#' of an ~ operator, and the terms, separated by + operators, on the right passed to \code{gam()}
#' or \code{bam()} from \code{mgcv}.
#' @param formula_qr Formula for linear quantile regression model for GAM residuals. Term \code{gam_pred}
#' is the prediction from the above GAM may be included in this formula.
#' @param model_res2 If \code{TRUE} also model squared residuals of GAM using a GAM. Defaults to \code{FALSE}.
#' @param formula_res2 Formula for GAM to predict squared residuals.
#' @param quantiles The quantiles to fit models for.
#' @param gbm_params List of parameters to be passed to \code{fit.gbm()}.
#' @param CVfolds Control for cross-validation if not supplied in \code{data}.
#' @param Sort \code{boolean} Sort quantiles using \code{SortQuantiles()}?
#' @param SortLimits \code{Limits} argument to be passed to \code{SortQuantiles()}. Constrains quantiles to upper and lower limits given by \code{list(U=upperlim,L=lowerlim)}.
#' @details Returns a \code{list} comprising predictive quantiles, GAM models, and deterministic predictions
#' from GAMs.
#' 
#' The returned predictive quantiles and GAM predictions are those produced out-of-sample for each
#' cross-validation fold (using models trained on the remaining folds but not "Test" data).
#' Predictive quantiles corresponding to "Test" data are produced using models trained on all
#' non-test data.
#' @return Quantile forecasts in a \code{MultiQR} object.
#' @keywords Quantile Regression
#' @export
# MQR_qreg_gam <- function(data,
#                          formula,
#                          formula_qr=NULL,
#                          model_res2=F,
#                          formula_res2 = formula,
#                          quantiles=c(0.25,0.5,0.75),
#                          CVfolds=NULL,
#                          ...,
#                          use_bam=T,
#                          w=rep(1,nrow(data)),
#                          Sort=T,SortLimits=NULL){
#   
#   ### Set-up Cross-validation
#   TEST<-F # Flag for Training (with CV) AND Test output
#   if("kfold" %in% colnames(data)){
#     if(!is.null(CVfolds)){warning("Using column \"kfold\" from data. Argument \"CVfolds\" is not used.")}
#     
#     if("Test" %in% data$kfold){
#       TEST<-T
#       nkfold <- length(unique(data$kfold))-1
#     }else{
#       nkfold <- length(unique(data$kfold))
#     }
#   }else if(is.null(CVfolds)){
#     data$kfold <- rep(1,nrow(data))
#     nkfold <- 1
#   }else{
#     data$kfold <- sort(rep(1:CVfolds,length.out=nrow(data)))
#     nkfold <- CVfolds
#   }
#   
#   if(!"BadData" %in% colnames(data)){
#     data[,BadData:=F]
#   }
#   
#   
#   ## GAM for conditional expectation (and possible squared residuals)
#   data[,gam_pred:=as.numeric(NA)]
#   gams <- list()
#   formula_res2 <- reformulate(attr(terms(formula_res2),"term.labels"),response = "gam_res2")
#   for(fold in unique(data$kfold)){
#     ## Fit gam
#     print(paste0("GAM, kfold=",fold))
#     
#     if(use_bam){
#       gams[[fold]] <- bam(data=data[kfold!=fold & kfold!="Test" & BadData==F,],
#                           formula = formula,
#                           weights = w,...)
#       if(model_res2){
#         gams[[paste0(fold,"_r")]] <- bam(data=cbind(data[kfold!=fold & kfold!="Test" & BadData==F,],data.table(gam_res2=residuals(gams[[fold]])^2)),
#                                          formula = formula_res2,
#                                          weights = w,...)
#       }
#     }else{
#       gams[[fold]] <- gam(data=data[kfold!=fold & kfold!="Test" & BadData==F,],
#                           formula = formula,
#                           weights = w,...)
#       if(model_res2){
#         gams[[paste0(fold,"_r")]] <- gam(data=cbind(data[kfold!=fold & kfold!="Test" & BadData==F,],data.table(gam_res2=residuals(gams[[fold]])^2)),
#                                          formula = formula_res2,
#                                          weights = w,...)
#       }
#     }
#     
#     ## Predictions
#     data$gam_pred[data$kfold==fold] <- predict(gams[[fold]],newdata = data[kfold==fold,])  
#     
#   }
#   
#   data[,gam_resid:=get(as.character(formula[[2]]))-gam_pred]
#   
#   ## Container for predictive quantiles
#   predqs <- data.frame(matrix(NA,ncol = length(quantiles), nrow = nrow(data)))
#   colnames(predqs) <- paste0("q",100*quantiles)
#   
#   
#   ## Qunatile regression
#   if(!is.null(formula_qr)){
#     ## Use user-specified model equation for quantile regression
#     formula_qr <- reformulate(attr(terms(formula_qr),"term.labels"),response = "gam_resid")
#     for(fold in unique(data$kfold)){
#       print(paste0("MQR, kfold=",fold))
#       for(i in 1:length(quantiles)){
#         
#         lqr_fit <- rq(formula_qr,
#                       tau = quantiles[i],
#                       data = data[kfold!=fold & kfold!="Test" & BadData==F,],
#                       weights = data[kfold!=fold & kfold!="Test" & BadData==F,w],
#                       method = "br")
#         
#         predqs[data$kfold==fold,i] <- data[kfold==fold,gam_pred] + predict.rq(lqr_fit,data[kfold==fold,])
#       }
#     }
#   }else{
#     ## QR with features from GAM
#     for(fold in unique(data$kfold)){
#       print(paste0("MQR, kfold=",fold))
#       ## Get training Data
#       train <- predict(gams[[fold]],newdata = data[kfold!=fold & kfold!="Test" & BadData==F,],type = "terms")
#       if(model_res2){
#         train2 <- predict(gams[[paste0(fold,"_r")]],newdata = data[kfold!=fold & kfold!="Test" & BadData==F,],type = "terms")
#         # Only need to retrun smooth terms as linear terms included already...
#         train2 <- train2[,grep("\\(",colnames(train2)),drop=F]
#         colnames(train2) <- paste0(colnames(train2),"_r")
#         train <- cbind(train,train2); rm(train2)
#       }
#       train <- cbind(as.data.table(train),data[kfold!=fold & kfold!="Test" & BadData==F,.(gam_resid)])
#       
#       ## Out-of-sample data
#       test_cv <- as.data.table(predict(gams[[fold]],newdata = data[kfold==fold,],type = "terms"))
#       if(model_res2){
#         test_cv2 <- predict(gams[[paste0(fold,"_r")]],newdata = data[kfold==fold,],type = "terms")
#         test_cv2 <- test_cv2[,grep("\\(",colnames(test_cv2)),drop=F]
#         colnames(test_cv2) <- paste0(colnames(test_cv2),"_r")
#         test_cv <- cbind(test_cv,test_cv2); rm(test_cv2)
#       }
#       
#       for(i in 1:length(quantiles)){
#         ## Fit QR model
#         lqr_fit <- rq(formula = gam_resid ~ .,
#                       tau = quantiles[i],
#                       data = train,
#                       weights = data[kfold!=fold & kfold!="Test" & BadData==F,w],
#                       method = "br")
#         
#         
#         
#         ## Make predictions
#         predqs[data$kfold==fold,i] <- data[kfold==fold,gam_pred] + predict.rq(lqr_fit,test_cv)
#         
#       }
#     }
#   }
#   
#   class(predqs) <- c("MultiQR","data.frame")
#   
#   
#   if(Sort){
#     predqs <- SortQuantiles(data = predqs,Limits = SortLimits)
#   }
#   
#   return(list(predqs=predqs,
#               gams=gams,
#               gam_pred=data$gam_pred))
#   
# }

