import numpy as np
import cxios
from unstructured import ker
from dynamico.meshes import radian
from mpi4py import MPI

#-----------------------------------------------------------------------------

def setup_curvilinear(cat,id,nx,ny,displ,index_own_i,index_own_j,lon,lat):
    mesh = cxios.Handle(cat,id)
    def log(name,data) : print 'setup_curvilinear : %s shape min max'%name, data.shape, data.min(), data.max()
    log('index_own_i',index_own_i)
    log('index_own_j',index_own_j)
    log('lon',lon)
    log('lat',lat)
    mesh.set_attr(type='curvilinear',
                  ni_glo=nx, i_index=index_own_i,
                  nj_glo=ny, j_index=index_own_j)
    mesh.set_attr(lonvalue_1d=lon.flatten(), latvalue_1d=lat.flatten())

def cf_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

def setup_unstructured(cat,id, degree,vertex,lon,lat,lon_v,lat_v, ncell_glo, index_own_glo, displ): # cat in 'domain','domaingroup'
    # ncell_glo (int) : global number of cells
    # index_own_glo : global indices of cells own by this MPI process
    # displ : min of index_own across all MPI processes
    bnd_lon, bnd_lat = cf_boundaries(degree, vertex, lon_v, lat_v)
    ncell, nvertex = bnd_lon.shape
    mesh = cxios.Handle(cat,id)
    mesh.set_attr(type='unstructured')
    mesh.set_attr(ni_glo=ncell_glo, ni=ncell, ibegin=displ, i_index=index_own_glo)
    mesh.set_attr(lonvalue_1d=lon, latvalue_1d=lat)
    mesh.set_attr(nvertex=nvertex,bounds_lon_1d=bnd_lon, bounds_lat_1d=bnd_lat)

#-----------------------------------------------------------------------------

class Client:
    def __enter__(self):
        ker.dynamico_setup_xios()
        self.comm=MPI.COMM_WORLD # FIXME : we should use the communicator returned by XIOS
        return self
    def __exit__(self, type, value, traceback):
        print 'xios_finalize()'
        cxios.finalize()

class Context:
    def __enter__(self):
        print 'cxios.context_close_definition()'
        cxios.context_close_definition()
        return self
    def __exit__(self, type, value, traceback):
        print 'xios_context_finalize()'
        cxios.context_finalize()
    def init_llm(self, mesh, nqtot):
        self.mesh=mesh
        # vertical axis, nq
        llm = mesh.llm
        lev, levp1, nq = [cxios.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))
    def init_calendar(self, step):
        calendar=cxios.Handle('calendar_wrapper')
        dtime = cxios.Duration(second=step)
        calendar.set_attr(timestep=dtime)
        calendar.update_timestep()
        std=cxios.Handle('fieldgroup','standard_output')
        std.set_attr(freq_op=dtime)
        freq_offset = cxios.Duration(second=0)
        std.set_attr(freq_offset=freq_offset)
    def update_calendar(self, i):
        ker.dynamico_xios_update_calendar(i)
    def send_field_primal(self,name, data):
        own_loc = self.mesh.primal_own_loc
        if data.ndim==1:
            data = data[own_loc]
        if data.ndim==2:
            data = data[own_loc,:]
            data = data.transpose() # XIOS expects contiguous horizontal slices
        data = np.ascontiguousarray(data)
        cxios.send_field(name, data)

class Context_Curvilinear(Context):
    def __init__(self, mesh, nqtot, step):
        self.init_llm(mesh, nqtot)
        # primal mesh
        own = mesh.primal_own_loc
        lon_i, lat_i = [ x[own] for x in mesh.lon_i, mesh.lat_i ]
        primal_i, primal_j = [ x[own] for x in mesh.primal_i, mesh.primal_j ]
        setup_curvilinear('domaingroup','i', mesh.nx, mesh.ny, mesh.displ,
                          primal_i, primal_j, lon_i, lat_i)
        # calendar
        self.init_calendar(step)

class Context_Unstructured(Context):
    def __init__(self, pmesh, mesh,nqtot, step):
        self.init_llm(mesh, nqtot)
        # primal mesh
        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]
        setup_unstructured('domaingroup','i', mesh.primal_own_deg, mesh.primal_vertex,
                           lon_i[mesh.primal_own_loc], lat_i[mesh.primal_own_loc],
                           lon_v, lat_v,
                           pmesh.dim_primal.n, mesh.primal_own_glo, mesh.displ)
        # calendar
        self.init_calendar(step)