#!/usr/bin/env Rscript 
#
# This script extract from any 1d (time) to 3d (time-lat-lon) field the 
# combined effect of greenhouse gases, ENSO, and any mode of variability 
# through a multilinear regression over the corresponding index 
# (greenhouse gas concentration, ENSO index ..)  
#
# Usage : ./regressedts.R <ncfile> <ncvar> <indexfile1> <indexvar1> 
#             <indexfile2> <indexvar2> (as many indices as wanted) <outfile> 
#
# History : Virginie Guemas - Initial version                    -  2011
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

library(ncdf)

args=commandArgs(TRUE)

filein=args[1] # Input netcdf file 
varin=args[2] # Input netcdf variable name
jind=3
indexes=c() # List of netcdf files containing the indices to regress on
varindex=c() # List of index variable names
while ( jind <= (length(args)-1) ) {
  indexes=c(indexes,args[jind])
  varindex=c(varindex,args[jind+1])
  jind=jind+2
}
fileout=args[jind] # Output file

fnc=open.ncdf(filein)
absent=T
for ( jvar in 1:fnc$nvars ) {
  if ( fnc$var[[jvar]]$name == varin ) {
    data=fnc$var[[jvar]]
    unit=data$units
    n0dims=data$ndims
    wdimsout=data$dim
    var=get.var.ncdf(fnc,varin)
    ndims=length(dim(var))
    miss=att.get.ncdf(fnc,varin,'_FillValue')
    if ( miss$hasatt == TRUE ) { var[var==miss$value]=NA }
    miss=att.get.ncdf(fnc,varin,'missing_value')
    if ( miss$hasatt == TRUE ) { var[var==miss$value]=NA }
    absent=F
  }
}
if (absent) { stop("Variable not in input file !") } 
close.ncdf(fnc)

system(paste('cdo showmon ',filein,'> mons 2>/dev/null'))
mons=read.table('mons')
system(paste('cdo showyear ',filein,'> years 2>/dev/null'))
years=read.table('years')
system('rm -f mons years')

if (is.null(ndims)) { ndims=1 }
dimsvar=switch(ndims,length(var),dim(var),c(dim(var)[1]*dim(var)[2],dim(var)[3]))
if ( ndims > 1 ) { tmpvar=t(array(var,dim=dimsvar)) } else { tmpvar=var }
var_ts=ts(tmpvar,start=c(years[[1]],mons[[1]]),end=c(years[[length(years)]],mons[[length(mons)]]),frequency=12)

lstind=list()
for (ind in 1:length(indexes) ) {
  fnc=open.ncdf(indexes[ind])
  vind=get.var.ncdf(fnc,varindex[ind])
  close.ncdf(fnc)

  system(paste('cdo showmon ',indexes[ind],'> mons 2>/dev/null'))
  mons=read.table('mons')
  system(paste('cdo showyear ',indexes[ind],'> years 2>/dev/null'))
  years=read.table('years')
  system('rm -f mons years')

  index_ts=ts(vind,start=c(years[[1]],mons[[1]]),end=c(years[[length(years)]],mons[[length(mons)]]),frequency=12)

  lstind[[ind]]=index_ts
}

def=TRUE
if ( ndims > 1 ) {
  for (jpt in 1:dimsvar[1]) {
    varjpt=var_ts[,jpt]
    if (length(sort(varjpt)) > 1 ) {
      toto = varjpt
      f='dataf[,1] ~ '
      for (ind in 1:length(indexes)) { 
        toto = ts.intersect(toto,lstind[[ind]], dframe=FALSE) 
        f=paste(f,'dataf[,',as.character(ind+1),'] +',sep="")
      }
      dataf=toto
      if (def==TRUE) {
        tmp=array(,dim=c(dimsvar[1],dim(toto)[1]))
        def=FALSE
      }
      f=substr(f,start=1,stop=(nchar(f)-1))
      dataf=toto
      lm.out=lm(as.formula(f),data=toto,na.action=NULL)
      tmp[jpt,]=lm.out$fitted.values
    }
  }
}else{
  toto = var_ts
  f='dataf[,1] ~ '
  for (ind in 1:length(indexes)) {
    toto = ts.intersect(toto,lstind[[ind]], dframe=FALSE)       
    f=paste(f,'dataf[,',as.character(ind),'] +',sep="")
  }
  f=substr(f,start=1,stop=(nchar(f)-1))
  dataf=toto
  lm.out=lm(as.formula(f),data=dataf,na.action=NULL)
  tmp=lm.out$fitted.values
}

tmp[is.na(tmp)]=1e20
wtime=dim.def.ncdf("time",paste("months since ",as.character(start(toto)[1]),"-",as.character(start(toto)[2]),"-15 12:00:00",sep=""),seq(0,dim(toto)[1]-1))
wdimsout[[n0dims]]=wtime
dimsout=switch(ndims,length(tmp),dim(tmp),c(dim(var)[1],dim(var)[2],dim(tmp)[2]))
wvarout=var.def.ncdf(varin,unit,wdimsout,1e20)
f=create.ncdf(fileout,wvarout)
put.var.ncdf(f,wvarout,array(tmp,dim=dimsout))
close.ncdf(f)

rm(list=ls())
quit()