#' ADAMONT TimeDisaggregate uses analogous days from ADAMONT Analog and higher frequency time series in the reference dataset to experiment data based on
#' precipitation patterns
#'
#'@author Paola Marson, \email{paola.marson@meteo.fr} for PROSNOW version
#'@author Lauriane Batté, \email{lauriane.batte@meteo.fr} for CSTools adaptation
#'
#'@param exp_corr output from QQ_Corr, dimensions c(dataset=1,member,sdate,ftime,nlat_o,nlon_o)
#'@param analog output from Analog, dimensions c(dataset=1,member,sdate,ftime,5) (5 analogs)
#'@param obs reference subdaily data, dimensions c(dataset=1,member_obs=1,sdate,ftime,subd,nlat_o,nlon_o)
#'       subd is number of subdaily time steps in obs data (24 for hourly, 4 for 6-hourly, etc.)
#'@param obs_daily reference daily data, dimensions c(dataset=1,member_obs=1,sdate,ftime,nlat_o,nlon_o)
#'@param var variable name
#'@param indx index of analog day used for precip correction, dimensions c(dataset=1,member,sdate,ftime)
#'       (if var=prec, NULL required, if not, should be provided)
#'@param thres precipitation threshold for definition of rainy day
#'
#'@import ncdf4, abind

.Disagg <- function(mod, ref, var){
   tres <- length(ref) # Length of subdaily data (24: hourly, 4: 6-hourly, etc.)
   D_mod <- array(dim=c(tres))
   if (mean(ref) > 10e-10){
      correction1 <- c("Qair", "SWdown", "prec", "LWdown")
      correction2 <- c("uas", "vas")
      ## CHANGE TO STANDARD VARIABLE NAMES ABOVE...
      if (var %in% correction1){
         a <- mod/mean(ref)
         b <- 0
      } else if (var %in% correction2) {
         a <- mod/mean(ref)
         if (a <= 1) {
            b <- 0
         } else {
            a <- 1
            b <- mod - mean(ref)
         }
      } else { # Standard case such as temperature
         a <- 1 
         b <- mod-mean(ref)
      }
   } else { # mean(ref) <= 10e-10
      correction1 <- c("LWdown")
      correction2 <- c("uas", "vas")
      if (var %in% correction2){
         a <- 0
         b <- rep(mod, length(ref))
      } else if (var %in% correction1) {
         a <- mod/mean(ref) 
         b <- 0
      } else {
         a <- 0 
         b <- 0
      }
   }
   D_mod <- a*ref + b
   return(D_mod)
}

.SelPrecAnalog <- function(mod, ref, analog5, thres) {
   ## Function to select which of the 5 possible analogs is chosen
   ## mod: array with corrected daily precip data at grid point
   ## ref: corresponding obs array with subdaily precip data at grid point for five analog days
   ## analog5: 5 analog days
   ## thres: threshold for precipitation
   day_indx <- 0
   repeat {
      day_indx <- day_indx+1
      ind <- analog5[day_indx]
      if (day_indx == 5 |  (sum(is.na(ref[day_indx,]))==0 & ((mod*86400 < thres ) | (mod*86400 >= thres & mean(ref[day_indx,])*86400 >= thres)))) { 
         break
      } 
   }	
   return(ind)
}


TimeDisaggregate <- function(exp_corr, analog, obs, obs_daily, var, indx, thres) {
   dim_exp <- dim(exp_corr)
   exp_dt <- array(dim=c(subd=dim(obs)[5],dim_exp[1:6]))
   ## obs dimensions: c(dataset=1,member_obs=1,sdate,ftime,subd,nlat_o,nlon_o)
   obs2 <- aperm(obs,c(5,1,2,3,4,6,7)) # Start with subdaily dimension
   ## Flatten obs2 array
   dim(obs2) <- c(dim(obs2)[1],prod(dim(obs2)[2:5]),dim(obs2)[6],dim(obs2)[7])
   if (var != "prec"){
      if (is.null(indx)){
         stop("TimeDisaggregate MUST first be run on precipitation to select appropriate analog index indx")
      } else {
         ## RUN Disagg function on data using analog day indx
         for (imb in 1:dim(exp_corr)[2]){
            for (sd in 1:dim(exp_corr)[3]){
               for (ltime in 1:dim(exp_corr)[4]){
                  for (ilat in 1:dim(exp_corr)[5]){
                     for (ilon in 1:dim(exp_corr)[6]){
                        exp_dt[,1,imb,sd,ltime,ilat,ilon] <- Disagg(exp_corr[1,imb,sd,ltime,ilat,ilon],obs2[,indx[1,imb,sd,ltime,ilat,ilon],ilat,ilon],var)
                     }
                  }
               }
            }
         }
      }
   } else {
      if (is.null(indx)){
      ## Select appropriate analog among 5 possible -> indx array
         indx <- array(dim=dim_exp) 
         ##for (imb in 1:dim(exp_corr)[2]){
         ##   indx[1,imb,,,,] <- mapply(SelPrecAnalog,exp_corr[1,imb,,,,],obs2,analog[1,imb,,,],thres)
         ##}
         for (imb in 1:dim(exp_corr)[2]){
            for (sd in 1:dim(exp_corr)[3]){
               for (ltime in 1:dim(exp_corr)[4]){
                  for (ilat in 1:dim(exp_corr)[5]){
                     for (ilon in 1:dim(exp_corr)[6]){
                        indx[1,imb,sd,ltime,ilat,ilon] <- SelPrecAnalog(exp_corr[1,imb,sd,ltime,ilat,ilon],obs2[,,ilat,ilon],analog[1,imb,sd,ltime,],thres)
                     }
                  }
               }
            }
         }
         ## RUN Disagg function on data
         for (imb in 1:dim(exp_corr)[2]){
            for (sd in 1:dim(exp_corr)[3]){
               for (ltime in 1:dim(exp_corr)[4]){
                  for (ilat in 1:dim(exp_corr)[5]){
                     for (ilon in 1:dim(exp_corr)[6]){
                        exp_dt[,1,imb,sd,ltime,ilat,ilon] <- Disagg(exp_corr[1,imb,sd,ltime,ilat,ilon],obs2[,indx[1,imb,sd,ltime,ilat,ilon],ilat,ilon],var)
                     }
                  }
               }
            }
         }
      } else {
         warning("Using previously defined indx array for TimeDisaggregate: precipitation disaggregation is run again")
         for (imb in 1:dim(exp_corr)[2]){
            for (sd in 1:dim(exp_corr)[3]){
               for (ltime in 1:dim(exp_corr)[4]){
                  for (ilat in 1:dim(exp_corr)[5]){
                     for (ilon in 1:dim(exp_corr)[6]){
                        exp_dt[,1,imb,sd,ltime,ilat,ilon] <- Disagg(exp_corr[1,imb,sd,ltime,ilat,ilon],obs2[,indx[1,imb,sd,ltime,ilat,ilon],ilat,ilon],var)
                     }
                  }
               }
            }
         }
      }
   }
   ##Rearrange dimension order in exp_dt before output
   exp_dt <- aperm(exp_dt,c(2,3,4,5,1,6,7)) # Put subdaily dimension back in place
   return(list(disagg=exp_dt, indx=indx))
}

