var1.getAxis(1).designateLatitude()
var1.getAxis(2).id = 'lon'
var1.getAxis(2).designateLongitude()
var1.units = 'm'
var1.id = 'ssh'
var2 = var1.clone()
var2[:] = N.random.random((nt, ny, nx))
var1[3:13, :1, :1] = MV2.masked
var2[5:15, -1:, -1:] = MV2.masked
var2.long_name = 'Sea level'
var2.id = 'sla'
mask = var1.mask|var2.mask # common mask
vmax = var2.max()
bins = N.linspace(-0.1*vmax, 0.9*vmax, 14)
nbins = len(bins)
restart_file5 = code_file_name(ext='5.nc')
restart_file7 = code_file_name(ext='7.nc')
print restart_file5

# Normal
sab = StatAccum(tall=True, sall=True, bins=bins)
sab += var1[:5], var2[:5]
#print sab.get_tmean()

# Dump
sab.dump(restart_file5)
sab += var1[5:7], var2[5:7]
sab.dump(restart_file7)

# Load from scratch
sa5 = StatAccum(restart=True, restart_file=restart_file5)

nx = ny = 50
np = 500
mtype = 'gauss'
dmax = 2.5

from vcmq import N, P, code_file_name, savefigs
from vacumm.misc.grid.kriging import gridded_gauss3, random_gauss3, variogram, variogram_fit

# Generate random field
xxg, yyg, zzg = gridded_gauss3(nx=nx, ny=ny)
x, y, z = random_gauss3(np=np)

# Variogram from data
d, v = variogram(x, y, z, dmax=dmax)

# Variogram fit with null nugget
vm = variogram_fit(x, y, z, mtype, n=0, dmax=dmax)
D = N.linspace(0, d.max())
V = vm(D)

# Compare
P.figure(figsize=(6, 4))
P.title('Variogram')
P.plot(d, v, 'ob', label='From data')
P.plot(D, V, '-r', label='Fitted model (%s)'%mtype)
P.legend(loc='best')
P.ylim(ymin=0)
savefigs(code_file_name(ext=False), pdf=True, verbose=False)
P.close()

# Interpolate
zzo = krig(xxi.ravel(), yyi.ravel(), zzi.ravel(), xxo.ravel(), yyo.ravel())
zzo.shape = xxo.shape

# Section
P.figure(figsize=(8, 4))
iyis = [3, 4]
for iyi in iyis:
    label = iyi==iyis[0]
    P.plot(xi, zzi[iyi], 'ob-', markersize=8, label='Original' if label else None)
    P.plot(xo, zzo[iyi*r], 'or-', markersize=5, lw=.8, label='Interpolated' if label else None)
    P.legend(loc='best', framealpha=0.5)
P.grid()
P.title('Section')
P.tight_layout()
savefigs(code_file_name(ext='_0.png'), verbose=False, pdf=True)

# Maps
P.figure(figsize=(8, 4))
levels = auto_scale(vmin=zzi.min(), vmax=zzi.max(), nmax=30)
P.subplot(121)
P.contourf(xxi, yyi, zzi, levels=levels)
P.contour(xxi, yyi, zzi, linewidths=0.1, levels=levels, colors='k')
add_grid((xxi, yyi), alpha=.3, centers=True)
for iyi in iyis:
    P.axhline(yi[iyi], linestyle='--', color='k')
P.title('Original')
P.subplot(122)
P.contourf(xxo, yyo, zzo, levels=levels)
P.contour(xxo, yyo, zzo, linewidths=0.1, levels=levels, colors='k')
add_grid((xxo, yyo), alpha=.2, centers=True)

"""Test :func:`~vacumm.misc.grid.misc.coord2slice`"""

from vcmq import code_file_name, P, os, N, create_lon, code_file_name
from vacumm.misc.grid import coord2slice, create_grid, create_axes2d, meshbounds
from vacumm.misc.plot import add_grid

figfiles = []
figfile = code_file_name(ext=False)+'_%i.png'
def plot(xx, yy, target, label, figfiles, figfile, lon=None, lat=None, show=False):
    xs, ys, mask = coord2slice(target, lon=lon, lat=lat)
    P.figure(figsize=(6, 3.5))
    P.title('Target=%(label)s / select: lon=%(lon)s, lat=%(lat)s'%locals())
    add_grid((xx, yy))
    xx = xx.asma()
    yy = yy.asma()
    if isinstance(lon, tuple): 
        P.axvline(lon[0], color='m', ls='--', lw=2)
        P.axvline(lon[1], color='m', ls='--', lw=2)
    elif isinstance(lon, slice):
        i, j, k = lon.indices(xx.shape[1])
        P.plot(xx[:, i], yy[:, i], 'c--', lw=2)
        P.plot(xx[:, j-1], yy[:, j-1], 'c--', lw=2)
    if isinstance(lat, tuple): 
        P.axhline(lat[0], color='m', ls='--', lw=2)
        P.axhline(lat[1], color='m', ls='--', lw=2)
    elif isinstance(lat, slice):
        i, j, k = lat.indices(yy.shape[0])
        P.plot(xx[i], yy[i], 'c--', lw=2)
        P.plot(xx[j-1], yy[j-1], 'c--', lw=2)
    P.xticks(N.arange(xx.min()-1, xx.max()+1))
    P.yticks(N.arange(yy.min()-1, yy.max()+1))

nens = 50
# - references
ref1 = MV2.sin(MV2.arange(nt, dtype='f'))*10.
ref2 = ref1*1.2
# - ensemble 1
dat1 = MV2.resize(ref1, (nens, nt))+N.resize(N.random.rand(nt*nens), (nens, nt))*20.
dat1.long_name = 'Set 1'
dat1.units = 'meters'
# - modele 2
dat2 = MV2.resize(ref2, (nens, nt))+N.resize(N.random.rand(nt*nens), (nens, nt))*40.
dat2.long_name = 'Set 2'

# Plot
rc('font', size=8)
kwplot = dict(figsize=(5,5), savefigs_pdf=False, show=False, close=True, savefigs_verbose=False)
basename = code_file_name(ext=False)+'_%i'
# - single variable
taylor(dat1[0], ref1, savefigs=basename%0, colors='cyan', title='Basic', **kwplot)
# - single variable without normalization
taylor(dat1[0], ref1, savefigs=basename%1, normalize=False, reflabel='Reference', 
    title='No normalization', **kwplot)
# - several variables with same ref
taylor([d for d in dat1], ref1, colors='cyan', savefigs=basename%2, 
    title='Several variables', **kwplot)
# - several variables with same ref with specified colors
taylor([d for d in dat1[:3]], ref1, colors = ['r', 'g', 'b'], 
    labels=['M#1', 'M#2', 'M#3'], savefigs=basename%3, 
    title='Explicit colors', **kwplot)
# - several variables with same ref with numeric scalar colors
taylor([d for d in dat1[:3]], ref1, colors = [0.,  10.,  5.], cmap='jet', 
    savefigs=basename%4, title='Scalar colors', **kwplot)

dx = N.diff(xi[0]).mean()
dy = N.diff(yi[:, 0]).mean()
xo = N.arange(xi.min()+10*dx, -30*dx+xi.max(), dx)
yo = N.arange(yi.min()-20*dy, yi.max()-20*dy, dy)
lono = cdms2.createAxis(xo)
lono.designateLongitude() ; lono.units= 'degrees_east'
lato = cdms2.createAxis(yo)
lato.designateLatitude() ; lato.units = 'degrees_north'
xob = bounds1d(lono) ; lono.setBounds(xob)
yob = bounds1d(lato) ; lato.setBounds(yob)
grido = cdms2.createRectGrid(lato, lono)
xmin, xmax = minmax(loni.asma(),lono)
ymin, ymax = minmax(lati.asma(), lato)
nyo,nxo = grido.shape
#print 'rank',rank
basefile = code_file_name(ext=False)
repfile = basefile+'.nt%(nt)s-nz%(nz)s-nyi%(nyi)s-nxi%(nxi)s.log'%locals()
if rank==0: 
    if os.path.exists(repfile): os.remove(repfile)
    f = open(repfile, 'w')
    if size:
#        print 'MPI', size
        print >>f, 'MPI: NPROC=%i'%size
    print >>f, 'NT=%(nt)i, NZ=%(nz)i'%locals()
    print >>f, 'NYI=%(nyi)i, NXI=%(nxi)i'%locals()
    print >>f, 'NYO=%(nyo)i, NXO=%(nxo)i'%locals()

#if rank==0: print 'regridding...'
vmin = vari.min()
vmax = vari.max()
for tool, methods in config.items():

from vacumm.misc.grid import rotate_grid
from time import time

# Input
nx = ny = 300
vari = MV2.array(N.arange(nx*ny*1.).reshape(ny, nx))
gridi = create_grid2d(vari.getAxis(1)[:]*50/nx,  vari.getAxis(0)[:]*50/nx)
set_grid(vari, gridi)

# Output grid
gridor = create_grid2d(vari.getAxis(1)[:]*0.09*50/nx,  
    vari.getAxis(0)[:]*0.09*50/nx)
gridoc = rotate_grid(gridi, 30)

# Log
logfile = code_file_name(ext='log')
if os.path.exists(logfile): os.remove(logfile)
f = open(logfile, 'w')
print >>f, 'NY=%(ny)i, NX=%(nx)i'%locals()

# Loop on methods
for tool, methods in config.items():
    for method in methods:
#        print tool.upper(), method
        print >>f, tool.upper(), method
        
        # rect...
        t0 = time()
        r = CDATRegridder(vari, gridor, tool=tool, method=method)
        t1 = time()
        dt = t1-t0

yyob, xxob  = meshcells(yyo, x)

varon = N.ma.masked_values(interp1dxx(vari.filled(), yyi, yyo, mv, 0, extrap=0), mv)
varol = N.ma.masked_values(interp1dxx(vari.filled(), yyi, yyo, mv, 1, extrap=0), mv)
varoh = N.ma.masked_values(interp1dxx(vari.filled(), yyi, yyo, mv, 3, extrap=0), mv)

kw = dict(vmin=vari.min(), vmax=vari.max())
axlims = [x[0], x[-1], yo[0], yo[-1]]
P.figure(figsize=(8, 8))
P.subplot(221)
P.pcolor(xxib, yyib, vari)
P.axis(axlims)
P.title('Original')
P.subplot(222)
P.pcolor(xxob, yyob, varon, **kw)
P.axis(axlims)
P.title('Nearest1dxx')
P.subplot(223)
P.pcolor(xxob, yyob, varol, **kw)
P.axis(axlims)
P.title('Linear1dxx')
P.subplot(224)
P.pcolor(xxob, yyob, varoh, **kw)
P.axis(axlims)
P.title('Hermit1dxx')
P.tight_layout()
figfile = code_file_name(ext='png')
if os.path.exists(figfile): os.remove(figfile)
P.savefig(figfile)
P.close()

var1.getAxis(1).designateLatitude()
var1.getAxis(2).id = 'lon'
var1.getAxis(2).designateLongitude()
var1.units = 'm'
var1.id = 'ssh'
var2 = var1.clone()
var2[:] += N.ma.sin(var1)*100
var1[3:13, 5:9, 3:7] = MV2.masked
var2[5:15, 7:12, 1:5] = MV2.masked
var2.long_name = 'Sea level'
var2.id = 'sla'
mask = var1.mask|var2.mask # common mask
vmax = var2.max()
bins = N.linspace(-0.1*vmax, 0.9*vmax, 14)
nbins = len(bins)
restart_file = code_file_name(ext='nc')
print restart_file

# Normal
sa0 = StatAccum(tall=True, sall=True, bins=bins)
sa0 += var1[:7], var2[:7]

# Dump
sa0.dump(restart_file)

# Load from scratach
sa1 = StatAccum(restart=True, restart_file=restart_file)

## Intermediate load
#sa1 = StatAccum(tall=True, sall=True, bins=bins)
#sa1 += var1[:5], var2[:5]

add_grid(gridci, **kwg)
# - dstwgt c2c / vacumm
varo = regrid2d(varci, gridco, method="dstwgt", tool="vacumm")
plot2d(varo, title="VACUMM / Dstwgt", subplot=(4, 3, ip), **kw)
ip += 1
add_grid(gridci, **kwg)
# - patch c2c / emsf
varo = regrid2d(varci, gridco, method="patch", tool="esmf")
plot2d(varo, title="ESMF / Patch", subplot=(4, 3, ip), **kw)
ip += 1
add_grid(gridci, **kwg)
# - cellave r2r / regrid2
varo = regrid2d(varri, gridro, method="cellave", tool="regrid2")
plot2d(varo, title="Regrid2 / Cellave", subplot=(4, 3, ip), **kw)
ip += 1
add_grid(gridri, **kwg)
# - cellave c2c / esmf
varo = regrid2d(varci, gridco, method="cellave", tool="esmf")
plot2d(varo, title="ESMF / Cellave", subplot=(4, 3, ip), **kw)
ip += 1
add_grid(gridci, **kwg)
# - conserv c2c / esmf
varo = regrid2d(varci, gridco, method="conserv", tool="esmf")
plot2d(varo, title="ESMF / Conserv", subplot=(4, 3, ip), **kw)
ip += 1
add_grid(gridci, **kwg)
P.tight_layout()
P.savefig(code_file_name(ext="png"))
rcdefaults()
P.close()

"""Test :meth:`~vacumm.data.misc.dataset.OceanDataset.get_uvgbt` on MENOR"""

# Inits
ncfile = "menor.nc"

# Imports
from vcmq import data_sample, DS, N, MV2, shapiro2d, map2, code_file_name

# Read data
ds = DS(data_sample(ncfile), 'mars', logger_level='critical')
u, v = ds.get_uvgbt()

# Mask huge values
mask = (N.ma.sqrt(u**2+v**2)>2.).filled(True)
for var in u, v:
    var[:] = MV2.masked_where(mask, var, copy=False)

# Shapiro filter
for var in u, v:
    var[:] = shapiro2d(var)

# Plot
map2((u[0], v[0]), title='Geostrophic velocity', 
    quiver_samp=3, xymasked=False, quiver_norm=3, contour=False, 
    fill=False, figsize=(7, 7), quiver_linewidth=.3, quiver_width=0.003, 
    quiver_scale=10, right=1, colorbar_shrink=0.8, bottom=0.05, show=False, 
    savefig=code_file_name(), close=True)

from vcmq import data_sample, code_file_name

# Time
#t0 = round_date(now(utc=False), 'day', 'floor')
t0 = round_date(nowtime, 'day', 'floor')
t1 = t0.add(ndayfore, cdtime.Day)
t0 = paris_to_utc(t0)
t1 = paris_to_utc(t1)

# Read
ds = DS(data_sample(ncpat), 'mars', time=(t0, t1), logger_level='critical')
data = ds.get(field)
utc_to_paris(data)

# Read some config values
cfgfile = code_file_name(ext='cfg')
load_cfg(cfgfile)
long_name = cfgget('long_name', field)
long_long_name = cfgget('long_name_model', field)
units = cfgget('units', field)
vmin = cfgget('vmin', field)
vmax = cfgget('vmax', field)
points = cfgget('points', field)

# Interpolate to points
lons = N.array([d[0] for d in points])
lats = N.array([d[1] for d in points])
ptdata = grid2xy(data, lons, lats)
datas = []
for i in xrange(ptdata.shape[1]):
    datas.append(ptdata[:, i])

    P.title('Cellave')
    P.tight_layout()
    P.savefig(figfile)


# 1d->1d
depi1d = create_dep(N.arange(-4500., 1, 500))
depo1d = create_dep(N.arange(-4000., 1, 333.33))
nzi = depi1d.shape[0]
vari = MV2.asarray(N.ma.resize(depi1d[:], (nt, ny, nx, nzi)).transpose([0, 3, 1, 2]))
vari.setAxis(1, depi1d)
varol1 = regrid1d(vari, depo1d, method='linear')
varol2 = regrid1d(vari, depo1d, method='linear', iaxi=0, iaxo=0, axi=depi1d)
result.append(('assertEqual', [(varol1-varol2).std(), 0]))
varoc = regrid1d(vari, depo1d, method='cellave')
myplot(vari, depi1d, varol1, varoc, depo1d, code_file_name(ext='_0.png'))


# 4d->1d
depi1d = N.arange(-4500., 1, 500)
nzi = depi1d.shape[0]
depi4d = N.resize(N.resize(depi1d, (nx, ny, nzi)).T, (nt, nzi, ny, nx))
depi4d += 500*(N.random.random(depi4d.shape)-0.5)
depo1d = create_dep(N.arange(-4000., 1, 333.33))
vari = MV2.array(depi4d, fill_value=1e20)
vari.getAxis(1).designateLevel()
depi4d = MV2.asarray(depi4d)
depi4d.getAxis(1).designateLevel()
varol1 = regrid1d(vari, depo1d, method='linear', axi=depi4d)
varol2 = regrid1d(vari, depo1d, method='linear', iaxi=1, axi=depi4d)
result.append(('assertEqual', [(varol1-varol2).std(), 0]))

yi = N.array(yi)
zi = N.array(zi)
xxg, yyg = N.meshgrid(xg, yg)
xo = xxg.ravel()
yo = yyg.ravel()
vgm = variogram_model('linear', n=0, s=sill, r=range)

# Setup the kriger
sck = SimpleCloudKriger(xi, yi, zi, vgf=vgm, farvalue=farvalue)

# Interpolate
zo = sck(xo, yo)

# Check a far value
zzg = zo.reshape(ny, nx)
N.testing.assert_allclose(zzg[-1, -1],farvalue)

# Plot
vmin = min(zi.min(), zo.min())
vmax = max(zi.max(), zo.max(), farvalue)
cmap = 'cmocean_ice'
kw = dict(vmin=vmin, vmax=vmax)
P.pcolormesh(xxg, yyg, zzg, cmap=cmap, **kw)
P.colorbar()
P.scatter(xi, yi, c=zi, s=100, cmap=cmap, **kw)
P.axis('image')
P.title('Simple kriging with fixed far value')
P.savefig(code_file_name(ext='.png'))
P.show()
P.close()

rr = P.linspace(0, P.sqrt(xg.ptp()**2+yg.ptp()**2), 100)
vv = ock.variogram_func(rr)

# Plot
P.figure(figsize=(8, 8))
axis = [xg.min(), xg.max(), yg.min(), yg.max()]
kwmm = dict(vmin=zi.min(), vmax=zi.max())
kwsc = dict(lw=0.2, **kwmm)
ax = P.subplot(221, aspect=1)
P.imshow(zzg,  extent=axis, interpolation='bilinear', origin='lower', alpha=.2, **kwmm)
P.scatter(xi, yi, c=zi, s=20, lw=0.2, **kwmm)
P.axis(axis)
P.title('Input points')
P.subplot(222, aspect=1)
P.scatter(xo, yo, c=zo, s=40, lw=0.2, **kwmm)
P.axis(axis)
P.title('Interpolated points')
P.subplot(223, aspect=1)
P.scatter(xo, yo, c=zoe, s=40, lw=0.2, vmin=0, vmax=100)
P.axis(axis)
P.colorbar(shrink=.8)
P.title('Relative interpolation error')
P.subplot(224)
P.plot(rr, vv)
P.xlabel('Distance')
P.ylabel('Error')
P.title('Fitted variogram')
P.tight_layout()
savefigs(code_file_name(), verbose=False)
P.close()

"""Test :meth:`~vacumm.misc.core_plot.Plot2D.plot_streamplot`"""

from vcmq import map2, data_sample, cdms2, code_file_name
import warnings
warnings.filterwarnings('ignore', 'Warning: converting a masked element to nan')

f = cdms2.open(data_sample('mars2d.xyt.nc'))
u = f('u', slice(0, 1), squeeze=1)
v = f('v', slice(0, 1), squeeze=1)
f.close()

map2((u, v), fill=False, contour=False, streamplot=True, streamplot_density=3, 
    streamplot_linewidth='modulus', streamplot_lwmod=3, streamplot_color='modulus', 
    title='Tidal stream lines', colorbar_shrink=0.7, savefig=code_file_name(), 
    right=1, show=False)

# Interpolations
for p in 'u', 'v', 'f', 'w':
    var[p] = ag.interp(var['t'], 't', p, mode='extrap')

# Surface plots
vmin, vmax = minmax(*[var[p][-1] for p in ['u', 'v', 'f']])
kw = dict(show=False, res=None, vmin=vmin, vmax=vmax, colorbar=False, grid=False, cmap='jet')
m = map2(var['t'][-1], fill='pcolor', 
    title='Interpolations on an Arakawa C grid: T->U/V/F', **kw)
add_grid(var['t'], linestyle='-')
kw.update(fill='scatter', contour=False, fill_s=60)
markers = dict(u='>', v='^', f='D', t='o')
for p in 't', 'u', 'v', 'f':
    m = map2(var[p][-1], fill_marker=markers[p], shadow=True, zorder=100, **kw)
m.savefig(code_file_name(ext='_1.png'))
m.close()

# Vertical plot
curve2(var['t'][:, 0, 0], 'o-b', ymax=0, show=False, 
    title='Interpolations on an Arakawa C grid: T->W')
curve2(var['w'][:, 0, 0], '^r', show=False, savefig=code_file_name(ext='_2.png'))


result = [
    ('assertEqual', [var['t'][0, 0, :2].mean(), var['u'][0, 0, 0]]), 
    ('assertEqual', [var['t'][0, :2, 0].mean(), var['v'][0, 0, 0]]), 
    ('assertEqual', [var['t'][0, :2, :2].mean(), var['f'][0, 0, 0]]), 
    ('assertEqual', [var['t'][:2, 0, 0].mean(), var['w'][0, 0, 0]]), 
    ]

"""Test :meth:`~vacumm.data.misc.OceanDataset.plot_hsection` with MFS"""

# Inits
ncfile = "mfs.nc"
depth = -1000.0

# Imports
from vcmq import DS, data_sample, os, code_file_name

# Setup dataset
ds = DS(data_sample(ncfile), "nemo", logger_level="critical")

# Plot hsection
figfile = code_file_name(ext="png")
if os.path.exists(figfile):
    os.remove(figfile)
ds.plot_hsection("temp", depth, savefig=figfile, fill="contourf", show=False, close=True)

# Output curved grid
nxo = 7
nyo = 7
xxo, yyo = N.meshgrid(N.arange(nxo)+.5, N.arange(nyo)-.5)
grido = create_grid2d(xxo, yyo) # output cdms grid
xxob,yyob = meshbounds(xxo,yyo) # coordinates of cell corners

# Input field
vari = MV2.array(N.arange(nyi*nxi).reshape(nyi,nxi))+10.
vari[1,1:3] = 100
vari[:2,1:3] = MV2.masked
set_grid(vari, gridi) # set grid and axes
#gridi.setMask(vari.mask)

# Define plot function
figfile = code_file_name(ext=False)+'_%(ifig)i.png'
#'%(tool)s_%(method)s.png'
figfiles = []
rc('font',size=9)
kw = dict(vmin=vari.min(),vmax=vari.max())
    
# Define logger
logfile = code_file_name(ext='log')
f = open(logfile, 'w')
def log(f, text): # logger
#    print text
    print >>f, text

# Loop on methods and tools
for tool, methods in configs.items():
    for method in methods: