#' ADAMONT Analog finds analogous days 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 wt corresponding weather types, dimensions c(dataset=1,member,sdate,ftime)
#'@param obs reference data, dimensions c(dataset=1,member_obs=1,sdate,ftime,nlat_o,nlon_o)
#'@param wt_obs corresponding weather types, dimensions c(dataset=1,member_obs=1,sdate,ftime)
#'@param var variable
#'@param method string indicating the method used for analog definition
#'       This needs to be detailed more... Options rain1, rain01, patt_corr
#'@param thres real number indicating the precipitation threshold to define a rainy day
#'@param exp_pr, needed if var is not precip
#'@param obs_pr, needed if var is not precip
#'
#'@import ncdf4, abind


Search_Analog_rain1 <- function(exp_day,obs,iwt,thres) {
        # exp_day: exp_day[1,imb,sdate,ftime,,]
        # obs: full obs array (all sdates and ftimes)
	# iwt: indices corresponding to the weather type of exp_day
	# thres: threshold for precipitation (rainy day or not)
	analog5 <- array(dim=c(5))
	accuracy <- array(dim=length(iwt))
	for (i in iwt){
		accuracy[which(iwt==i)]<-sum((obs[i,,]>=thres)*(exp_day>=thres))
	}
	best_accuracy <- sort(accuracy,decreasing=TRUE)[1:5]
	vals_accuracy <- unique(best_accuracy)
	nunique <- length(vals_accuracy)
	count <- array(dim=nunique+1)
	count[1] <- 0
	for (ss in 1:nunique){
          count[ss+1] <- count[ss]+sum(best_accuracy==vals_accuracy[ss])
	  analog5[(count[ss]+1):count[ss+1]] <- sample(iwt[which(accuracy == vals_accuracy[ss])],sum(best_accuracy == vals_accuracy[ss]))
	}
	return(analog5)
}

Search_Analog_rain01 <- function(exp_day,obs,iwt,thres) {
        # exp_day: exp_day[1,imb,sdate,ftime,,]
        # obs: full obs array (all sdates and ftimes)
	# iwt: indices corresponding to the weather type of exp_day
	# thres: threshold for precipitation (rainy day or not)
	# Output: analog5 -> 5 indices corresponding to analog days of exp_day
	analog5 <- array(dim=c(5))
	accuracy <- array(dim=length(iwt))
	for (i in iwt){
		accuracy[which(iwt==i)]<-sum(diag(table((obs[i,,]>=thres),(exp_day>=thres))))
	}
	best_accuracy <- sort(accuracy,decreasing=TRUE)[1:5]
	vals_accuracy <- unique(best_accuracy)
	nunique <- length(vals_accuracy)
	count <- array(dim=nunique+1)
	count[1] <- 0
	for (ss in 1:nunique){
          count[ss+1] <- count[ss]+sum(best_accuracy==vals_accuracy[ss])
	  analog5[(count[ss]+1):count[ss+1]] <- sample(iwt[which(accuracy == vals_accuracy[ss])],sum(best_accuracy == vals_accuracy[ss]))
	}
	return(analog5)
}

Search_Analog_pattcorr <- function(exp_day,obs,iwt,thres) {
        # exp_day: exp_day[1,imb,sdate,ftime,,]
        # obs: full obs array (all sdates and ftimes)
	# iwt: indices corresponding to the weather type of exp_day
	# thres: threshold for precipitation (rainy day or not)
	# Output: analog5 -> 5 indices corresponding to analog days of exp_day
	analog5 <- array(dim=c(5))
	accuracy <- array(dim=length(iwt))
	for (i in iwt){
		accuracy[which(iwt==i)]<-cor(as.vector(obs[i,,]),as.vector(exp_day))
	}
	best_accuracy <- sort(accuracy,decreasing=TRUE)[1:5]
	vals_accuracy <- unique(best_accuracy)
	nunique <- length(vals_accuracy)
	count <- array(dim=nunique+1)
	count[1] <- 0
	for (ss in 1:nunique){
          count[ss+1] <- count[ss]+sum(best_accuracy==vals_accuracy[ss])
	  analog5[(count[ss]+1):count[ss+1]] <- sample(iwt[which(accuracy == vals_accuracy[ss])],sum(best_accuracy == vals_accuracy[ss]))
	}
	return(analog5)
}


Analog <- function(exp_corr, wt, obs, wt_obs, var, method='rain1', thres=1., exp_pr=NULL, obs_pr=NULL) {

	
	# WORK IN PROGRESS...	
	
	# CHECKS NEEDED BEFORE RUNNING
	# Do the dimensions correspond between obs_pr and obs? exp_pr and exp?
	# INCLUDING CHECKS ON DIMENSIONS
	# CHECK ON VALUE OF PRECIPITATION
	# ONE OPTION IS TO GIVE THE NAME OF PRECIP VARIABLE IN THE WRAPPER FUNCTION

        search_analog <- switch(method, 'rain1'=Search_Analog_rain1, 'rain01'=Search_Analog_rain01, 'pattcorr'=Search_Analog_pattcorr, stop("Adamont Analog function supports methods 'rain1', 'rain01', 'pattcorr' only"))

        analog <- array(dim=c(1,member=dim(exp_corr)[2],sdate=dim(exp_corr)[3],ftime=dim(exp_corr)[4],ndays=5))

	for (imb in 1:dim(exp_corr)[2]){
          for (isdate in 1:dim(exp_corr)[3]){
            for (itime in 1:dim(exp_corr)[4]){
              iwt <- wt[1,imb,isdate,itime]
              index <- which(wt_obs == iwt)
              # No rain day case : pick five random days among corresponding weather type
              if (sum(exp_pr[1,imb,isdate,itime,,]>=thres) == 0) {
                analog[1,imb,isdate,itime,] <- obs[sample(index,5)]
              } else {
                # Find five closest analogs among corresponding weather type
                if (length(index)!=0){
                  obs_pr2 <- obs_pr
                  dim(obs_pr2) <- c(1*1*dim(obs_pr)[3]*dim(obs_pr[4]),dim(obs_pr[5]),dim(obs_pr[6]))
                  analog[1,imb,isdate,itime,] <- search_analog(exp_pr[1,imb,isdate,itime,,],obs_pr2,index,thres)
                }
              }
	    }
	  }
	}

	return(analog)
}

