#!/usr/bin/env cdat
# -*- coding: iso-8859-15 -*-

import cdms, cdutil, genutil, os, sys, getopt
from cdms import MV
from dataproc.utils import cdat
from regrid import Regridder

from Numeric import ones

def usage():
    """Display usage information."""
    
    man = """
    This script will get 2 SECHIBA or STOMATE restarts for ORCHIDEE code and will
    produce a regridded restart with the grid of the first restart (and other non gridded variables)
    and the values of the second one.
    
Usage : 
    prompt> restart_interpol.py restart_outgrid.nc restart_invalues.nc'
        This script regrids values in file 2 (namely restart_invlaues.nc) with
        grid in restart_outgrid.nc.

Examples :
    > restart_interpol.py CF_CM1414A_20001231_sechiba_rest.nc CF_CM5PIRC11_29901231_sechiba_rest.nc
    > restart_interpol.py CF_CM1414A_20001231_stomate_rest.nc CF_CM5PIRC11_29901231_stomate_rest.nc
    > restart_interpol.py CF_CM5PIRC8_26601231_sechiba_rest.nc CF_OOL-CTRLPI-1_9_5_SPIN_v2ORC_43_01001231_sechiba_rest.nc
    """
    print man

options = [
    "help",
    "version" ]

# Extract arguments from the command line
try:
    myopts, myargs = getopt.getopt(sys.argv[1:], 'Vh', longopts = options)
except getopt.GetoptError:
    print 'getopt error'
    usage()
    sys.exit(1)

variable=""
if (len(myargs) < 2):
    print 'I need at least 2 files.'
    usage()
    sys.exit(1)
else:
    file = myargs[:]
    lf = len(file)
    for i in range(2):
    	if (os.path.exists(file[i])==False):
            print 'file number ',i+1,' : ',file[i],\
                  ' doesn''t exist !' 
            sys.exit(1)

f1=cdms.open( file[0],'r')
f2=cdms.open( file[1],'r')
listv=f1.listvariables()
listv2=f2.listvariables()
#listv.sort()
#listv2.sort()
print listv
print listv2


f3=cdms.open( "regrid_"+file[1],'w')

onet=None

missing_value = 1e+20

for var in listv:
    print var
    if (var in ["nav_lon","nav_lat","nav_lev","time_steps","time",
                "lat","bounds_lat","lon","bounds_lon","lev","day_counter",
                "dt_days", "routingcounter", "veget_year", "date" ]):
        var1=f1(var)
        f3.write(var1)

        print "remplacé."
        continue
        
    var1=f1(var)
    try:
        var2=f2(var)
    except:
        f3.write(var1)
        print "remplacé."
        continue

#    var2.info()
    if (len(var1.shape) > 3):
       var_1=f1(var, squeeze=1)
       tmp=cdms.createVariable(var1)
       var_2=f2(var, squeeze=1)
#       print var1.id,var

#        if (var1.id == 'day_counter'):
#            var1[0,0,0,0]=var_2
#            continue
       
#       print tmp.shape,var_2,type(var_2)
#       axlist2=var_2.getAxisList()
#       print tmp.shape,var_2,len(axlist2)
#       if (len(axlist2) > 0 ):

       if (len(var_2)!=1) :
           if (len(var_2.shape) > 2):
#               print tmp.shape, var_2.shape,var2.shape
               if ( len(tmp.getAxis(0)) == 1 ):
#                   print tmp.getAxis(0),var2.getAxis(0)

                   if ( len(var_2.getAxis(0)) > 0 ):
#                       print var_2.getAxis(0)[:]
                       shift=( var_2.getAxis(0)[0] > 0 )
                       outgrid=var_1[0,:,:].getGrid()
                       ingrid=var_2[0,:,:].getGrid()
                       regridfunc = Regridder(ingrid, outgrid)
                       sumpft=var_1[0,:,:]*0.
                       for fpft in var_2.getAxis(0):
                           if ( shift ):
                               pft=int(fpft) - int(var_2.getAxis(0)[0])
                           else:
                               pft=int(fpft)
                               
#                           print pft
                           tmp_ = tmp[0,pft,:,:]
                           v2 = var_2[pft,:,:]
                           v1 = var_1[pft,:,:]
 #                           mask_undef = MV.where(v2 > missing_value, 1, 0)
#                            v2_mask=MV.masked_array(v2, mask_undef)
#                            varmax2=max(max(abs(v2_mask)))
#                            v2=v2_mask/varmax2
                           tmp__ = regridfunc(v2)
#                            tmp[0,pft,:,:]=tmp__*varmax2                           
                           tmp[0,pft,:,:]=tmp__
                           if (var in ["veget_max","veget"]):
                               sumpft[:,:]=sumpft[:,:]+tmp__

                       if (var in ["veget_max"]):
                           # On veut :
                           #somm(tmp[0,pft,:,:])+frac_nobio=1
                           vegnorm=1-frac_nobio
                           veget_max=tmp[0,:,:,:]
                           for fpft in var_2.getAxis(0):
                               if ( shift ):
                                   pft=int(fpft) - int(var_2.getAxis(0)[0])
                               else:
                                   pft=int(fpft)
                               veget_max[pft,:,:]=MV.where(sumpft[:,:] == 0., 0., veget_max[pft,:,:] * vegnorm / sumpft[:,:])
                               tmp[0,pft,:,:]= veget_max[pft,:,:]
                               
                       if (var in ["veget"]):
                           for fpft in var_2.getAxis(0):
                               if ( shift ):
                                   pft=int(fpft) - int(var_2.getAxis(0)[0])
                               else:
                                   pft=int(fpft)
                                   
                               tmp[0,pft,:,:]= MV.where(tmp[0,pft,:,:] > veget_max[pft,:,:], veget_max[pft,:,:], tmp[0,pft,:,:])

                   else:
                       tmp[0,0,:,:]=var_2.regrid( var_1.getGrid() )
               else:
                   print "Erreur ! "
                   print tmp.getAxis(0),var2.getAxis(0)
#                   print var1.getAxisList()
                   
                   if ( len(var_2.getAxis(0)) > 0 ):
#                       print var_2.getAxis(0)[:]
                       shift=( var_2.getAxis(0)[0] > 0 )
                       for fpft in var_2.getAxis(0):
                           if ( shift ):
                               pft=int(fpft) - int(var_2.getAxis(0)[0])
                           else:
                               pft=int(fpft)
                               
#                           print pft
                           tmp_ = tmp[0,pft,:,:]
                           v2 = var_2[pft,:,:]
                           v1 = var_1[0,pft,:,:]
#                           mask_undef = MV.where(v2 > missing_value, 1, 0)
#                           v2_mask=MV.masked_array(v2, mask_undef)
#                           varmax2=max(abs(v2_mask))
#                           v2=v2_mask/varmax2
                           tmp__ = v2.regrid( v1.getGrid() )
#                           tmp[0,pft,:,:]=tmp__*varmax2
                           tmp[0,pft,:,:]=tmp__
                           
           else:
               tmp[0,0,:,:]=var_2.regrid( var_1.getGrid() )
               if (var in ["frac_nobio"]):
                   frac_nobio=tmp[0,0,:,:]
       else:
           tmp[0,0,0,0]=var_2

       var1=tmp

       #       var1.info()
       t=var1.getTime()
       print 'len of t:',t
       if t is not None:
           if len(t)!=1:
               raise 'Error too much time'
       if onet is None:
           onet=t.clone()
       else:
           if t[0]!=onet[0]:
               raise 'different time'

       f3.write(var1)
    else:
        print "impossible ?"
#        var2=f2(var)
#        f3.write(var2)