#!/usr/bin/env python
#
# Author: Patrick Brockmann 
# Contact: Patrick.Brockmann@ipsl.jussieu.fr 
# Release: 1.02 
# Date: 13/03/2003
#

import sys,re
import cdms,vcs,Numeric
 
##################################
usage = """

Usage: vcs_demo.py [-h] [-p projection] [-m] [-n nblevels] gridCF_file var_file var 

Options:
        -h, -?, --help, -help
                Print this manual.
        -p, --projection
                Projection to choose in ('linear','polar','mollweide').
        -m, --mesh
                View mesh of the grid.
        -n, --nblevels
                Number of levels should be in [2:100] 

"""

##################################
option_projection='linear'
option_mesh_visible=0
option_nblevels=20
while len(sys.argv[1:]) != 0 :
        if sys.argv[1] in ('-h','--help') :
                del(sys.argv[1])
                print usage 
                sys.exit(1)
        elif sys.argv[1] in ('-p','--projection') :
                option_projection=sys.argv[2]
                del(sys.argv[1])
                del(sys.argv[1])
        elif sys.argv[1] in ('-m','--mesh') :
                option_mesh_visible=1
                del(sys.argv[1])
        elif sys.argv[1] in ('-n','--nblevels') :
                option_nblevels=int(sys.argv[2])
                del(sys.argv[1])
                del(sys.argv[1])
        elif re.match('-',sys.argv[1]) :
                print 'option inconnu'
                break
        else:
                break

if len(sys.argv[1:]) != 3 :
        print usage
        sys.exit(1)

if option_projection not in ('linear','polar','mollweide'):
	print usage
        sys.exit(1)
        
if not(1<option_nblevels<101) :
	print usage
        sys.exit(1)
        
################################################################
# Limit excursion of the vertex of the cell according to 
# its center position
def limit_lon(blon,clon,CM):
    clon=(clon+360*10)%360
    clon=Numeric.where(Numeric.greater(clon,CM+180),clon-360,clon)
    clon=Numeric.where(Numeric.less(clon,CM-180),clon+360,clon)
    clon1=Numeric.zeros(blon.shape,Numeric.Float)
    clon1[:,:,0]=clon1[:,:,1]=clon1[:,:,2]=clon1[:,:,3]=clon

    blon=(blon+360*10)%360
    blon=Numeric.where(Numeric.greater(blon,CM+180),blon-360,blon)
    blon=Numeric.where(Numeric.less(blon,CM-180),blon+360,blon)

    blon=Numeric.where(Numeric.greater(abs(blon-clon1),abs(blon+360-clon1)),blon+360,blon)
    blon=Numeric.where(Numeric.greater(abs(blon-clon1),abs(blon-360-clon1)),blon-360,blon)
    return blon

################################################################
# Create mesh according to meshfill method
def make_mesh(blon,blat):
    sh=list(blon.shape)
    sh.insert(0,2)
    M=Numeric.zeros(sh,Numeric.Float)
    M[0]=blat
    M[1]=blon
    reorder=(1,2,0,3)
    M=Numeric.transpose(M,reorder)
    return M

################################################################
print '********************************'
print 'Program vcs_demo.py'

x=vcs.init()

x.setcolormap('default')
x.open()

#############################
if option_projection == 'polar' :
	t=x.createtemplate('new','polar')
	x.geometry(640,640,0,0)
elif option_projection == 'mollweide' :
	t=x.createtemplate('new','molwd')
	x.geometry(640,480,0,0)
else :
	t=x.createtemplate('new','default')
	x.geometry(640,480,0,0)

#############################
m=x.createmeshfill('new')
m.mesh=option_mesh_visible
m.wrap=[0,360]

if option_projection == 'polar' :
	m.projection = 'polar'
elif option_projection == 'mollweide' :
	m.projection = 'mollweide'
else:
	m.projection = 'linear'
	m.datawc_x1=0.
	m.datawc_x2=450.
	m.datawc_y1=-90.
	m.datawc_y2=90.

print '********************************'
print 'Reading grid file'
try:
	f1 =cdms.open(sys.argv[1])
	klevel=0
	mask=f1('mask',depth=slice(klevel,klevel+1))[0]
	print 'mask rank: ', mask.rank()
	print 'mask axis: ', mask.getAxisIds() 
	print 'mask shape: ', mask.shape
	blon=f1('bounds_lon').filled()
	blat=f1('bounds_lat').filled()
        clon=f1('lon').filled() 
except:
	print 'Error in grid file processing'
        sys.exit(1)

print '********************************'
print 'Reading var file'
try:
	f2=cdms.open(sys.argv[2])
	#var=f2(sys.argv[3])
	var=f2[sys.argv[3]]
	print 'var rank: ', var.rank() 
	print 'var shape: ', var.shape
	print 'var axis: ', var.getAxisIds() 
	if var.rank() == 3:
		print '3'
		var=var(slice(0,1),squeeze=1)
	elif var.rank() == 4:
		print '4'
		var=var(slice(0,1),slice(0,1),squeeze=1)
	print 'var extrema: ', vcs.minmax(var)
	extr = vcs.minmax(var)
	levels= vcs.mkscale(extr[0],extr[1],nc=option_nblevels)
	levels.insert(0,1E+20)
	levels.append(1E+20)
	m.levels= levels
	m.fillareacolors=vcs.getcolors(levels)	

except:
	print 'Error in var file processing'
        sys.exit(1)

print '********************************'
print 'Limiting vertex'
blon=limit_lon(blon,clon,0)
print 'Doing mesh'
M=make_mesh(blon,blat)
print 'Making variable'
M=cdms.asVariable(M)

print '********************************'
print 'Masking mesh'
mask_ocean=cdms.MV.masked_not_equal(mask,1)
selection_frommask=cdms.MV.getmask(mask_ocean)

selection_fromvar=cdms.MV.getmask(var)
selection_both=cdms.MV.mask_or(selection_frommask,selection_fromvar)

for i in range(2):
	for j in range(4):
                M[:,:,i,j]=cdms.MV.masked_where(selection_both,M[:,:,i,j])

M.setMissing(1E20)

print '********************************'
print 'Reshaping mesh and var'
print 'Before ', M.shape, var.shape
M=cdms.MV.reshape(M,[M.shape[0]*M.shape[1],2,4])
var=cdms.MV.reshape(var,[var.shape[0]*var.shape[1]])
print 'After ', M.shape, var.shape 

print '********************************'
print 'Modify default continents'
ln=x.getline('continents')
ln.width=1
ln.type='solid'
ln.color=160

print '********************************'
print 'Plotting'
x.plot(var,M,m,t,continents=1)

print '********************************'
print 'Saving in various formats'
x.cgm( 'example.cgm','r')
x.gif('example.gif','r','l',geometry='640x480')
x.postscript('example.ps')

print '********************************'
print 'Press a key'
ligne=raw_input()