""" Sum factorized geometry evaluations """

from dune.perftool.generation import (domain,
                                      get_backend,
                                      get_counted_variable,
                                      iname,
                                      instruction,
                                      kernel_cached,
                                      temporary_variable,
                                      )
from dune.perftool.loopy.buffer import get_buffer_temporary
from dune.perftool.pdelab.geometry import (local_dimension,
                                           world_dimension,
                                           )
from dune.perftool.sumfact.symbolic import SumfactKernelInputBase
from dune.perftool.sumfact.vectorization import attach_vectorization_info

from pytools import ImmutableRecord

import pymbolic.primitives as prim


@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):
        name = get_buffer_temporary(sf.buffer,
                                    shape=(2 ** local_dimension(), sf.vector_width),
                                    name="input_{}".format(sf.buffer)
                                    )

        ciname = corner_iname()

        from dune.perftool.pdelab.geometry import name_geometry
        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)}),
                    )


@kernel_cached
def pymbolic_spatial_coordinate(visitor_indices):
    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 = (BasisTabulationMatrix(quadrature_size=quadrature_size, basis_size=2),) * local_dimension()
    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)

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