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

from dynamico import getargs
args=getargs.parse()
log_master, log_world = getargs.getLogger(__name__)
INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error
INFO_ALL, DEBUG_ALL = log_world.info, log_world.debug

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

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):
        INFO('xios_finalize()')
        DEBUG_ALL('xios_finalize')
        cxios.finalize()
    def min(self,data): return self.comm.allreduce(data, op=MPI.MIN)
    def max(self,data): return self.comm.allreduce(data, op=MPI.MAX)
    
class Context:
    def __enter__(self):
        INFO('cxios.context_close_definition()')
        cxios.context_close_definition()
        return self
    def __exit__(self, type, value, traceback):
        INFO('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):
        self.send_field(self.mesh.primal_own_loc, name, data)
    def send_field_dual(self,name, data):
        self.send_field(self.mesh.dual_own_loc, name, data)
    def send_field(self,own_loc, name, data):
        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)

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

def setup_curvilinear(cat,id,nx,ny,own,i_index,j_index,lon,lat):
    mesh = cxios.Handle(cat,id)
    def log(name,data) : 
        DEBUG('setup_curvilinear : %s shape min max %s %s %s'%(name, data.shape, data.min(), data.max()) )
    i_index, j_index, lon, lat = [ x[own] for x in i_index, j_index, lon, lat ]
    log('i_index',i_index)
    log('j_index',j_index)
    log('lon',lon)
    log('lat',lat)
    mesh.set_attr(type='curvilinear', ni_glo=nx, i_index=i_index, nj_glo=ny, j_index=j_index)
    mesh.set_attr(lonvalue_1d=lon.flatten(), latvalue_1d=lat.flatten())

class Context_Curvilinear(Context):
    def __init__(self, mesh, nqtot, step):
        self.init_llm(mesh, nqtot)
        # primal mesh
        setup_curvilinear('domaingroup','i', mesh.nx, mesh.ny, mesh.primal_own_loc, 
                          mesh.primal_i, mesh.primal_j, mesh.lon_i, mesh.lat_i)
        # dual mesh
        setup_curvilinear('domaingroup','v', mesh.nx, mesh.ny, mesh.dual_own_loc, 
                          mesh.dual_i, mesh.dual_j, mesh.lon_v, mesh.lat_v)
        # calendar
        self.init_calendar(step)

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

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 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)