import numpy as nm

from sfepy.base.base import Struct, assert_
from sfepy.discrete.fem.mappings import VolumeMapping, SurfaceMapping
from poly_spaces import PolySpace
from fe_surface import FESurface

def set_mesh_coors(domain, fields, coors, update_fields=False, actual=False,
                   clear_all=True):
    if actual:
        domain.mesh.coors_act = coors.copy()
    else:
        domain.mesh.coors = coors.copy()

    if update_fields:
        for field in fields.itervalues():
            field.setup_coors(coors)
            field.clear_mappings(clear_all=clear_all)

def eval_nodal_coors(coors, mesh_coors, region, poly_space, geom_poly_space,
                     econn, ig, only_extra=True):
    """
    Compute coordinates of nodes corresponding to `poly_space`, given
    mesh coordinates and `geom_poly_space`.
    """
    if only_extra:
        iex = (poly_space.nts[:,0] > 0).nonzero()[0]
        if iex.shape[0] == 0: return

        qp_coors = poly_space.node_coors[iex, :]
        econn = econn[:, iex].copy()

    else:
        qp_coors = poly_space.node_coors

    ##
    # Evaluate geometry interpolation base functions in (extra) nodes.
    bf = geom_poly_space.eval_base(qp_coors)
    bf = bf[:,0,:].copy()

    ##
    # Evaluate extra coordinates with 'bf'.
    group = region.domain.groups[ig]
    cells = region.get_cells(ig)

    ecoors = nm.dot(bf, mesh_coors[group.conn[cells]])
    coors[econn] = nm.swapaxes(ecoors, 0, 1)


##
# 04.08.2005, c
def _interp_to_faces( vertex_vals, bfs, faces ):
    dim = vertex_vals.shape[1]
    n_face = faces.shape[0]
    n_qp = bfs.shape[0]
    
    faces_vals = nm.zeros( (n_face, n_qp, dim), nm.float64 )
    for ii, face in enumerate( faces ):
        vals = vertex_vals[face,:dim]
        faces_vals[ii,:,:] = nm.dot( bfs[:,0,:], vals )

    return( faces_vals )

class Interpolant( Struct ):
    """A simple wrapper around PolySpace."""

    def __init__(self, name, gel, space='H1', base='lagrange',
                 approx_order=1, force_bubble=False):
        self.name = name
        self.gel = gel

        self.poly_spaces = poly_spaces = {}
        poly_spaces['v'] = PolySpace.any_from_args(name, gel, approx_order,
                                                   base=base,
                                                   force_bubble=force_bubble)
        gel = gel.surface_facet
        if gel is not None:
            ps = PolySpace.any_from_args(name, gel, approx_order,
                                         base=base,
                                         force_bubble=False)
            skey = 's%d' % ps.n_nod
            poly_spaces[skey] = ps

    def describe_nodes( self ):
        ps = self.poly_spaces['v']
        node_desc = ps.describe_nodes()

        return node_desc

    ##
    # 16.11.2007, c
    def get_n_nodes( self ):
        nn = {}
        for key, ps in self.poly_spaces.iteritems():
            nn[key] = ps.nodes.shape[0]
        return nn

    def get_geom_poly_space(self, key):
        if key == 'v':
            ps = self.gel.interp.poly_spaces['v']

        elif key[0] == 's':
            n_v = self.gel.surface_facet.n_vertex
            ps = self.gel.interp.poly_spaces['s%d' % n_v]

        else:
            raise ValueError('bad polynomial space key! (%s)' % key)

        return ps

class SurfaceInterpolant(Interpolant):
    """
    Like Interpolant, but for use with SurfaceField and
    SurfaceApproximation.
    """

    def __init__(self, name, gel, space='H1', base='lagrange',
                 approx_order=1, force_bubble=False):
        Interpolant.__init__(self, name, gel, space=space, base=base,
                             approx_order=approx_order,
                             force_bubble=force_bubble)

        # Make alias 'v' <-> 's#'.
        ps = self.poly_spaces['v']
        self.poly_spaces['s%d' % ps.n_nod] = ps

    def get_geom_poly_space(self, key):
        
assert_(key[0] == 's')