#!/usr/bin/env Rscript 

#
# This script filters out from any 1d (time) to 3d (time-lat-lon) field the 
# effect of greenhouse gases, ENSO, or any mode of variability through a 
# multilinear regression over the corresponding index (greenhouse gas 
# concentration, ENSO  index ..) at a set of lags you can choose. 
#
# Usage : ./regressedts.R <ncfile> <ncvar> <indexfile> <indexvar> <outfile> 
#                         <lag1> <lag2> ... (as many lags as wanted)
#
# History : Virginie Guemas - Initial version                    -  2011
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

library(ncdf)

args=commandArgs(TRUE)

filein=args[1] # Input netcdf file 
varin=args[2] # Input netcdf variable name
index=args[3] # Netcdf file containing the index which to filter the signal
varindex=args[4] # Index variable name
fileout=args[5] # Output file
lags=c()        # List of lags
jind=6
while ( jind <= length(args) ) {
  lags=c(lags,as.numeric(args[jind]))
  jind=jind+1
}
if ( is.null(lags) ) { lags=c(0) } 

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)

fnc=open.ncdf(index)
vind=get.var.ncdf(fnc,varindex)
close.ncdf(fnc)

system(paste('cdo showmon ',index,'> mons 2>/dev/null'))
mons=read.table('mons')
system(paste('cdo showyear ',index,'> 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=list()
for (jlag in lags) { lstind=list(lstind,lag(index_ts,-jlag)) }

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] ~ '
      jind=2
      for (jlag in lags) { 
        toto = ts.intersect(toto,lag(index_ts,-jlag), dframe=FALSE) 
        f=paste(f,'dataf[,',as.character(jind),'] +',sep="")
        jind=jind+1
      }
      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)
      #lm.out=lm(varjpt ~ index_ts,na.action=na.omit)
      #tmp[jpt,]=varjpt-(lm.out$coefficients[1]+index_ts*lm.out$coefficients[2])
      tmp[jpt,]=toto[,1]-lm.out$fitted.values
    }
  }
}else{
  toto = var_ts
  f='dataf[,1] ~ '
  jind=2
  for (jlag in lags) {
    toto = ts.intersect(toto,lag(index_ts,-jlag), dframe=FALSE)
    f=paste(f,'dataf[,',as.character(jind),'] +',sep="")
    jind=jind+1
  }
  f=substr(f,start=1,stop=(nchar(f)-1))
  dataf=toto
  lm.out=lm(as.formula(f),data=dataf,na.action=NULL)
  tmp=toto[,1]-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()