""" Sum factorized geometry evaluations """

from dune.perftool.generation import (backend,
                                      class_member,
                                      domain,
                                      get_backend,
                                      get_counted_variable,
                                      iname,
                                      instruction,
                                      kernel_cached,
                                      preamble,
                                      temporary_variable,
                                      globalarg,
                                      )
from dune.perftool.pdelab.geometry import (local_dimension,
                                           world_dimension,
                                           name_geometry,
                                           )
from dune.perftool.sumfact.switch import get_facedir
from dune.perftool.sumfact.symbolic import SumfactKernelInputBase
from dune.perftool.sumfact.vectorization import attach_vectorization_info
from dune.perftool.options import get_form_option, option_switch
from dune.perftool.ufl.modified_terminals import Restriction

from pytools import ImmutableRecord

import pymbolic.primitives as prim
import numpy as np


@iname
def corner_iname():
    name = get_counted_variable("corneriname")
    domain(name, 2 ** local_dimension())
    return name


class GeoCornersInput(SumfactKernelInputBase, ImmutableRecord):
    def __init__(self, dir):
        ImmutableRecord.__init__(self, dir=dir)

    def realize(self, sf, index, insn_dep):
        from dune.perftool.sumfact.realization import name_buffer_storage
        name = name="input_{}".format(sf.buffer)
        temporary_variable(name,
                           shape=(2 ** local_dimension(), sf.vector_width),
                           custom_base_storage=name_buffer_storage(sf.buffer, 0),
                           decl_method=buffer_decl(storage, get_sumfact_dtype(sf)),
                           managed=True,
                           )

        ciname = corner_iname()
        geo = name_geometry()

        # NB: We need to realize this as a C instruction, because the corner
        #     method does return a non-scalar, which does not fit into the current
        #     loopy philosophy for function calls. This problem will be solved once
        #     #11 is resolved. Admittedly, the code looks *really* ugly until that happens.
        code = "{}[{}*{}+{}] = {}.corner({})[{}];".format(name,
                                                          sf.vector_width,
                                                          ciname,
                                                          index,
                                                          geo,
                                                          ciname,
                                                          self.dir,
                                                          )

        instruction(code=code,
                    within_inames=frozenset({ciname}),
                    assignees=(name,),
                    tags=frozenset({"sumfact_stage{}".format(sf.stage)}),
                    )


@backend(interface="spatial_coordinate", name="default")
def pymbolic_spatial_coordinate_multilinear(do_predicates, visitor):
    assert len(visitor.indices) == 1

    # Construct the matrix sequence for the evaluation of the global coordinate.
    # We need to manually construct this one, because on facets, we want to use the
    # geometry embedding of the facet into the global space directly without going
    # through the neighboring cell geometries. That matrix sequence will only have
    # dim-1 matrices!
    from dune.perftool.sumfact.tabulation import quadrature_points_per_direction, BasisTabulationMatrix
    quadrature_size = quadrature_points_per_direction()
    matrix_sequence = tuple(BasisTabulationMatrix(direction=i, basis_size=2) for i in range(world_dimension()) if i != get_facedir(visitor.restriction))
    inp = GeoCornersInput(visitor.indices[0])

    from dune.perftool.sumfact.symbolic import SumfactKernel
    sf = SumfactKernel(matrix_sequence=matrix_sequence,
                       input=inp,
                       )

    vsf = attach_vectorization_info(sf)

    # Add a sum factorization kernel that implements the evaluation of
    # the basis functions at quadrature points (stage 1)
    from dune.perftool.sumfact.realization import realize_sum_factorization_kernel
    var, _ = realize_sum_factorization_kernel(vsf)

    visitor.indices = None
    return prim.Subscript(var, vsf.quadrature_index(sf, visitor))


@preamble
def define_corner(name, low):
    geo = name_geometry()
    return "auto {} = {}.corner({});".format(name,
                                             geo,
                                             0 if low else 2 ** local_dimension() - 1)


@class_member(classtag="operator")
def define_mesh_width(name):
    from dune.perftool.pdelab.localoperator import lop_template_range_field
    rft = lop_template_range_field()
    define_mesh_width_eval(name)
    return "Dune::FieldVector<{}, {}> {};".format(rft, world_dimension(), name)


def define_mesh_width_eval(name):
    from dune.perftool.pdelab.localoperator import name_ansatz_gfs_constructor_param
    gfs = name_ansatz_gfs_constructor_param()
    code = ["{",
            "  auto e = *({}.gridView().template begin<0>());".format(gfs),
            "  {} = e.geometry().corner((1<<{}) - 1);".format(name, world_dimension()),
            "  {} -= e.geometry().corner(0);".format(name),
            "}",
            ]
    instruction(code="\n".join(code),
                kernel="operator")


def name_lowerleft_corner():
    name = "lowerleft_corner"
    globalarg(name, shape=(world_dimension(),))
    define_corner(name, True)
    return name


def name_meshwidth():
    name = "meshwidth"
    globalarg(name, shape=(world_dimension(),))
    define_mesh_width(name)
    return name


@backend(interface="spatial_coordinate", name="diagonal_transformation_matrix")
def pymbolic_spatial_coordinate_axiparallel(do_predicates, visitor):
    assert len(visitor.indices) == 1
    index, = visitor.indices

    # Urgh: *SOMEHOW* construct a face direction
    from dune.perftool.pdelab.restriction import Restriction
    restriction = Restriction.NONE
    from dune.perftool.generation import get_global_context_value
    if get_global_context_value("integral_type") == "interior_facet":
        restriction = Restriction.POSITIVE
    from dune.perftool.sumfact.switch import get_facedir
    face = get_facedir(restriction)

    lowcorner = name_lowerleft_corner()
    meshwidth = name_meshwidth()

    # If we have to decide which boundary condition to take for this
    # intersection we always take the boundary condition of the center
    # of the intersection. We assume that there are no cells with more
    # than one boundary condition.
    if do_predicates:
        x = 0.5
    elif index == face:
        x = 0
    else:
        iindex = index
        if face is not None and index > face:
            iindex = iindex - 1
        from dune.perftool.sumfact.quadrature import pymbolic_quadrature_position
        x = pymbolic_quadrature_position(iindex, visitor)

    visitor.indices = None
    return prim.Subscript(prim.Variable(lowcorner), (index,)) + x * prim.Subscript(prim.Variable(meshwidth), (index,))


def pymbolic_unit_outer_normal(visitor_indices):
    index, = visitor_indices
    assert isinstance(index, int)
    if get_form_option("diagonal_transformation_matrix"):
        from dune.perftool.sumfact.switch import get_facedir, get_facemod
        if index == get_facedir(Restriction.POSITIVE):
            if get_facemod(Restriction.POSITIVE):
                return 1, None
            else:
                return -1, None
        else:
            return 0, None
    else:
        from dune.perftool.pdelab.geometry import pymbolic_unit_outer_normal as _norm
        return _norm(), visitor_indices


def pymbolic_unit_inner_normal(visitor_indices):
    index, = visitor_indices
    assert isinstance(index, int)
    if get_form_option("diagonal_transformation_matrix"):
        from dune.perftool.sumfact.switch import get_facedir, get_facemod
        if index == get_facedir(Restriction.POSITIVE):
            if get_facemod(Restriction.POSITIVE):
                return -1, None
            else:
                return 1, None
        else:
            return 0, None
    else:
        from dune.perftool.pdelab.geometry import pymbolic_unit_inner_normal as _norm
        return _norm(), visitor_indices


def pymbolic_facet_jacobian_determinant():
    if get_form_option("constant_transformation_matrix"):
        return pymbolic_constant_facet_jacobian_determinant()
    else:
        from dune.perftool.pdelab.geometry import pymbolic_facet_jacobian_determinant as _norm
        return _norm()


def pymbolic_constant_facet_jacobian_determinant():
    facedir = get_facedir(Restriction.POSITIVE)
    assert isinstance(facedir, int)

    name = "fdetjac"
    define_constant_facet_jacobian_determinant(name)
    globalarg(name, shape=(world_dimension(),))
    return prim.Subscript(prim.Variable(name), (facedir,))


@class_member(classtag="operator")
def define_constant_facet_jacobian_determinant(name):
    define_constant_facet_jacobian_determinant_eval(name)
    from dune.perftool.pdelab.localoperator import lop_template_ansatz_gfs
    gfst = lop_template_ansatz_gfs()
    return "std::array<typename {}::Traits::GridView::template Codim<0>::Geometry::ctype, {}> {};".format(gfst, world_dimension(), name)


@kernel_cached(kernel="operator")
def define_constant_facet_jacobian_determinant_eval(name):
    from dune.perftool.pdelab.localoperator import name_ansatz_gfs_constructor_param, lop_template_range_field
    gfs = name_ansatz_gfs_constructor_param()
    rft = lop_template_range_field()
    code = ["{",
            "  auto e = *({}.gridView().template begin<0>());".format(gfs),
            "  int dir=0;",
            "  for(auto is = {0}.gridView().ibegin(e); is != {0}.gridView().iend(e); ++(++is), ++dir)".format(gfs),
            "    {}[dir] = is->geometry().integrationElement(Dune::FieldVector<{}, {}>());".format(name, rft, world_dimension() - 1),
            "}",
            ]
    instruction(code="\n".join(code),
                kernel="operator",
                )


def pymbolic_facet_area():
    if get_form_option("constant_transformation_matrix"):
        return pymbolic_facet_jacobian_determinant()
    else:
        from dune.perftool.pdelab.geometry import pymbolic_facet_area as _norm
        return _norm()