library(doParallel)            
library(doSNOW)  
library(randomForest)
library(qmap)
library(RColorBrewer)
library(ggplot2)
library(VIC5)
library(hydroGOF)
library(abind)

calculate_metrics <- function(qsim_tr, qobs_tr, qsim_ts, qobs_ts, prefix) {
  
  metrics <- list()
  
  metrics[[paste0("mae.tr.", prefix, ".pp")]] = mae(qsim_tr, qobs_tr)
  metrics[[paste0("mae.ts.", prefix, ".pp")]] = mae(qsim_ts, qobs_ts)
  
  metrics[[paste0("nse.tr.", prefix, ".pp")]] = NSE(qsim_tr, qobs_tr)
  metrics[[paste0("nse.ts.", prefix, ".pp")]] = NSE(qsim_ts, qobs_ts)
  
  metrics[[paste0("kge.tr.", prefix, ".pp")]] = KGE(qsim_tr, qobs_tr, method= "2009", out.type="full")
  metrics[[paste0("kge.ts.", prefix, ".pp")]] = KGE(qsim_ts, qobs_ts, method= "2009", out.type="full")
  
  metrics[[paste0("lns.tr.", prefix, ".pp")]] = logNSE(qsim_tr, qobs_tr)
  metrics[[paste0("lns.ts.", prefix, ".pp")]] = logNSE(qsim_ts, qobs_ts)
  
  return(metrics)
}




ptm <- proc.time()
ncores =31 
cl <- makePSOCKcluster(ncores,outfile='') # Register cluster
checkCluster(cl)
clusterSetRNGStream(cl, c(1:ncores))
registerDoSNOW(cl)  
step <- ceiling(Station.N/ncores)
smpopts <- c('hydroGOF','qmap','randomForest','neuralnet','VIC5','abind')
clusterCall(cl, ".libPaths", usedlib)

Outputs <- list()
Outputs <- foreach(ix_sub=seq(1, Station.N, step), .combine = 'c', .packages = smpopts) %dopar% {
  Outputs_ = list()
for(i in ix_sub:min(ix_sub+step-1,Station.N)){

  setwd(fold.out)
  if(!file.exists(paste("pp_station_",i,".RDATA",sep=""))){
  print(i)
  qobs=obsCOUT[,i]
  qsim=simCOUT[,i]


  na.idx=which(is.na(qobs))
  if(length(na.idx)!=0){
    qobs=qobs[-na.idx]
    qsim = asub(qsim, -na.idx, dims = 1)
  }
  
  if(length(qobs)!=0){

    
    ## spit 80% 20%
    tr.idx=seq(1,ceiling(length(qobs)*0.8))
    
    qobs.tr <- qobs[tr.idx]
    qobs.ts <- qobs[-tr.idx]
    qsim.tr <- asub(qsim, tr.idx, dims = 1)  
    qsim.ts <- asub(qsim, -tr.idx, dims = 1)  
    
   ##rescale 
    scale.=max(qsim.tr)-min(qsim.tr)
    center.=min(qsim.tr)
    scale.pp=scale.
    center.pp=center.
    
    ## raw -----------------------------------------------------------------
    prefix <- "rw"
    rw.metrics.pp = calculate_metrics(qsim.tr, qobs.tr, qsim.ts, qobs.ts, prefix)
    
    
    ## QM-------------------------------------------------------------------
    qm.fit <- fitQmapQUANT(qobs.tr,qsim.tr,qstep=0.01)
    model.qm.pp=qm.fit
    qm.tr <- doQmapQUANT(qsim.tr,qm.fit,type='tricub')
    qm.ts <- doQmapQUANT(qsim.ts,qm.fit,type='tricub')
    prefix <- "qm"
    qm.metrics.pp = calculate_metrics(qm.tr, qobs.tr, qm.ts, qobs.ts, prefix)

    
    # RF---------------------------------------------------------------------
    Q.tr=as.data.frame(cbind(qobs.tr,qsim.tr))
    colnames(Q.tr)=c("qobs","qsim")
    Q.ts=as.data.frame(cbind(qobs.ts,qsim.ts))
    colnames(Q.ts)=c("qobs","qsim")


    rf <- randomForest(qobs ~ ., data = Q.tr, 
                       importance = TRUE,  type='regression',keep.forest=TRUE)
    
    model.rf.pp=rf
    rf.tr=predict(rf,Q.tr)
    rf.ts=predict(rf,Q.ts)
    prefix <- "rf"
    rf.metrics.pp = calculate_metrics(rf.tr, qobs.tr, rf.ts, qobs.ts, prefix)
    Outputs_ <- list()
    Outputs_[["model.rf.pp"]]<-model.rf.pp
    #GLM--------------------------------------------------------------------

    glm.model=glm(qobs~.,data=Q.tr)
    model.gl.pp=glm.model
    glm.tr <- predict(object = glm.model, newdata = Q.tr, type = "response")
    glm.ts <- predict(object = glm.model, newdata = Q.ts, type = "response")
   
    prefix <- "gl"
    gl.metrics.pp = calculate_metrics(glm.tr, qobs.tr, glm.ts, qobs.ts, prefix)
    
    
    scores.list = ls(pattern = '.pp')
    for(sco in scores.list){
      
      out = eval(parse(text=sco)) 
      Outputs_[[sco]] = out
    }
    setwd(fold.out)
    save(Outputs_,file=paste("pp_station_",i,".RDATA",sep=""))
  }
}
  }
 
  cat(' ... RETURN ')
}



cat(' ... STOP CLUSTER ')
stopCluster(cl)