from mpi4py import MPI
comm = MPI.COMM_WORLD
mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
print '%d/%d starting'%(mpi_rank,mpi_size)

import sys
import numpy as np
#import dynamico.wrap as wrap
print 'import dynamico.unstructured'
import dynamico.unstructured as unst
print 'import dynamico.xios'
import dynamico.xios as xios
print 'Done.'

def boundaries(degree,points,lon,lat):
    n, nvertex = len(degree), degree.max()
    bnd_lon, bnd_lat = np.zeros((n,nvertex)),np.zeros((n,nvertex))
    for ij in range(n):
        nb=degree[ij]
        for k in range(nb):
            vertex = points[ij,k]
            bnd_lon[ij,k]=lon[vertex]
            bnd_lat[ij,k]=lat[vertex]
        for k in range(nb,nvertex):   # repeat last vertex if nb<nvertex
            bnd_lon[ij,k]=bnd_lon[ij,nb-1]
            bnd_lat[ij,k]=bnd_lat[ij,nb-1]
    return bnd_lon, bnd_lat

#----------------------------- read MPAS mesh --------------------------------

radian=unst.radian
grid, llm, nqdyn, nqtot = 10242, 1,1,1 # 2562, 10242, 40962  
Omega, radius, g, gh0 = 2.*np.pi/86400., 6.4e6, 1., 2.94e4
N, T, courant = 40, 10800., 1.2 # simulation length = N*T
print 'Omega, planetary PV', Omega, 2*Omega/gh0
sys.stdout.flush()

def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter                                                                                                                
print 'Reading MPAS mesh ...'
sys.stdout.flush()
mesh = unst.MPAS_Mesh('grids/x1.%d.grid.nc'%grid, llm, nqdyn, radius, f)
print '...Done'
sys.stdout.flush()

lon_i, lat_i, lon_v, lat_v = [x*radian for x in mesh.lon_i, mesh.lat_i, mesh.lon_v, mesh.lat_v]
#lon_i, lat_i, lon_v, lat_v = mesh.lon_i, mesh.lat_i, mesh.lon_v, mesh.lat_v

#--------------------------------- initialize XIOS --------------------------

def setup_domain(cat,id, degree,vertex,lon,lat,lon_v,lat_v): # cat in 'domain','domaingroup'
    bnd_lon, bnd_lat = boundaries(degree, vertex, lon_v, lat_v)
    ncell_tot, nvertex = bnd_lon.shape
    mesh = xios.Handle(cat,id)
    mesh.set_attr(type='unstructured')
    mesh.set_attr(ni_glo=ncell_tot, ibegin=0, ni=ncell_tot)
    mesh.set_attr(i_index=np.arange(ncell_tot,dtype=np.int32))
    mesh.set_attr(lonvalue_1d=lon, latvalue_1d=lat)
    mesh.set_attr(nvertex=nvertex,bounds_lon_1d=bnd_lon, bounds_lat_1d=bnd_lat)

unst.ker.dynamico_setup_xios()

if True:
    lev, levp1, nq = [xios.Handle('axis',name) for name in 'lev', 'levp1', 'nq']
    lev.set_attr(n_glo=llm, value=np.arange(llm))
    levp1.set_attr(n_glo=llm+1, value=np.arange(llm+1))
    nq.set_attr(n_glo=nqtot, value=np.arange(nqtot))
    mesh_i=setup_domain('domaingroup','i', mesh.primal_deg, mesh.primal_vertex-1, lon_i, lat_i, lon_v, lat_v) # primal mesh
#    mesh_v=setup_domain('domain','v', mesh.dual_deg, mesh.dual_vertex-1, lon_v, lat_v, lon_i, lat_i) # dual mesh

#   CALL xios_set_domain_attr("u",ni_glo=ncell_tot, ibegin=displ, ni=ncell)                                                                              
#   CALL xios_set_domain_attr("u", data_dim=1, type='unstructured' , nvertex=2, i_index=ind_glo)                                                         
#   CALL xios_set_domain_attr("u",lonvalue_1d=lon, latvalue_1d=lat, bounds_lon_1d=bounds_lon, bounds_lat_1d=bounds_lat)                                   
#mesh_e = xios.Handle('domain','u')

# works up to that point !

dtime = xios.Duration(second=3600.)

#   CALL xios_set_timestep(dtime)                                                                                                                        
if True:
    calendar=xios.Handle('calendar_wrapper')
    print (dtime.year, dtime.month, dtime.day, dtime.hour, dtime.second)
    calendar.set_attr(timestep=dtime)
    calendar.update_timestep()
else:
    unst.ker.dynamico_xios_set_timestep(3600.)

#   CALL xios_set_fieldgroup_attr("standard_output", freq_op=itau_out*xios_timestep, freq_offset=(itau_out-1)*xios_timestep)                             
if False:
    std=xios.Handle('fieldgroup','standard_output')
    std.set_attr(freq_op=dtime)
    freq_offset = xios.Duration(second=0)
    std.set_attr(freq_offset=freq_offset)

if True:
    print 'xios.context_close_definition()'
    xios.context_close_definition()

no_error=True
try:
    for i in range(100):
        if True:
            unst.ker.dynamico_xios_update_calendar(i)
        else:
            xios.update_calendar(i)
        print 'send_field', i
        xios.send_field('ps', lat_i)
except:
    print 'An error has occured, trying to close XIOS properly'
    no_error=False
# CALL xios_recv_field(name,field)                                                                                                                       
# CALL xios_send_field(name,field_tmp)                                                                                                                   

print 'xios.context_finalize()'
xios.context_finalize()

print 'xios.finalize()'
xios.finalize()

print 'Done'

if not no_error : raise