# This function interpolates vertically a (x,y,z) field using the
# conservative method.
#
# Usage   : python interp_vert.py <input file> <input variable name> 
#    <input grid description file> <input level thickness in grid 
#    file> <input mask file> <input mask variable> <output grid 
#    description file> <output level thickness in grid file>  
#    <output level depth in grid file> <output file>
#
# History : Virginie Guemas - Initial version adpated
#                             from PhD tools        -       2012
#                             Cleaning              - March 2014
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#import cdms,regrid,sys,string
import cdms2 as cdms
import sys,string
#import MLab,MV,MA,Numeric as N
from cdms2 import MV
import numpy as N
from numpy import ma
#import vcs,time
#x=vcs.init()
#
# 1. Input arguments
# ===================
# 
# Input file and var names :
# ---------------------------
fileIN=sys.argv[1]
varIN=sys.argv[2]
#
# Input grid file : 
# ------------------
fileIG=sys.argv[3]
varIT=sys.argv[4]
#
# Input mask file :
# ------------------
fileM=sys.argv[5]
varM=sys.argv[6]
#
# Output grid file : 
# -------------------
fileOG=sys.argv[7]
varOT=sys.argv[8]
varOD=sys.argv[9]
#
# Output file name :
# -------------------
fileOUT=sys.argv[10]
#
# 2. Get the input files 
# =======================
#
f=cdms.open(fileIN)
var0_a=f(varIN,squeeze=1)
f.close()

if var0_a.rank() != 3 :
    print "3d input file please !"
dims_var0=var0_a.getOrder()
if dims_var0[0] == 'x' : 
    lx0=var0_a.shape[0]
    if dims_var0[1] == 'y' : 
        ly0=var0_a.shape[1]
        lz0=var0_a.shape[2]
    else :
        ly0=var0_a.shape[2]
        lz0=var0_a.shape[1]
else :
    if dims_var0[0] == 'y' :
        ly0=var0_a.shape[0]
        if dims_var0[1] == 'x' :
            lx0=var0_a.shape[1]
            lz0=var0_a.shape[2]
        else : 
            lx0=var0_a.shape[2]
            lz0=var0_a.shape[1]	
    else :
        lz0=var0_a.shape[0]
        if dims_var0[1] == 'x' :
            lx0=var0_a.shape[1]
            ly0=var0_a.shape[2]
        else :
            lx0=var0_a.shape[2]
            ly0=var0_a.shape[1]
#	    
var0=N.zeros((lz0,ly0,lx0))
var0=var0.astype('d')
#
if dims_var0[0] == 'x' : 
    if dims_var0[1] == 'y' : 
        for jk in N.arange(lz0) :
            var0[jk,:,:]=N.transpose(var0_a[:,:,jk])	
    else :
        for ji in N.arange(lx0) :
            var0[:,:,ji]=var0_a[ji,:,:]
else :
    if dims_var0[0] == 'y' :
        if dims_var0[1] == 'x' :
            for jk in N.arange(lz0) :
                var0[jk,:,:]=var0_a[:,:,jk]	
        else : 
            for jk in N.arange(lz0) :
                var0[jk,:,:]=var0_a[:,jk,:]
    else :
        if dims_var0[1] == 'x' :
            for ji in N.arange(lx0) :
                var0[:,:,ji]=var0_a[:,ji,:]	    
        else :
            var0[:,:,:]=var0_a[:,:,:]
var0=ma.masked_array(var0,mask=None,fill_value=None)
#
f=cdms.open(fileIG)
thick_in=f(varIT,squeeze=1)
f.close()
#
f=cdms.open(fileM)
mask=f(varM,squeeze=1)
f.close()
if mask.rank() != 2 :
    print "2d input mask please !"
dims_mask=mask.getOrder()
if dims_mask[0] == 'x' : 
    mask=N.transpose(mask)
if mask.shape[1] != lx0 or mask.shape[0] != ly0 :
    print "same dims for mask and input file please !"
#
f=cdms.open(fileOG)
thick_out=f(varOT,squeeze=1)
depth_out=f(varOD,squeeze=1)
f.close()
lz1=thick_out.shape[0]
if depth_out.shape[0] != lz1 :
    print "same dims for depth and thickness please !"
#
# 2. 2D to 3D mask file
# =======================
#
mask_in=N.zeros((lz0,ly0,lx0))
mask_in=mask_in.astype('d')

for i in N.arange(lx0) :
    for j in N.arange(ly0) :
        if mask[j,i] > 0 :
            down=mask[j,i]
            for k in N.arange(down) :
                 mask_in[k,j,i]=1.	    

# 3. Vertical interpolation
# ==========================
#
var1=N.zeros((lz1,ly0,lx0),'d')
thick_bis=N.zeros((ly0,lx0),'d')
thick_jk1=N.zeros((ly0,lx0),'d')
mask_out=N.zeros((lz1,ly0,lx0),'d')
#
jk0=0
for jk1 in N.arange(lz1) :
    while ( (jk0 <= (lz0-1)) & (N.add.reduce(thick_in[0:(jk0+1)]) <= N.add.reduce(thick_out[0:(jk1+1)])) ) :
        thick_bis=N.multiply(N.minimum(thick_in[jk0],N.add.reduce(thick_in[0:(jk0+1)])-N.add.reduce(thick_out[0:jk1])),mask_in[jk0,:,:])
        var1[jk1,:,:]=var1[jk1,:,:]+N.multiply(var0[jk0,:,:],thick_bis)
        thick_jk1=thick_jk1+thick_bis
        jk0=jk0+1
    if ( jk0 <= (lz0-1)) :
      thick_bis=N.multiply(N.minimum(thick_out[jk1],N.add.reduce(thick_out[0:(jk1+1)])-N.add.reduce(thick_in[0:jk0])),mask_in[jk0,:,:])
      var1[jk1,:,:]=var1[jk1,:,:]+N.multiply(var0[jk0,:,:],thick_bis)
      thick_jk1=thick_jk1+thick_bis
    mask_out[jk1,:,:]=N.where(thick_jk1>0.01,1.,0.)
    thick_jk1=N.where(thick_jk1==0.,1.,thick_jk1)
    var1[jk1,:,:]=var1[jk1,:,:]/thick_jk1
    thick_jk1[:,:]=0

lon=cdms.createAxis(N.arange(lx0),id='x')
lat=cdms.createAxis(N.arange(ly0),id='y')
depth=cdms.createAxis(depth_out,id='z')

var2=cdms.createVariable(var1,id=var0_a.id)
var2.getAxis(0)[:]=depth_out
var2.getAxis(0).id='z'
var2.getAxis(1).id='y'
var2.getAxis(2).id='x'
var3=MV.where(mask_out>0.,var1,1e20)
var3.id=var0_a.id
var3.getAxis(1).id='y'
var3.getAxis(2).id='x'
var3.getAxis(0).id='z'
var3.getAxis(0)[:]=depth_out

g=cdms.open(fileOUT,'w')
g.write(var3)
g.close()