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

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_domain(cat,id, degree,vertex,lon,lat,lon_v,lat_v, ncell_glo, index_glo, displ): # cat in 'domain','domaingroup'
    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_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 XIOS_Context:
    def __init__(self, pmesh, mesh,nqtot):
        self.mesh=mesh
        ker.dynamico_setup_xios()
        # 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))
        # 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]
        mesh_i=setup_domain('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
        calendar=cxios.Handle('calendar_wrapper')
        dtime = cxios.Duration(second=3600.)
        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)
        print 'cxios.context_close_definition()'
        cxios.context_close_definition()
    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 = np.ascontiguousarray(data)
        cxios.send_field(name, data)
    def update_calendar(self, i):
        ker.dynamico_xios_update_calendar(i)
    def finalize(self):
        cxios.context_finalize()