# -*- coding: utf-8 -*-
"""
Fields for Discontinous Galerkin method
"""
import numpy as nm
import six
from numpy.lib.stride_tricks import as_strided
from six.moves import range

from sfepy.base.base import (output, assert_, Struct)
from sfepy.discrete import Integral, PolySpace
from sfepy.discrete.common.fields import parse_shape
from sfepy.discrete.fem.fields_base import FEField
from sfepy.discrete.fem.mappings import VolumeMapping


def get_unraveler(n_el_nod, n_cell):
    """Returns function for unraveling i.e. unpacking dof data from
    serialized array from shape (n_el_nod*n_cell, 1) to (n_cell, n_el_nod, 1).
    
    The unraveler returns non-writeable view into the input array.

    Parameters
    ----------
    n_el_nod : int
        expected dimensions of dofs array
    n_cell : int
        

    Returns
    -------
    unravel : callable

    """
    def unravel(u):
        """Returns non-writeable view into the input array reshaped  to
        (n*m, 1) to (m, n, 1) .

        Parameters
        ----------
        u : array_like
            solution in shape (n*m, 1)

        Returns
        -------
        u : ndarray
            unraveledsolution in shape (m, n, 1)

        """
        ustride1 = u.strides[0]
        ur = as_strided(u,
                        shape=(n_cell, n_el_nod, 1),
                        strides=(n_el_nod * ustride1, ustride1, ustride1),
                        writeable=False)
        return ur

    return unravel


def get_raveler(n_el_nod, n_cell):
    """Returns function for raveling i.e. packing dof data from
    two dimensional array of shape (n_cell, n_el_nod, 1) to (n_el_nod*n_cell, 1)
    
    The raveler returns view into the input array.

    Parameters
    ----------
    n_el_nod :
        param n_el_nod, n_cell: expected dimensions of dofs array
    n_cell : int
        

    Returns
    -------
    ravel : callable
    """
    def ravel(u):
        """Returns view into the input array reshaped from (m, n, 1) to (n*m, 1)
        to (m, n, 1) .

        Parameters
        ----------
        u : array_like

        Returns
        -------
        u : ndarray
        """

        # ustride1 = u.strides[0]
        # ur = as_strided(u, shape=(n_el_nod*n_cell, 1),
        #                     strides=(n_cell*ustride1, ustride1))
        ur = nm.ravel(u)[:, None]
        # possibly use according to
        #  https://docs.scipy.org/doc/numpy/reference/generated/numpy.ravel.html
        # ur = u.reshape(-1)
        return ur

    return ravel


# mapping between geometry element types
# and their facets types
# TODO move to sfepy/discrete/fem/geometry_element.py?
cell_facet_gel_name = {
    "1_2": "0_1",
    "2_3": "1_2",
    "2_4": "1_2",
    "3_4": "2_3",
    "3_8": "2_4"
}


def get_gel(region):
    """
    Parameters
    ----------
    region : sfepy.discrete.common.region.Region

    Returns
    -------
    gel :
        base geometry element of the region
    """
    cmesh = region.domain.cmesh
    for key, gel in six.iteritems(region.domain.geom_els):
        ct = cmesh.cell_types
        if (ct[region.cells] == cmesh.key_to_index[gel.name]).all():
            return gel
        else:
            raise ValueError('Region {} contains multiple'
                             ' reference geometries!'.format(region))


class DGField(FEField):
    """Class for usage with DG terms, provides functionality for Discontinous
    Galerkin method like neighbour look up, projection to discontinuous basis
    and correct DOF treatment.
    """
    family_name = 'volume_DG_legendre_discontinuous'
    is_surface = False

    def __init__(self, name, dtype, shape, region, space="H1",
                 poly_space_base="legendre", approx_order=1, integral=None):
        """
        Creates DGField, with Legendre polyspace and default integral
        corresponding to 2 * approx_order.

        Parameters
        ----------

        name : string

        dtype : type

        shape : string
        'vector', 'scalar' or something else

        region : sfepy.discrete.common.region.Region

        space : string
            default "H1"

        poly_space_base : PolySpace
            optionally force polyspace

        approx_order : 0 for FVM, default 1

        integral : Integral
            if None integral of order 2*approx_order is created
        """
        shape = parse_shape(shape, region.domain.shape.dim)
        Struct.__init__(self, name=name, dtype=dtype, shape=shape,
                        region=region)

        if isinstance(approx_order, tuple):
            self.approx_order = approx_order[0]
        else:
            self.approx_order = approx_order

        # geometry
        self.domain = region.domain
        self.region = region
        self.dim = region.tdim
        self._setup_geometry()
        self._setup_connectivity()
        # TODO treat domains embedded into higher dimensional spaces?
        self.n_el_facets = self.dim + 1 if self.gel.is_simplex else 2**self.dim

        # approximation space
        self.space = space
        self.poly_space_base = poly_space_base
        self.force_bubble = False
        self._create_interpolant()

        # DOFs
        self._setup_shape()
        self._setup_all_dofs()

        self.ravel_sol = get_raveler(self.n_el_nod, self.n_cell)
        self.unravel_sol = get_unraveler(self.n_el_nod, self.n_cell)

        # integral
        self.clear_qp_base()
        self.clear_facet_qp_base()
        if integral is None:
            self.integral = Integral("dg_fi", order = 2 * self.approx_order)
        else:
            self.integral = integral

        self.ori = None
        self.basis_transform = None

        # mapping
        self.mappings = {}
        self.mapping = self.create_mapping(self.region, self.integral, "volume",
                                           return_mapping=True)[1]
        self.mappings0 = {}

        # neighbour facet mapping and data caches
        # TODO use lru cache or different method?
        self.clear_facet_neighbour_idx_cache()
        self.clear_normals_cache()
        self.clear_facet_vols_cache()
        self.boundary_facet_local_idx = {}

    def _create_interpolant(self):
        name = self.gel.name + '_DG_legendre'
        ps = PolySpace.any_from_args(name, self.gel, self.approx_order,
                                     base=self.poly_space_base,
                                     force_bubble=False)
        self.poly_space = ps
        # 'legendre_simplex' is created for '1_2'.
        if self.gel.name in ["2_4", "3_8"]:
            self.extended = True
        else:
            self.extended = False

    def _setup_all_dofs(self):
        """Sets up all the differet kinds of DOFs, for DG only bubble DOFs"""
        self.n_el_nod = self.poly_space.n_nod
        self.n_vertex_dof = 0  # in DG we will propably never need vertex DOFs
        self.n_edge_dof = 0  # use facets DOFS for AFS methods
        self.n_face_dof = 0  # use facet DOF for AFS methods

        (self.n_bubble_dof,
         self.bubble_remap,
         self.bubble_dofs) = self._setup_bubble_dofs()

        self.n_nod = self.n_vertex_dof + self.n_edge_dof \
                     + self.n_face_dof + self.n_bubble_dof

    def _setup_bubble_dofs(self):
        """Creates DOF information for so called element, cell or bubble DOFs
        - the only DOFs used in DG
        n_dof = n_cells * n_el_nod
        remap optional remapping between cells
        dofs is mapping between dofs and cells

        Returns
        -------
        n_dof : int

        remap : ndarray

        dofs : ndarray
        """
        self.n_cell = self.region.get_n_cells(self.is_surface)
        n_dof = self.n_cell * self.n_el_nod

        dofs = nm.arange(n_dof, dtype=nm.int32)\
                   .reshape(self.n_cell, self.n_el_nod)
        remap = nm.arange(self.n_cell)
        self.econn = dofs
        self.dofs2cells = nm.repeat(nm.arange(self.n_cell), self.n_el_nod)

        return n_dof, remap, dofs

    def _setup_shape(self):
        """What is shape used for and what it really means.
        Does it represent shape of the problem?
        """
        self.n_components = nm.prod(self.shape)
        self.val_shape = self.shape

    def _setup_geometry(self):
        """Setup the field region geometry."""
        # get_gel extracts the highest dimension geometry from self.region
        self.gel = get_gel(self.region)

    def _setup_connectivity(self):
        """Forces self.domain.mesh to build necessary conductivities
        so they are available in self.get_nbrhd_dofs
        """
        self.region.domain.mesh.cmesh.setup_connectivity(self.dim, self.dim)
        self.region.domain.mesh.cmesh.setup_connectivity(self.dim - 1, self.dim)
        self.region.domain.mesh.cmesh.setup_connectivity(self.dim, self.dim - 1)

    def get_coor(self, nods=None):
        """Returns coors for matching nodes
        # TODO revise DG_EPBC and EPBC matching?

        Parameters
        ----------
        nods :
            if None use all nodes (Default value = None)

        Returns
        -------
        coors : ndarray
            coors on surface

        """

        if nods is None:
            nods = self.bubble_dofs

        cells = self.dofs2cells[nods]
        coors = self.domain.mesh.cmesh.get_centroids(self.dim)[cells]
        eps = min(self.domain.cmesh.get_volumes(self.dim)) / (self.n_el_nod + 2)
        if self.dim == 1:
            extended_coors = nm.zeros(nm.shape(coors)[:-1] + (2,))
            extended_coors[:, 0] = coors[:, 0]
            coors = extended_coors
        # shift centroid coors to lie within cells but be different for each dof
        # use coors of facet QPs?
        coors += eps * nm.repeat(nm.arange(self.n_el_nod),
                                 len(nm.unique(cells)))[:, None]
        return coors

    def clear_facet_qp_base(self):
        """Clears facet_qp_base cache"""
        self.facet_bf = {}
        self.facet_qp = None
        self.facet_whs = None

    def _transform_qps_to_facets(self, qps, geo_name):
        """Transforms points given in qps to all facets of the reference element
        with geometry geo_name.

        Parameters
        ----------
        qps :
            qps corresponding to facet dimension to be transformed
        geo_name :
            element type

        Returns
        -------
        tqps : ndarray
            tqps is of shape shape(qps) + (n_el_facets, geo dim)

        """
        if geo_name == "1_2":
            tqps = nm.zeros(nm.shape(qps) + (2, 1,))
            tqps[..., 0, 0] = 0.
            tqps[..., 1, 0] = 1.
        elif geo_name == "2_3":
            tqps = nm.zeros(nm.shape(qps) + (3, 2,))
            # 0.
            tqps[..., 0, 0] = qps  # x = 0 + t
            tqps[..., 0, 1] = 0.  # y = 0
            # 1.
            tqps[..., 1, 0] = 1 - qps  # x = 1 - t
            tqps[..., 1, 1] = qps  # y = t
            # 2.
            tqps[..., 2, 0] = 0  # x = 0
            tqps[..., 2, 1] = 1 - qps  # y = 1 - t
        elif geo_name == "2_4":
            tqps = nm.zeros(nm.shape(qps) + (4, 2,))
            # 0.
            tqps[..., 0, 0] = qps  # x = t
            tqps[..., 0, 1] = 0.  # y = 0
            # 1.
            tqps[..., 1, 0] = 1  # x = 1
            tqps[..., 1, 1] = qps  # y = t
            # 2.
            tqps[..., 2, 0] = 1 - qps  # x = 1 -t
            tqps[..., 2, 1] = 1  # y = 1
            # 3.
            tqps[..., 3, 0] = 0  # x = 0
            tqps[..., 3, 1] = 1 - qps  # y = 1 - t
        elif geo_name == "3_4":
            # tqps = nm.zeros(nm.shape(qps) + (4, 3,))
            raise NotImplementedError("Geometry {} not supported, yet"
                                      .format(geo_name))
        elif geo_name == "3_8":
            # tqps = nm.zeros(nm.shape(qps) + (8, 3,))
            raise NotImplementedError("Geometry {} not supported, yet"
                                      .format(geo_name))
        else:
            raise NotImplementedError("Geometry {} not supported, yet"
                                      .format(geo_name))
        return tqps

    def get_facet_qp(self):
        """Returns quadrature points on all facets of the reference element in
        array of shape (n_qp, 1 , n_el_facets, dim)

        Returns
        -------
        qps : ndarray
            quadrature points
        weights : ndarray
            Still needs to be transformed to actual facets!
        """

        if self.dim == 1:
            facet_qps = self._transform_qps_to_facets(nm.zeros((1, 1)), "1_2")
            weights = nm.ones((1, 1, 1))
        else:
            qps, weights = self.integral.get_qp(cell_facet_gel_name[self.gel.name])
            weights = weights[None, :, None]
            facet_qps = self._transform_qps_to_facets(qps, self.gel.name)

        return facet_qps, weights

    def get_facet_base(self, derivative=False, base_only=False):
        """
        Returns values of base in facets quadrature points, data shape is a bit
        crazy right now:
            (number of qps, 1, n_el_facets, 1, n_el_nod)
        end for derivatine:
            (1, number of qps, (dim,) * derivative, n_el_facets, 1, n_el_nod)

        Parameters
        ----------
        derivative: truthy or integer
        base_only: do not return weights

        Returns
        -------
        facet_bf : ndarray
            values of basis functions in facet qps
        weights : ndarray, optionally
            weights of qps
        """
        if derivative:
            diff = int(derivative)
        else:
            diff = 0

        if diff in self.facet_bf:
            facet_bf = self.facet_bf[diff]
            whs = self.facet_whs
        else:
            qps, whs = self.get_facet_qp()
            ps = self.poly_space
            self.facet_qp = qps
            self.facet_whs = whs
            if derivative:
                facet_bf = nm.zeros((1,) + nm.shape(qps)[:-1] +
                                    (self.dim,) * diff + (self.n_el_nod,))
            else:
                facet_bf = nm.zeros(nm.shape(qps)[:-1] + (1, self.n_el_nod,))

            for i in range(self.n_el_facets):
                facet_bf[..., i, :, :] = \
                    ps.eval_base(qps[..., i, :], diff=diff,
                                 transform=self.basis_transform)
            self.facet_bf[diff] = facet_bf

        if base_only:
            return facet_bf
        else:
            return facet_bf, whs

    def clear_facet_neighbour_idx_cache(self, region=None):
        """
        If region is None clear all!

        Parameters
        ----------
        region : sfepy.discrete.common.region.Region
            If None clear all.
        """
        if region is None:
            self.facet_neighbour_index = {}
        else:
            self.facet_neighbour_index.pop(region.name)

    def get_facet_neighbor_idx(self, region=None, eq_map=None):
        """
        Returns index of cell neighbours sharing facet, along with local index
        of the facet within neighbour, also treats periodic boundary conditions
        i.e. plugs correct neighbours for cell on periodic boundary.
        Where there are no neighbours specified puts -1  instead of neighbour
        and facet id

        Cashes neighbour index in self.facet_neighbours

        Parameters
        ----------
        region : sfepy.discrete.common.region.Region
            Main region, must contain cells.
        eq_map :
            eq_map from state variable containing information on
            EPBC and DG EPBC. (Default value = None)

        Returns
        -------
        facet_neighbours : ndarray
             Shape is
                (n_cell, n_el_facet, 2),
             first value is index of the neighbouring cell,
             the second is index of the facet in said nb. cell.
        """
        if region is None or eq_map is None:
            # HOTFIX enabling limiter to obtain connectivity data without
            # knowing eq_map or region
            if self.region.name in self.facet_neighbour_index:
                return self.facet_neighbour_index[self.region.name]
            else:
                raise ValueError("No facet neighbour mapping for main " +
                                 "region {}".format(self.region.name) +
                                 " cached yet, call with region and " +
                                 "eq_map first.")

        if region.name in self.facet_neighbour_index:
            return self.facet_neighbour_index[region.name]

        dim, n_cell, n_el_facets = self.get_region_info(region)

        cmesh = region.domain.mesh.cmesh
        cells = region.cells

        facet_neighbours = nm.zeros((n_cell, n_el_facets, 2), dtype=nm.int32)

        c2fi, c2fo = cmesh.get_incident(dim - 1, cells, dim, ret_offsets=True)

        for ic, o1 in enumerate(c2fo[:-1]):  # loop over cells
            o2 = c2fo[ic + 1]

            # get neighbours per facet of the cell
            c2ci, c2co = cmesh.get_incident(dim, c2fi[o1:o2], dim - 1,
                                            ret_offsets=True)
            ii = cmesh.get_local_ids(c2fi[o1:o2], dim - 1, c2ci, c2co, dim)
            fis = nm.c_[c2ci, ii]

            nbrs = []
            for ifa, of1 in enumerate(c2co[:-1]):  # loop over facets
                of2 = c2co[ifa + 1]
                if of2 == (of1 + 1):  # facet has only one cell
                    # Surface facet.
                    nbrs.append([-1, -1])  # c2ci[of1])  # append no neighbours
                else:
                    if c2ci[of1] == cells[ic]:  # do not append the cell itself
                        nbrs.append(fis[of2 - 1])
                    else:
                        nbrs.append(fis[of1])
            facet_neighbours[ic, :, :] = nbrs

        facet_neighbours = \
            self._set_fem_periodic_facet_neighbours(facet_neighbours, eq_map)

        facet_neighbours = \
            self._set_dg_periodic_facet_neighbours(facet_neighbours, eq_map)

        # cache results
        self.facet_neighbour_index[region.name] = facet_neighbours

        return facet_neighbours

    def _set_dg_periodic_facet_neighbours(self, facet_neighbours, eq_map):
        """

        Parameters
        ----------
        facet_neighbours : array_like
            Shape is
            (n_cell, n_el_facet, 2),
            first value is index of the neighbouring cell
            the second is index of the facet in said nb. cell.
        eq_map :
            must contain dg_ep_bc a List with pairs of slave and master boundary
            cell boundary facet mapping

        Returns
        -------
        facet_neighbours : ndarray
            Updated incidence array.

        
        """
        # if eq_map.
        # treat DG EPBC - these are definitely preferred
        if eq_map.n_dg_epbc > 0 and self.gel.name not in ["1_2", "2_4", "3_6"]:
            raise ValueError(
                "Periodic boundary conditions not supported " +
                "for geometry {} elements.".format(self.gel.name))

        dg_epbc = eq_map.dg_epbc

        for master_bc2bfi, slave_bc2bfi in dg_epbc:
            # set neighbours of periodic cells to one another
            facet_neighbours[master_bc2bfi[:, 0], master_bc2bfi[:, 1], 0] = \
                slave_bc2bfi[:, 0]
            facet_neighbours[slave_bc2bfi[:, 0], slave_bc2bfi[:, 1], 0] = \
                master_bc2bfi[:, 0]

            # set neighbours facets
            facet_neighbours[slave_bc2bfi[:, 0], slave_bc2bfi[:, 1], 1] = \
                master_bc2bfi[:, 1]
            facet_neighbours[master_bc2bfi[:, 0], master_bc2bfi[:, 1], 1] =\
                slave_bc2bfi[:, 1]

        return facet_neighbours

    def _set_fem_periodic_facet_neighbours(self, facet_neighbours, eq_map):
        """Maybe remove after DG EPBC revision in self.get_coor

        Parameters
        ----------
        facet_neighbours : array_like
            Shape is (n_cell, n_el_facet, 2), first value is index of the
            neighbouring cell the second is index of the facet in said nb. cell.
        eq_map :
            eq_map from state variable containing information on
            EPBC and DG EPBC.

        Returns
        -------
        facet_neighbours : ndarray
            Updated incidence array.
        """

        # treat classical FEM EPBCs - we need to correct neighbours
        if eq_map.n_epbc > 0:
            # set neighbours of periodic cells to one another
            mcells = nm.unique(self.dofs2cells[eq_map.master])
            scells = nm.unique(self.dofs2cells[eq_map.slave])
            mcells_facets = nm.array(
                nm.where(facet_neighbours[mcells] == -1))[1, 0]  # facets mcells
            scells_facets = nm.array(
                nm.where(facet_neighbours[scells] == -1))[1, 0]  # facets scells
            # [1, 0]  above, first we need second axis to get axis on which
            # facet indices are stored, second we drop axis with neighbour
            # local facet index,
            #
            # for multiple s/mcells this will have to be
            # something like 1 + 2*nm.arange(len(mcells)) - to skip double
            # entries for -1 tags in neighbours and  neighbour local facet idx

            # set neighbours of mcells to scells
            facet_neighbours[mcells, mcells_facets, 0] = scells
            # set neighbour facets to facets of scell missing neighbour
            facet_neighbours[
                mcells, mcells_facets, 1] = scells_facets
            # we do not need to distinguish EBC and EPBC cells, EBC overwrite
            # EPBC, we only need to fix shapes

            # set neighbours of scells to mcells
            facet_neighbours[scells, scells_facets, 0] = mcells
            # set neighbour facets to facets of mcell missing neighbour0
            facet_neighbours[
                scells, scells_facets, 1] = mcells_facets

        return facet_neighbours

    @staticmethod
    def get_region_info(region):
        """
        Extracts information about region needed in various methods of DGField

        Parameters
        ----------
        region : sfepy.discrete.common.region.Region
            

        Returns
        -------
            dim, n_cell, n_el_facets
        """
        if not region.has_cells():
            raise ValueError("Region {} has no cells".format(region.name))
        n_cell = region.get_n_cells()
        dim = region.tdim
        gel = get_gel(region)
        n_el_facets = dim + 1 if gel.is_simplex else 2 ** dim
        return dim, n_cell, n_el_facets

    def get_both_facet_state_vals(self, state, region,
                                  derivative=None, reduce_nod=True):
        """Computes values of the variable represented by dofs in
        quadrature points located at facets, returns both values -
        inner and outer, along with weights.

        Parameters
        ----------
        state : state variable containing BC info
            
        region : sfepy.discrete.common.region.Region
            
        derivative : compute derivative if truthy,
            compute n-th derivative if a number (Default value = None)
        reduce_nod : if False DOES NOT sum nodes into values at QPs
             (Default value = True)

        Returns
        -------
        inner_facet_values (n_cell, n_el_facets, n_qp),
                 outer facet values (n_cell, n_el_facets, n_qp),
                 weights,
                 if derivative is True:
                    inner_facet_values (n_cell, n_el_facets, dim, n_qp),
                    outer_facet values (n_cell, n_el_facets, dim, n_qp)

        
        """
        if derivative:
            diff = int(derivative)
        else:
            diff = 0
        unreduce_nod = int(not reduce_nod)

        inner_base_vals, outer_base_vals, whs = \
            self.get_both_facet_base_vals(state, region, derivative=derivative)
        dofs = self.unravel_sol(state.data[0])

        n_qp = whs.shape[-1]
        outputs_shape = (self.n_cell, self.n_el_facets) + \
                        (self.n_el_nod,) * unreduce_nod + \
                        (self.dim,) * diff + \
                        (n_qp,)

        inner_facet_vals = nm.zeros(outputs_shape)
        if unreduce_nod:
            inner_facet_vals[:] = nm.einsum('id...,idf...->ifd...',
                                            dofs, inner_base_vals)
        else:
            inner_facet_vals[:] = nm.einsum('id...,id...->i...',
                                            dofs, inner_base_vals)

        per_facet_neighbours = self.get_facet_neighbor_idx(region, state.eq_map)

        outer_facet_vals = nm.zeros(outputs_shape)
        for facet_n in range(self.n_el_facets):
            if unreduce_nod:
                outer_facet_vals[:, facet_n, :] = \
                    nm.einsum('id...,id...->id...',
                              dofs[per_facet_neighbours[:, facet_n, 0]],
                              outer_base_vals[:, :, facet_n])
            else:
                outer_facet_vals[:, facet_n, :] = \
                    nm.einsum('id...,id...->i...',
                              dofs[per_facet_neighbours[:, facet_n, 0]],
                              outer_base_vals[:, :, facet_n])

        boundary_cells = nm.array(nm.where(per_facet_neighbours[:, :, 0] < 0)).T
        outer_facet_vals[boundary_cells[:, 0], boundary_cells[:, 1]] = 0.0
        # TODO detect and print boundary cells without defined BCs?
        for ebc, ebc_vals in zip(state.eq_map.dg_ebc.get(diff, []),
                                 state.eq_map.dg_ebc_val.get(diff, [])):
            if unreduce_nod:
                raise NotImplementedError(
                    "Unreduced DOFs are not available for boundary " +
                    "outerfacets")
                outer_facet_vals[ebc[:, 0], ebc[:, 1], :] = \
                    nm.einsum("id,id...->id...",
                              ebc_vals, inner_base_vals[0, :, ebc[:, 1]])
            else:
                #  fix flipping qp order to accomodate for
                #  opposite facet orientation of neighbours
                outer_facet_vals[ebc[:, 0], ebc[:, 1], :] = ebc_vals[:, ::-1]

        # flip outer_facet_vals moved to get_both_facet_base_vals
        return inner_facet_vals, outer_facet_vals, whs