from dynamico import cxios
from dynamico import xios
from dynamico import unstructured as unst
from dynamico import meshes
import numpy as np

# client.__exit__() and context.__exit__() are called at the end of the 'with' blocks
# even if an exception is raised in the body of the block

with xios.Client() as client:
    comm = client.comm
    mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
    print '%d/%d starting'%(mpi_rank,mpi_size)

    nqdyn,nqtot,llm = 1,1,4

    if False:
        nx,ny,llm = 64,64
        Lx,Ly,f = 1.,1.,1.
        mesh = meshes.Cartesian_Mesh(nx,ny,llm,nqdyn,Lx,Ly,f)
    else:
        filename, g, f, radius = 'cart_128_128.nc', 1., 1., None
    #    filename = 'cart_008_008.nc'
        unst.setvar('g',g)
        print 'Reading Cartesian mesh ...'
        def coriolis(lon,lat):
            return f+0.*lon
        meshfile = meshes.DYNAMICO_Format(filename)
        nx,ny=meshfile.nx,meshfile.ny

        pmesh = meshes.Unstructured_PMesh(comm,meshfile)
    #    pmesh.partition_metis()
        pmesh.partition_curvilinear(2,2)
        mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis)

#--------------------------------- write some data ----------------------------------------

    with xios.Context_Curvilinear(mesh,nqtot, 3600) as context:
        lon_i,lat_i = mesh.lon_i, mesh.lat_i
        lat_ik, junk  = np.meshgrid(lat_i, np.arange(llm), indexing='ij')
        lon_ik, junk  = np.meshgrid(lon_i, np.arange(llm), indexing='ij')

        for i in range(48):
            context.update_calendar(i)
#            print 'send_field', i, lat_ik.shape
            context.send_field_primal('ps', lat_i+lon_i)
            context.send_field_primal('theta', lat_ik+lon_ik)