import dynamico.unstructured as unst
import dynamico.xios as xios
from dynamico.meshes import radian, MPAS_Format, Unstructured_PMesh as PMesh, Local_Mesh as Mesh

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 numpy as np
import cProfile
import re

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

grid, llm, nqdyn, nqtot = 2562, 4,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

def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter                                                                                                                
print 'Reading MPAS mesh ...'
meshfile = MPAS_Format('grids/x1.%d.grid.nc'%grid)
pmesh = PMesh(comm,meshfile)
mesh = Mesh(pmesh, llm, nqdyn, radius, f)
print '...Done'

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

context=xios.XIOS_Context_Unstructured(pmesh,mesh,nqtot, 3600)

lat_i = radian*mesh.lat_i
lat_ik, junk  = np.meshgrid(lat_i, np.arange(llm), indexing='ij')

no_error=True
try:
    for i in range(100):
        context.update_calendar(i)
        print 'send_field', i, lat_ik.shape
        context.send_field_primal('ps', lat_i)
        context.send_field_primal('theta', lat_ik)
except:
    print 'An error has occured, trying to close XIOS properly'
    no_error=False

print 'xios.context_finalize()'
context.finalize()
print 'xios.finalize()'
xios.finalize()
print 'Done'

if not no_error : raise