#!/bin/ksh 

module load cdo
set -exuv 

##======================================
##  Script to calculate the rescaling of ERAI precipitation using GPCP 
## Emanuel, November 2013 


##======================================
## From command line 
YYYYMM=$1  ## date to calculate the rescaling 
TPMIN=0.05 ## minimum precipitation accepted in mm/day 
RESmin=0.1 # minimum rescaling allowed
RESmax=10  # maximum rescaling allowed 

## USAGE: rescale_erai_gpcp_mm YYYYMM 

##======================================
## File locations
FGPCP="/vol/seasonal_data/odb/rain/gpcp_2.2/"  #GPCPsgobs201210
# FGPCP="/scratch/mo/moi/GPCP_2.2/"

ECFSDIR="need/tmp" #"/need/ERAI_SFC_201307/"
## Location of ERAI forcing files as saves in 201307 and onwards from experiment fxml (user: need)

ECFSDIRSTOR=$ECFSDIR/ #GPCP_resc/  
## To store rescaled precipitation and coef.

TMPSTOR=$SCRATCH/tmp/EI_GPCP2_RESC_NEW/
## To store monthly means in SCRATCH for diagnostics 

WDIR=$TMPDIR/resc_erai_${YYYYMM}_NEW
## temporary working directory 

##======================================
## Function definitions 

## calculate daily and monthly mean of total precipitation 
calc_dd_mm_tp_mean() {
frainf=$1  #rainfall 3h file 
fsnowf=$2  # snowfall 3h file 
foutDD=$3  # total precipitation daily means
foutMM=$4  # total precipitaiton monthly means

## remove 1th field of grib files 
grib_copy -w count!=1 $frainf tmp_rainf.grb
grib_copy -w count!=1 $fsnowf tmp_snowf.grb

## add rainf+snowf
cat > bar << EOF
read,source=tmp_rainf.grb,fieldset=tp
read,source=tmp_snowf.grb,fieldset=sf
compute,formula="tp+sf",target="tmp_tp.grb"
EOF
mars < bar 

## generate mars script for the daily and monthly means 
rm -f bar
touch bar
# for st in 03 06 09 12 15 18 21 24
# do
# cat >> bar << EOF
# read, source="tmp_tp.grb",time=00,step=${st}/to/${LSTEP}/by/24,fieldset=rf_${st}
# EOF
# done
cat >> bar << EOF
read, source="tmp_tp.grb",time=00,step=12,fieldset=rf_12
read, source="tmp_tp.grb",time=00,step=15,fieldset=rf_15
read, source="tmp_tp.grb",time=00,step=18,fieldset=rf_18
read, source="tmp_tp.grb",time=00,step=21,fieldset=rf_21
read, source="tmp_tp.grb",time=12,step=12,fieldset=rf_24
read, source="tmp_tp.grb",time=12,step=15,fieldset=rf_03
read, source="tmp_tp.grb",time=12,step=18,fieldset=rf_06
read, source="tmp_tp.grb",time=12,step=21,fieldset=rf_09
EOF

cat >> bar << EOF
compute,formula="max(0,3600*3*((rf_24+rf_21+rf_18+rf_15+rf_12+rf_09+rf_06+rf_03)))",fieldset=tp_dd
compute,formula="mean(tp_dd)",fieldset=tp_mm

write,fieldset=tp_dd,target="$foutDD"
write,fieldset=tp_mm,target="$foutMM"
EOF
mars < bar
rm -f tmp_rainf.grb tmp_snowf.grb tmp_tp.grb  rules.txt tmp_tp.grb bar
grib_set -s indicatorOfParameter=228 $foutDD tmp.grb ; mv tmp.grb $foutDD
grib_set -s indicatorOfParameter=228 $foutMM tmp.grb ; mv tmp.grb $foutMM
}

## apply 3h rescaling 
apply_resc() {
frainf_IN=$1
fsnowf_IN=$2
fresc=$3
frainf_OUT=$4
fsnowf_OUT=$5

nfields=`grib_ls $frainf | tail -1 | awk '{print $1}'`

frescF=tmp_rest.grb 
rm -f $frescF
touch $frescF
ikf=1
while [ $ikf -le $nfields ];
do
    cat $fresc >> $frescF
    ikf=$(( $ikf + 1 ))
done

cat > bar << EOF
read,source="$frainf_IN",fieldset=rf
read,source="$fsnowf_IN",fieldset=sf
read,source="$frescF",fieldset=resc
compute,formula="(rf*resc)",target="$frainf_OUT"
compute,formula="(sf*resc)",target="$fsnowf_OUT"
EOF
mars < bar
rm -f bar $frescF
}


## extract ERAI-forcing from experiment 
extract_ei_data() {
## extract data
# expver="g2us"
# if [[ ${YYYYMM} -ge 201401 ]] ; then
#   expver=g93o
# fi
# mars << EOF
#   retrieve,
#   stream=oper,levtype=sfc,type=fc,
#   time=0,date=${YYYYMM}01,
#   step=${STEP},
#   expver=$expver,class=rd,
#   param=142.128/240.128,
#   target="tmp.grb"
# EOF
# grib_copy -w paramId=142 tmp.grb Rainf_ei${YYYYMM}.grb
# grib_copy -w paramId=240 tmp.grb Snowf_ei${YYYYMM}.grb
# rm -f tmp.grb 

mkdir -p tmp
cd tmp 
ecp ec:/ptb/OSM_FORCING/forcing_ei_1_oper_3_${YYYYMM}00.tar.gz forcing_${YYYYMM}.tar.gz
# ecp ec:/need/E2O/ERAI3H/ei_1_${YYYYMM}00.tar.gz forcing_${YYYYMM}.tar.gz
tar -zxvf forcing_${YYYYMM}.tar.gz
rm -f forcing_${YYYYMM}.tar.gz
mv Rainf.grb ../Rainf_ei${YYYYMM}.grb 
mv Snowf.grb ../Snowf_ei${YYYYMM}.grb
cd ..
rm -rfv tmp 

}

##======================================
## MAin Script 

CWD=`pwd`
emkdir -p $ECFSDIRSTOR
mkdir -p $WDIR/
cd $WDIR 
rm -rfv *

## get ERAI grib files 
LASTDATE=$(newdate -D ${YYYYMM}01 +33 | cut -c1-6)01
LSTEP=$(/usr/local/bin/datediff -h ${LASTDATE}00 ${YYYYMM}0100 )
STEP=0/to/${LSTEP}/by/3
extract_ei_data


## get gaussian reduced cdo grid description 
/home/rd/need/work/efas/gen_gg_grid_descr.py -i Rainf_ei${YYYYMM}.grb -o gdes.txt 

## Calc ERAI original daily and monthly means of total precipitation - original data in kg m-2 s-1 
calc_dd_mm_tp_mean Rainf_ei${YYYYMM}.grb Snowf_ei${YYYYMM}.grb TP_ei_DD_${YYYYMM}.grb TP_ei_MM_${YYYYMM}.grb


## Get GPCP data 
# cdo -f nc copy $FGPCP/GPCPsgobs${YYYYMM} GPCP.nc

# cdo -f nc seldate,${YYYYMM}01 /hugetmp/need/data/precip/precip.mon.mean.nc GPCP.nc
cdo -f nc seldate,${YYYYMM}01 /perm/rd/need/work/precip_cmap/precip.mon.mean.nc GPCP.nc
ncrename -v precip,var228 GPCP.nc 


## Interpolate the ERAI monthly mean to the GPCP grid 
cdo -R -t ecmwf -f nc -remapcon,GPCP.nc TP_ei_MM_${YYYYMM}.grb TP_ERAI_UP.nc

## Calculate rescalling factor:
cdo -setrtoc,0,$RESmin,$RESmin -setrtoc,$RESmax,100000000,$RESmax -setmisstoc,1 -div GPCP.nc -setrtomiss,-10000,$TPMIN TP_ERAI_UP.nc RESC_${YYYYMM}.nc

## interpolate rescalling factor to gaussian grid 
cdo -t ecmwf -f grb -setgrid,Rainf_ei${YYYYMM}.grb -remapbil,gdes.txt RESC_${YYYYMM}.nc RESC_GG_${YYYYMM}.grb 
grib_set -s table2Version=128,indicatorOfParameter=142 RESC_GG_${YYYYMM}.grb tmp.grb ; mv tmp.grb RESC_GG_${YYYYMM}.grb 

## Apply rescaling 
apply_resc Rainf_ei${YYYYMM}.grb Snowf_ei${YYYYMM}.grb RESC_GG_${YYYYMM}.grb  Rainf_eigpcp2${YYYYMM}.grb Snowf_eigpcp2${YYYYMM}.grb

# Calc New ERAI daily and monthly means of total precipitation
calc_dd_mm_tp_mean Rainf_eigpcp2${YYYYMM}.grb Snowf_eigpcp2${YYYYMM}.grb TP_eigpcp2_DD_${YYYYMM}.grb TP_eigpcp2_MM_${YYYYMM}.grb

for ff in TP_ei_DD_ TP_ei_MM_  RESC_GG_ Rainf_eigpcp2 Snowf_eigpcp2 TP_eigpcp2_DD_ TP_eigpcp2_MM_
do 
    ecp -o ${ff}${YYYYMM}.grb ec:/$ECFSDIRSTOR/${ff}${YYYYMM}.grb
done
ecp -o RESC_${YYYYMM}.nc ec:/$ECFSDIRSTOR/RESC_${YYYYMM}.nc
ecp -n gdes.txt ec:/$ECFSDIRSTOR/

mkdir -p $TMPSTOR
cp TP_ei_MM_${YYYYMM}.grb $TMPSTOR/
cp TP_eigpcp2_MM_${YYYYMM}.grb $TMPSTOR/
cp RESC_${YYYYMM}.nc $TMPSTOR/

cd $CWD
rm -rf $WDIR 
exit 0 


##=====================================
## for checking only:
# cdo -R -t ecmwf -f nc -remapcon,$FGPCP -mulc,86400  -add -timmean -mul Rainf_ei${YYYYMM}.grb RESC_GG_${YYYYMM}.grb  -timmean -mul Snowf_ei${YYYYMM}.grb RESC_GG_${YYYYMM}.grb  TP_ERAI_UP_NEW.nc


