# This function extrapolates vertically a (x,y,z) field extending
# linearly the vertical slope of the profile.
#
# Usage   : python vertextrap.py <input file> <input variable name> 
#    <input grid description file> <input depth in grid file>
#    <output file>
#
# History : Virginie Guemas - Initial version     -  March 2014
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
import cdms2 as cdms
import sys,string
from cdms2 import MV
import numpy as N
from numpy import ma
#
# 1. Input arguments
# ===================
# 
# Input file and var names :
# ---------------------------
fileIN=sys.argv[1]
varIN=sys.argv[2]
#
# Meshmask file : 
# ----------------
fileM=sys.argv[3]
varD=sys.argv[4]
#
# Output file name :
# -------------------
fileOUT=sys.argv[5]
#
# 2. Get the input files 
# =======================
#
f=cdms.open(fileIN)
var0_a=f(varIN,squeeze=1)
f.close()
mask3d=var0_a.mask
(lz0,ly0,lx0)=var0_a.shape
#	    
f=cdms.open(fileM)
dep=f(varD,squeeze=1)
f.close()
#
varout=N.zeros((lz0,ly0,lx0))
varout=varout.astype('d')
varout=N.where(mask3d==False,var0_a,varout)
maskout=N.zeros((lz0,ly0,lx0))

for jz in N.arange(lz0) :
  if mask3d[jz,:,:].mean() == 1 :
    ratio=(dep[jz]-dep[jz-1])/(dep[jz-1]-dep[jz-2])
    varout[jz,:,:]=varout[jz-1,:,:]+(varout[jz-1,:,:]-varout[jz-2,:,:])*ratio
    maskout[jz,:,:]=maskout[jz-1,:,:]
  else :
    maskout[jz,:,:]=mask3d[jz,:,:]

varout=cdms.createVariable(varout,id=var0_a.id)
varout=MV.where(maskout<0.5,varout,1e20)
varout.getAxis(0)[:]=var0_a.getAxis(0)[:]
varout.getAxis(0).id='z'
varout.getAxis(1).id='y'
varout.getAxis(2).id='x'
varout.id=var0_a.id

h=cdms.open(fileOUT,'w')
h.write(varout)
h.close()