""" Sum factorized geometry evaluations """

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

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

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

@kernel_cached
@backend(interface="spatial_coordinate", name="default")
def pymbolic_spatial_coordinate_multilinear(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


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


@preamble
def define_mesh_width(name):
    define_corner(name, False)
    lower = name_lowerleft_corner()
    return "{} -= {};".format(name, lower)


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


def name_meshwidth():
    name = "meshwidth"
    globalarg(name, dtype=np.float64, shape=(world_dimension(),))
    define_mesh_width(name)
    return name


@kernel_cached
@backend(interface="spatial_coordinate", name="diagonal_transformation_matrix")
def pymbolic_spatial_coordinate_axiparallel(visitor_indices):
    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.NEGATIVE
    from dune.perftool.sumfact.switch import get_facedir
    face = get_facedir(restriction)

    lowcorner = name_lowerleft_corner()
    meshwidth = name_meshwidth()

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

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